def test_torchmodel_loader_with_compressor(self):
model_config = ModelConfig(
**self.storage_config,
model_name="test_torchmodel_loader_with_compressor",
additional={"framework": "pytorch"},
compressor={
"complevel": 4,
"complib": "zlib"
})
self.test_torchmodel_saver(model_config=model_config)
self.model_manager = ModelManager(config=self.model_config)
model = Model()
self.model_manager.load(model, step=0)
def test_mnist_train_process(self):
(train_images, train_labels), _ = tf.keras.datasets.mnist.load_data()
train_labels = train_labels[:1000]
train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0
model = self.create_model()
model.fit(train_images, train_labels, epochs=5)
self.model_config = ModelConfig(
endpoint="127.0.0.1:9000",
access_key="minio",
secret_key="miniosecretkey",
model_name="test_tf_mnist",
additional={"version": "1.0.1"},
)
self.model_manager = ModelManager(config=self.model_config)
self.model_manager.save(model, epoch=1)
def test_tfmodel_saver(self, model_config=None, save_to_json_file=False):
self.model = self.create_model()
if model_config is None:
self.model_config = ModelConfig(**self.storage_config,
model_name="test_tfmodel_saver",
additional={
"framework": "tensorflow"
})
else:
self.model_config = model_config
if save_to_json_file:
self.model_config_file = "model_config_file.json"
self.model_config.to_json_file(self.model_config_file)
self.model_manager = ModelManager(config=self.model_config)
self.model_manager.save(self.model, step=0)
def test_mnist_reloaded_nas(self):
self.flag = True
import torch.optim as optim
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_dataset = datasets.MNIST("/tmp/data",
train=True,
download=True,
transform=self.transform)
model = Model().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()
train_loader = DataLoader(train_dataset, batch_size=64, num_workers=4)
for batch_idx, (image, target) in enumerate(tqdm(train_loader)):
image, target = image.to(device), target.to(device)
optimizer.zero_grad()
output = model(image)
loss = criterion(output, target)
loss.backward()
optimizer.step()
self.model_config = ModelConfig(**self.nas_config,
model_name="test_torch_mnist_nas",
additional={"version": "1.0.1"})
self.model_manager = ModelManager(config=self.model_config)
self.model_manager.save(model, epoch=1)
pretrained_model = Model().to(device)
correct = self.test_mnist_eval(model=pretrained_model, device=device)
self.model_manager.load(pretrained_model, epoch=1)
pretrained_correct = self.test_mnist_eval(model=pretrained_model,
device=device)
assert correct < pretrained_correct
def test_mnist_eval_process(self):
_, (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
test_labels = test_labels[:1000]
test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0
model = self.create_model()
_, correct = model.evaluate(test_images, test_labels, verbose=2)
self.model_config = ModelConfig(
endpoint="127.0.0.1:9000",
access_key="minio",
secret_key="miniosecretkey",
model_name="test_tf_mnist",
additional={"version": "1.0.1"},
)
self.model_manager = ModelManager(config=self.model_config)
self.model_manager.load(model, epoch=1)
_, pretrained_correct = model.evaluate(test_images,
test_labels,
verbose=2)
assert correct < pretrained_correct