def test_writer(self, mock_is_primary_func: mock.MagicMock) -> None: """ Tests that the tensorboard writer writes the correct scalars to SummaryWriter iff is_primary() is True. """ for phase_idx, master in product([0, 1, 2], [True, False]): train, phase_type = ((True, "train") if phase_idx % 2 == 0 else (False, "test")) mock_is_primary_func.return_value = master # set up the task and state config = get_test_task_config() config["dataset"]["train"]["batchsize_per_replica"] = 2 config["dataset"]["test"]["batchsize_per_replica"] = 5 task = build_task(config) task.prepare() task.advance_phase() task.phase_idx = phase_idx task.train = train losses = [1.23, 4.45, 12.3, 3.4] sample_fetch_times = [1.1, 2.2, 3.3, 2.2] summary_writer = SummaryWriter(self.base_dir) # create a spy on top of summary_writer summary_writer = mock.MagicMock(wraps=summary_writer) # create a loss lr tensorboard hook tensorboard_plot_hook = TensorboardPlotHook(summary_writer) # run the hook in the correct order tensorboard_plot_hook.on_phase_start(task) # test tasks which do not pass the sample_fetch_times as well disable_sample_fetch_times = phase_idx == 0 for loss, sample_fetch_time in zip(losses, sample_fetch_times): task.losses.append(loss) step_data = ({} if disable_sample_fetch_times else { "sample_fetch_time": sample_fetch_time }) task.last_batch = LastBatchInfo(None, None, None, None, step_data) tensorboard_plot_hook.on_step(task) tensorboard_plot_hook.on_phase_end(task) if master: # add_scalar() should have been called with the right scalars if train: learning_rate_key = f"Learning Rate/{phase_type}" summary_writer.add_scalar.assert_any_call( learning_rate_key, mock.ANY, global_step=mock.ANY, walltime=mock.ANY, ) avg_loss_key = f"Losses/{phase_type}" summary_writer.add_scalar.assert_any_call(avg_loss_key, mock.ANY, global_step=mock.ANY) for meter in task.meters: for name in meter.value: meter_key = f"Meters/{phase_type}/{meter.name}/{name}" summary_writer.add_scalar.assert_any_call( meter_key, mock.ANY, global_step=mock.ANY) if step_data: summary_writer.add_scalar.assert_any_call( f"Speed/{phase_type}/cumulative_sample_fetch_time", mock.ANY, global_step=mock.ANY, walltime=mock.ANY, ) else: # add_scalar() shouldn't be called since is_primary() is False summary_writer.add_scalar.assert_not_called() summary_writer.add_scalar.reset_mock()
def test_writer(self, mock_is_master_func: mock.MagicMock) -> None: """ Tests that the tensorboard writer writes the correct scalars to SummaryWriter iff is_master() is True. """ for phase_idx, master in product([0, 1, 2], [True, False]): train, phase_type = ((True, "train") if phase_idx % 2 == 0 else (False, "test")) mock_is_master_func.return_value = master # set up the task and state config = get_test_task_config() config["dataset"]["train"]["batchsize_per_replica"] = 2 config["dataset"]["test"]["batchsize_per_replica"] = 5 task = build_task(config) task.prepare() task.phase_idx = phase_idx task.train = train losses = [1.23, 4.45, 12.3, 3.4] local_variables = {} summary_writer = SummaryWriter(self.base_dir) # create a spy on top of summary_writer summary_writer = mock.MagicMock(wraps=summary_writer) # create a loss lr tensorboard hook tensorboard_plot_hook = TensorboardPlotHook(summary_writer) # test that the hook logs a warning and doesn't write anything to # the writer if on_phase_start() is not called for initialization # before on_update() is called. with self.assertLogs() as log_watcher: tensorboard_plot_hook.on_update(task, local_variables) self.assertTrue( len(log_watcher.records) == 1 and log_watcher.records[0].levelno == logging.WARN and "learning_rates is not initialized" in log_watcher.output[0]) # test that the hook logs a warning and doesn't write anything to # the writer if on_phase_start() is not called for initialization # if on_phase_end() is called. with self.assertLogs() as log_watcher: tensorboard_plot_hook.on_phase_end(task, local_variables) self.assertTrue( len(log_watcher.records) == 1 and log_watcher.records[0].levelno == logging.WARN and "learning_rates is not initialized" in log_watcher.output[0]) summary_writer.add_scalar.reset_mock() # run the hook in the correct order tensorboard_plot_hook.on_phase_start(task, local_variables) for loss in losses: task.losses.append(loss) tensorboard_plot_hook.on_update(task, local_variables) tensorboard_plot_hook.on_phase_end(task, local_variables) if master: # add_scalar() should have been called with the right scalars if train: loss_key = f"{phase_type}_loss" learning_rate_key = f"{phase_type}_learning_rate_updates" summary_writer.add_scalar.assert_any_call( loss_key, mock.ANY, global_step=mock.ANY, walltime=mock.ANY) summary_writer.add_scalar.assert_any_call( learning_rate_key, mock.ANY, global_step=mock.ANY, walltime=mock.ANY, ) avg_loss_key = f"avg_{phase_type}_loss" summary_writer.add_scalar.assert_any_call(avg_loss_key, mock.ANY, global_step=mock.ANY) for meter in task.meters: for name in meter.value: meter_key = f"{phase_type}_{meter.name}_{name}" summary_writer.add_scalar.assert_any_call( meter_key, mock.ANY, global_step=mock.ANY) else: # add_scalar() shouldn't be called since is_master() is False summary_writer.add_scalar.assert_not_called() summary_writer.add_scalar.reset_mock()