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) tboard.on_end(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({}) self.assertEqual(mock_board.return_value.close.call_count, 1)
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_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_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_single_channel_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, 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=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_writer.return_value.add_image.assert_called_once_with('test1', 10, 1) self.assertTrue(mock_grid.call_args[0][0].size() == torch.ones(18, 1, 10, 10).size()) tboard.on_end({})
def test_multi_batch(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=36, 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) 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( 36, 3, 10, 10).size())