def test_writer(self, mock_is_primary_func: mock.MagicMock) -> None: """ Tests that the tensorboard writer calls SummaryWriter with the model iff is_primary() is True. """ mock_summary_writer = mock.create_autospec(SummaryWriter, instance=True) task = get_test_classy_task() task.prepare() for primary in [False, True]: mock_is_primary_func.return_value = primary model_configs = get_test_model_configs() for model_config in model_configs: model = build_model(model_config) task.base_model = model # create a model tensorboard hook model_tensorboard_hook = ModelTensorboardHook( mock_summary_writer) model_tensorboard_hook.on_start(task) if primary: # SummaryWriter should have been init-ed with the correct # add_graph should be called once with model as the first arg mock_summary_writer.add_graph.assert_called_once() self.assertEqual( mock_summary_writer.add_graph.call_args[0][0], model) else: # add_graph shouldn't be called since is_primary() is False mock_summary_writer.add_graph.assert_not_called() mock_summary_writer.reset_mock()
def test_get_model_dummy_input(self): for config in get_test_model_configs(): model = build_model( config) # pass in a dummy model for the cuda check batchsize = 8 # input_key is list input_key = ["audio", "video"] input_shape = [[3, 40, 100], [4, 16, 223, 223]] # dummy input shapes result = util.get_model_dummy_input(model, input_shape, input_key, batchsize) self.assertEqual(result.keys(), {"audio", "video"}) for i in range(len(input_key)): self.assertEqual(result[input_key[i]].size(), tuple([batchsize] + input_shape[i])) # input_key is string input_key = "video" input_shape = [4, 16, 223, 223] result = util.get_model_dummy_input(model, input_shape, input_key, batchsize) self.assertEqual(result.keys(), {"video"}) self.assertEqual(result[input_key].size(), tuple([batchsize] + input_shape)) # input_key is None input_key = None input_shape = [4, 16, 223, 223] result = util.get_model_dummy_input(model, input_shape, input_key, batchsize) self.assertEqual(result.size(), tuple([batchsize] + input_shape))
def test_model_complexity_hook(self) -> None: model_configs = get_test_model_configs() task = get_test_classy_task() task.prepare() # create a model complexity hook model_complexity_hook = ModelComplexityHook() for model_config in model_configs: model = build_model(model_config) task.base_model = model with self.assertLogs(): model_complexity_hook.on_start(task)
def test_model_complexity(self) -> None: """ Test that the number of parameters and the FLOPs are calcuated correctly. """ model_configs = get_test_model_configs() expected_mega_flops = [4122, 4274, 106152] expected_params = [25557032, 25028904, 43009448] local_variables = {} task = get_test_classy_task() task.prepare() # create a model complexity hook model_complexity_hook = ModelComplexityHook() for model_config, mega_flops, params in zip(model_configs, expected_mega_flops, expected_params): model = build_model(model_config) task.base_model = model with self.assertLogs() as log_watcher: model_complexity_hook.on_start(task, local_variables) # there should be 2 log statements generated self.assertEqual(len(log_watcher.output), 2) # first statement - either the MFLOPs or a warning if mega_flops is not None: match = re.search( r"FLOPs for forward pass: (?P<mega_flops>[-+]?\d*\.\d+|\d+) MFLOPs", log_watcher.output[0], ) self.assertIsNotNone(match) self.assertEqual(mega_flops, float(match.group("mega_flops"))) else: self.assertIn("Model contains unsupported modules", log_watcher.output[0]) # second statement match = re.search( r"Number of parameters in model: (?P<params>[-+]?\d*\.\d+|\d+)", log_watcher.output[1], ) self.assertIsNotNone(match) self.assertEqual(params, float(match.group("params")))
def test_complexity_calculation_resnext(self) -> None: model_configs = get_test_model_configs() # make sure there are three configs returned self.assertEqual(len(model_configs), 3) # expected values which allow minor deviations from model changes # we only test at the 10^6 scale expected_m_flops = [4122, 7850, 8034] expected_m_params = [25, 44, 44] expected_m_activations = [11, 16, 21] for model_config, m_flops, m_params, m_activations in zip( model_configs, expected_m_flops, expected_m_params, expected_m_activations ): model = build_model(model_config) self.assertEqual(compute_activations(model) // 10 ** 6, m_activations) self.assertEqual(compute_flops(model) // 10 ** 6, m_flops) self.assertEqual(count_params(model) // 10 ** 6, m_params)
class TestClassyModel(unittest.TestCase): model_configs = get_test_model_configs() def _get_config(self, model_config): return { "name": "classification_task", "num_epochs": 12, "loss": { "name": "test_loss" }, "dataset": { "name": "imagenet", "batchsize_per_replica": 8, "use_pairs": False, "num_samples_per_phase": None, "use_shuffle": { "train": True, "test": False }, }, "meters": [], "model": model_config, "optimizer": { "name": "test_opt" }, } def _compare_model_state(self, state, state2): compare_model_state(self, state, state2) def test_build_model(self): for cfg in self.model_configs: config = self._get_config(cfg) model = build_model(config["model"]) self.assertTrue(isinstance(model, ClassyModel)) self.assertTrue( type(model.input_shape) == tuple and len(model.input_shape) == 3) self.assertTrue( type(model.output_shape) == tuple and len(model.output_shape) == 2) self.assertTrue(type(model.model_depth) == int) def test_get_set_state(self): config = self._get_config(self.model_configs[0]) model = build_model(config["model"]) fake_input = torch.Tensor(1, 3, 224, 224).float() model.eval() state = model.get_classy_state() with torch.no_grad(): output = model(fake_input) model2 = build_model(config["model"]) model2.set_classy_state(state) # compare the states state2 = model2.get_classy_state() self._compare_model_state(state, state2) model2.eval() with torch.no_grad(): output2 = model2(fake_input) self.assertTrue(torch.allclose(output, output2)) # test deep_copy by assigning a deep copied state to model2 # and then changing the original model's state state = model.get_classy_state(deep_copy=True) model3 = build_model(config["model"]) state3 = model3.get_classy_state() # assign model2's state to model's and also re-assign model's state model2.set_classy_state(state) model.set_classy_state(state3) # compare the states state2 = model2.get_classy_state() self._compare_model_state(state, state2) def test_get_set_head_states(self): config = copy.deepcopy(self._get_config(self.model_configs[0])) head_configs = config["model"]["heads"] config["model"]["heads"] = [] model = build_model(config["model"]) trunk_state = model.get_classy_state() heads = defaultdict(dict) for head_config in head_configs: head = build_head(head_config) heads[head_config["fork_block"]][head.unique_id] = head model.set_heads(heads) model_state = model.get_classy_state() # the heads should be the same as we set self.assertEqual(len(heads), len(model.get_heads())) for block_name, hs in model.get_heads().items(): self.assertEqual(hs, heads[block_name]) model._clear_heads() self._compare_model_state(model.get_classy_state(), trunk_state) model.set_heads(heads) self._compare_model_state(model.get_classy_state(), model_state)