def test_batch_writer_visdom(self, mock_visdom, mock_writer, _):
tboard = TensorBoardText(visdom=True,
write_epoch_metrics=False,
write_batch_metrics=True,
log_trial_summary=False)
metrics = {'test_metric_1': 1, 'test_metric_2': 1}
state = {
torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)),
torchbearer.EPOCH: 1,
torchbearer.BATCH: 100,
torchbearer.METRICS: metrics
}
metric_string = TensorBoardText.table_formatter(str(metrics))
metric_string = '<h3>Epoch {} - Batch {}</h3>'.format(
state[torchbearer.EPOCH], state[torchbearer.BATCH]) + metric_string
tboard.on_start(state)
tboard.on_start_training(state)
tboard.on_start_epoch(state)
tboard.on_step_training(state)
mock_writer.return_value.add_text.assert_called_once_with(
'batch', metric_string, 1)
tboard.on_end_epoch(state)
tboard.on_end(state)
def test_epoch_writer(self, mock_writer, _):
tboard = TensorBoardText(log_trial_summary=False)
metrics = {'test_metric_1': 1, 'test_metric_2': 1}
state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)),
torchbearer.EPOCH: 1, torchbearer.METRICS: metrics}
metric_string = TensorBoardText.table_formatter(str(metrics))
tboard.on_start(state)
tboard.on_start_training(state)
tboard.on_start_epoch(state)
tboard.on_end_epoch(state)
mock_writer.return_value.add_text.assert_called_once_with('epoch', metric_string, 1)
tboard.on_end(state)