def test_batch_metrics_visdom(self, mock_visdom, mock_writer, _):
mock_writer.return_value = Mock()
mock_writer.return_value.add_scalar = Mock()
state = {
torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)),
torchbearer.EPOCH: 0,
torchbearer.METRICS: {
'test': 1
},
torchbearer.BATCH: 0,
torchbearer.TRAIN_STEPS: 0
}
tboard = TensorBoard(visdom=True,
write_batch_metrics=True,
write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_start_epoch(state)
tboard.on_step_training(state)
mock_writer.return_value.add_scalar.assert_called_once_with(
'test', 1, 0, main_tag='batch')
mock_writer.return_value.add_scalar.reset_mock()
tboard.on_step_validation(state)
mock_writer.return_value.add_scalar.assert_called_once_with(
'test', 1, 0, main_tag='batch')
tboard.on_end_epoch(state)
tboard.on_end(state)
def test_log_dir(self, mock_board):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoard(write_epoch_metrics=False)
tboard.on_start(state)
mock_board.assert_called_once_with(log_dir=os.path.join('./logs', 'Sequential_torchbearer'))
def test_log_dir(self, mock_board, _):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoard(write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_end(state)
mock_board.assert_called_once_with(log_dir=os.path.join('./logs', 'Sequential_torchbearer'))
def test_batch_log_dir(self, mock_board):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0}
tboard = TensorBoard(write_batch_metrics=True, write_graph=False, write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_start_epoch(state)
mock_board.assert_called_with(log_dir=os.path.join('./logs', 'Sequential_torchbearer', 'epoch-0'))
def test_batch_log_dir(self, mock_board, _):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0}
tboard = TensorBoard(write_batch_metrics=True, write_graph=False, write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_start_epoch(state)
tboard.on_end_epoch(state)
tboard.on_end(state)
mock_board.assert_called_with(log_dir=os.path.join('./logs', 'Sequential_torchbearer', 'epoch-0'))
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 = TensorBoard(write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_end({})
mock_board.return_value.close.assert_called_once()
def test_epoch_metrics(self, mock_board):
mock_board.return_value = Mock()
mock_board.return_value.add_scalar = Mock()
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0, torchbearer.METRICS: {'test': 1}}
tboard = TensorBoard(write_batch_metrics=False, write_epoch_metrics=True)
tboard.on_start(state)
tboard.on_end_epoch(state)
mock_board.return_value.add_scalar.assert_called_once_with('epoch/test', 1, 0)
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 = TensorBoard(write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_end({})
mock_board.return_value.close.assert_called_once()
def test_writer_closed_on_end_visdom(self, mock_visdom, mock_writer, _):
mock_writer.return_value = Mock()
mock_writer.return_value.close = Mock()
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoard(visdom=True, write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_end({})
self.assertEqual(mock_writer.return_value.close.call_count, 1)
def test_batch_log_dir_visdom(self, mock_visdom, mock_writer, _):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)),
torchbearer.EPOCH: 0, torchbearer.METRICS: {'test': 1}, torchbearer.BATCH: 0}
tboard = TensorBoard(visdom=True, write_batch_metrics=True, write_graph=False, write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_start_epoch(state)
tboard.on_end_epoch(state)
tboard.on_end(state)
self.assertTrue(mock_visdom.call_args[1]['log_to_filename'] == os.path.join('./logs', 'Sequential_torchbearer', 'epoch', 'log.log'))
def test_writer_closed_on_end_visdom_visdom(self, mock_visdom, mock_writer,
_):
mock_writer.return_value = Mock()
mock_writer.return_value.close = Mock()
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoard(visdom=True)
tboard.on_start(state)
tboard.on_end({})
mock_writer.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 = TensorBoard(visdom=True, write_epoch_metrics=False)
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('./logs', 'Sequential_torchbearer',
'log.log'))
def test_batch_writer_closed_on_end_epoch(self, mock_board, _):
mock_board.return_value = Mock()
mock_board.return_value.close = Mock()
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0}
tboard = TensorBoard(write_batch_metrics=True, write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_start_epoch(state)
tboard.on_end_epoch({})
self.assertEqual(mock_board.return_value.close.call_count, 1)
tboard.on_end(state)
def test_epoch_metrics(self, mock_board, _):
mock_board.return_value = Mock()
mock_board.return_value.add_scalar = Mock()
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0,
torchbearer.METRICS: {'test': 1}}
tboard = TensorBoard(write_batch_metrics=False, write_epoch_metrics=True)
tboard.on_start(state)
tboard.on_start_epoch(state)
tboard.on_end_epoch(state)
mock_board.return_value.add_scalar.assert_called_once_with('epoch/test', 1, 0)
tboard.on_end(state)
def test_write_graph(self, mock_rand, mock_board):
mock_board.return_value = Mock()
mock_board.return_value.add_graph = Mock()
mock_rand.return_value = 1
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.X: torch.zeros(1, 1, 9, 9)}
tboard = TensorBoard(write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_sample(state)
mock_rand.assert_called_once_with(state[torchbearer.X].size(), requires_grad=False)
mock_board.return_value.add_graph.assert_called_once()
self.assertEqual(str(state[torchbearer.MODEL]), str(mock_board.return_value.add_graph.call_args_list[0][0][0]))
self.assertNotEqual(state[torchbearer.MODEL], mock_board.return_value.add_graph.call_args_list[0][0][0])
def test_write_graph(self, mock_rand, mock_board, _):
mock_board.return_value = Mock()
mock_board.return_value.add_graph = Mock()
mock_rand.return_value = 1
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.X: torch.zeros(1, 1, 9, 9)}
tboard = TensorBoard(write_epoch_metrics=False)
tboard.on_start(state)
tboard.on_sample(state)
tboard.on_end(state)
mock_rand.assert_called_once_with(state[torchbearer.X].size(), requires_grad=False)
self.assertEqual(mock_board.return_value.add_graph.call_count, 1)
self.assertEqual(str(state[torchbearer.MODEL]), str(mock_board.return_value.add_graph.call_args_list[0][0][0]))
self.assertNotEqual(state[torchbearer.MODEL], mock_board.return_value.add_graph.call_args_list[0][0][0])
def test_visdom_error(self, mock_visdom, mock_writer, _, mock_warn):
import unittest.mock
orig_import = __import__
def import_mock(name, *args):
if name == 'visdom':
raise ImportError('Test error')
return orig_import(name, *args)
with unittest.mock.patch('builtins.__import__', side_effect=import_mock):
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))}
tboard = TensorBoard(visdom=True, write_epoch_metrics=False)
tboard.on_start(state)
self.assertEqual(mock_warn.call_count, 1)
