def test_create_restore_delete(self):
# Create the hyperparameters and objects to save.
hp = models.registry.get_default_hparams('cifar_resnet_20')
model = models.registry.get(hp.model_hparams)
optimizer = optimizers.get_optimizer(hp.training_hparams, model)
dataloader = datasets.registry.get(hp.dataset_hparams)
step = Step.from_epoch(13, 27, 400)
# Run one step of SGD.
examples, labels = next(iter(dataloader))
optimizer.zero_grad()
model.train()
model.loss_criterion(model(examples), labels).backward()
optimizer.step()
# Create a fake logger.
logger = MetricLogger()
logger.add('test_accuracy', Step.from_epoch(0, 0, 400), 0.1)
logger.add('test_accuracy', Step.from_epoch(10, 0, 400), 0.5)
logger.add('test_accuracy', Step.from_epoch(100, 0, 400), 0.8)
# Save a checkpoint.
checkpointing.save_checkpoint_callback(self.root, step, model, optimizer, logger)
self.assertTrue(os.path.exists(paths.checkpoint(self.root)))
# Create new models.
model2 = models.registry.get(hp.model_hparams)
optimizer2 = optimizers.get_optimizer(hp.training_hparams, model)
# Ensure the new model has different weights.
sd1, sd2 = model.state_dict(), model2.state_dict()
for k in model.prunable_layer_names:
self.assertFalse(np.array_equal(sd1[k].numpy(), sd2[k].numpy()))
self.assertIn('momentum_buffer', optimizer.state[optimizer.param_groups[0]['params'][0]])
self.assertNotIn('momentum_buffer', optimizer2.state[optimizer.param_groups[0]['params'][0]])
# Restore the checkpointt.
step2, logger2 = checkpointing.restore_checkpoint(self.root, model2, optimizer2, 400)
self.assertTrue(os.path.exists(paths.checkpoint(self.root)))
self.assertEqual(step, step2)
self.assertEqual(str(logger), str(logger2))
# Ensure the new model is now the same.
sd1, sd2 = model.state_dict(), model2.state_dict()
self.assertEqual(set(sd1.keys()), set(sd2.keys()))
for k in sd1:
self.assertTrue(np.array_equal(sd1[k].numpy(), sd2[k].numpy()))
# Ensure the new optimizer is now the same.
mom1 = optimizer.state[optimizer.param_groups[0]['params'][0]]['momentum_buffer']
mom2 = optimizer2.state[optimizer.param_groups[0]['params'][0]]['momentum_buffer']
self.assertTrue(np.array_equal(mom1.numpy(), mom2.numpy()))
def create_logger():
logger = MetricLogger()
logger.add('train_accuracy', Step.from_iteration(0, 400), 0.5)
logger.add('train_accuracy', Step.from_iteration(1, 400), 0.6)
logger.add('test_accuracy', Step.from_iteration(0, 400), 0.4)
return logger