def test_optimizer_params_handler_wrong_setup():
with pytest.raises(TypeError):
OptimizerParamsHandler(optimizer=None)
optimizer = MagicMock(spec=torch.optim.Optimizer)
handler = OptimizerParamsHandler(optimizer=optimizer)
mock_logger = MagicMock()
mock_engine = MagicMock()
with pytest.raises(
RuntimeError,
match="Handler OptimizerParamsHandler works only with WandBLogger"
):
handler(mock_engine, mock_logger, Events.ITERATION_STARTED)
def test_wandb_close():
optimizer = torch.optim.SGD([torch.Tensor(0)], lr=0.01)
wrapper = OptimizerParamsHandler(optimizer=optimizer, param_name="lr")
mock_logger = MagicMock(spec=WandBLogger)
mock_logger.log = MagicMock()
mock_engine = MagicMock()
wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
mock_logger.close()
def test_optimizer_params():
optimizer = torch.optim.SGD([torch.Tensor(0)], lr=0.01)
wrapper = OptimizerParamsHandler(optimizer=optimizer, param_name="lr")
mock_logger = MagicMock(spec=WandBLogger)
mock_logger.log = MagicMock()
mock_engine = MagicMock()
mock_engine.state = State()
mock_engine.state.iteration = 123
wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
mock_logger.log.assert_called_once_with({"lr/group_0": 0.01},
step=123,
sync=None)
wrapper = OptimizerParamsHandler(optimizer,
param_name="lr",
tag="generator")
mock_logger = MagicMock(spec=WandBLogger)
mock_logger.log = MagicMock()
wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
mock_logger.log.assert_called_once_with({"generator/lr/group_0": 0.01},
step=123,
sync=None)