def test_set_parameters(self): learning_rate = 1 optim = Optimizer(torch.optim.SGD, lr=learning_rate) params = [torch.nn.Parameter(torch.randn(2, 3, 4))] optim.set_parameters(params) self.assertTrue(type(optim.optimizer) is torch.optim.SGD) self.assertEquals(optim.optimizer.param_groups[0]['lr'], learning_rate)
def test_update(self): optim = Optimizer(torch.optim.SGD, lr=1, decay_after_epoch=5, lr_decay=0.5) params = [torch.nn.Parameter(torch.randn(2, 3, 4))] optim.set_parameters(params) optim.update(0, 10) self.assertEquals(optim.optimizer.param_groups[0]['lr'], 0.5)
def test_step(self, mock_clip_grad_norm): optim = Optimizer(torch.optim.Adam, max_grad_norm=5) params = [torch.nn.Parameter(torch.randn(2, 3, 4))] optim.set_parameters(params) optim.step() mock_clip_grad_norm.assert_called_once()