def test_log_dir(self, mock_board):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoardImages(log_dir='./test', comment='torchbearer')
tboard.on_start(state)
mock_board.assert_called_once_with(log_dir=os.path.join('./test', 'Sequential_torchbearer'))
def test_log_dir(self, mock_board):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoardImages(log_dir='./test', comment='torchbearer')
tboard.on_start(state)
mock_board.assert_called_once_with(
log_dir=os.path.join('./test', 'Sequential_torchbearer'))
def test_writer_closed_on_end(self, mock_board):
mock_board.return_value = Mock()
mock_board.return_value.close = Mock()
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoardImages()
tboard.on_start(state)
tboard.on_end({})
mock_board.return_value.close.assert_called_once()
def test_single_channel(self, mock_board, mock_grid):
mock_board.return_value = Mock()
mock_board.return_value.add_image = Mock()
mock_grid.return_value = 10
state = {'x': torch.ones(18, 10, 10), torchbearer.EPOCH: 1, torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoardImages(name='test', key='x', write_each_epoch=True, num_images=18, nrow=9, padding=3, normalize=True, range='tmp', scale_each=True, pad_value=1)
tboard.on_start(state)
tboard.on_step_validation(state)
mock_grid.assert_called_once_with(ANY, nrow=9, padding=3, normalize=True, range='tmp', scale_each=True, pad_value=1)
mock_board.return_value.add_image.assert_called_once_with('test', 10, 1)
self.assertTrue(mock_grid.call_args[0][0].size() == torch.ones(18, 1, 10, 10).size())
def test_odd_batches_visdom(self, mock_visdom, mock_writer, _, mock_grid):
mock_writer.return_value = Mock()
mock_writer.return_value.add_image = Mock()
mock_grid.return_value = 10
state = {
'x': torch.ones(18, 3, 10, 10),
torchbearer.EPOCH: 1,
torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))
}
tboard = TensorBoardImages(visdom=True,
name='test',
key='x',
write_each_epoch=True,
num_images=40,
nrow=9,
padding=3,
normalize=True,
norm_range='tmp',
scale_each=True,
pad_value=1)
tboard.on_start(state)
tboard.on_step_validation(state)
tboard.on_step_validation(state)
tboard.on_step_validation(state)
mock_grid.assert_called_once_with(ANY,
nrow=9,
padding=3,
normalize=True,
range='tmp',
scale_each=True,
pad_value=1)
mock_writer.return_value.add_image.assert_called_once_with(
'test1', 10, 1)
self.assertTrue(mock_grid.call_args[0][0].size() == torch.ones(
40, 3, 10, 10).size())
tboard.on_end({})
def test_writer_closed_on_end(self, mock_board, _):
mock_board.return_value = Mock()
mock_board.return_value.close = Mock()
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoardImages()
tboard.on_start(state)
tboard.on_end({})
self.assertEqual(mock_board.return_value.close.call_count, 1)
def test_log_dir_visdom(self, mock_visdom, mock_writer, _):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
mock_writer.__delete__ = Mock()
tboard = TensorBoardImages(visdom=True, log_dir='./test', comment='torchbearer')
tboard.on_start(state)
tboard.on_end(state)
self.assertEqual(mock_visdom.call_count, 1)
self.assertTrue(mock_visdom.call_args[1]['log_to_filename'] == os.path.join('./test', 'Sequential_torchbearer',
'log.log'))
def test_simple_case(self, mock_board, mock_grid, _):
mock_board.return_value = Mock()
mock_board.return_value.add_image = Mock()
mock_grid.return_value = 10
state = {'x': torch.ones(18, 3, 10, 10), torchbearer.EPOCH: 1,
torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoardImages(name='test', key='x', write_each_epoch=False, num_images=18, nrow=9, padding=3,
normalize=True, norm_range='tmp', scale_each=True, pad_value=1)
tboard.on_start(state)
tboard.on_step_validation(state)
mock_grid.assert_called_once_with(ANY, nrow=9, padding=3, normalize=True, range='tmp', scale_each=True,
pad_value=1)
mock_board.return_value.add_image.assert_called_once_with('test', 10, 1)
self.assertTrue(mock_grid.call_args[0][0].size() == state['x'].size())
tboard.on_end({})
