class TFModelTest(ModelTest, unittest.TestCase): flag = False def create_model(self): model = tf.keras.models.Sequential([ keras.layers.Dense(512, activation="relu", input_shape=(784, )), keras.layers.Dropout(0.2), keras.layers.Dense(10), ]) model.compile( optimizer="adam", loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=["accuracy"], ) return model 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_tfmodel_saver_from_json_file(self): self.test_tfmodel_saver(save_to_json_file=True) self.model_config = None self.model_manager = None self.model_config = ModelConfig.from_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_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_tfmodel_saver_from_json_file(self): self.test_tfmodel_saver(save_to_json_file=True) self.model_config = None self.model_manager = None self.model_config = ModelConfig.from_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_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
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)
class TorchModelTest(ModelTest, unittest.TestCase): transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, ))]) flag = False def test_torchmodel_saver(self, model_config=None, save_to_json_file=False): if model_config is None: self.model_config = ModelConfig(**self.storage_config, model_name="test_torchmodel_saver", additional={ "framework": "pytorch" }) 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) model = Model() self.model_manager.save(model, step=0) def test_torchmodel_saver_from_json_file(self): self.test_torchmodel_saver(save_to_json_file=True) self.model_config = None self.model_manager = None self.model_config = ModelConfig.from_json_file(self.model_config_file) self.model_manager = ModelManager(config=self.model_config) model = Model() self.model_manager.save(model, step=0) def test_torchmodel_loader(self): self.test_torchmodel_saver() model = Model() self.model_manager.load(model, step=0) 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_torchmodel_layer_loader(self): self.test_torchmodel_saver() self.model_manager = ModelManager(config=self.model_config) self.model_manager.load("f.weight", step=0) def test_mnist_eval(self, model=None, device=None): if not self.flag: raise SkipTest test_dataset = datasets.MNIST("/tmp/data", train=False, transform=self.transform) test_loader = DataLoader(test_dataset, batch_size=64, num_workers=4) model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for image, target in test_loader: image, target = image.to(device), target.to(device) output = model(image) test_loss += F.nll_loss(output, target, reduction="sum").item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) return correct def test_mnist_reloaded(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.storage_config, model_name="testmodel", 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_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
class TFModelTest(ModelTest, unittest.TestCase): def create_model(self): model = tf.keras.models.Sequential([ keras.layers.Dense(512, activation="relu", input_shape=(784, )), keras.layers.Dropout(0.2), keras.layers.Dense(10), ]) model.compile( optimizer="adam", loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=["accuracy"], ) return model 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_tfmodel_saver_from_json_file(self): self.test_tfmodel_saver(save_to_json_file=True) self.model_config = None self.model_manager = None self.model_config = ModelConfig.from_json_file(self.model_config_file) self.model_manager = ModelManager(config=self.model_config) self.model_manager.save(self.model, step=0) @unittest.skip("skip") 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) @unittest.skip("skip") 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 def test_mnist_reloaded(self): import multiprocessing process_train = multiprocessing.Process( target=self.test_mnist_train_process) process_train.start() process_train.join() process_eval = multiprocessing.Process( target=self.test_mnist_eval_process) process_eval.start() process_eval.join()