def test_invalid_weights_summmary(): """Test that invalid value for weights_summary raises an error.""" with pytest.raises(MisconfigurationException, match="`mode` can be None, .* got temp"): summarize(UnorderedModel, mode="temp") with pytest.raises(MisconfigurationException, match="`weights_summary` can be None, .* got temp"): Trainer(weights_summary="temp")
def test_max_depth_equals_mode_interface(): """Test summarize(model, full/top) interface mapping matches max_depth.""" model = DeepNestedModel() summary_top = summarize(model, mode="top") summary_0 = summarize(model, max_depth=1) assert str(summary_top) == str(summary_0) summary_full = summarize(model, mode="full") summary_minus1 = summarize(model, max_depth=-1) assert str(summary_full) == str(summary_minus1)
def test_empty_model_summary_shapes(mode: str): """Test that the summary works for models that have no submodules.""" model = EmptyModule() summary = summarize(model, mode=mode) assert summary.in_sizes == [] assert summary.out_sizes == [] assert summary.param_nums == []
def test_summary_layer_types(mode): """Test that the summary displays the layer names correctly.""" model = UnorderedModel() summary = summarize(model, mode=mode) assert summary.layer_types == [ "Linear", "Linear", "Linear", "ReLU", "Conv2d" ]
def test_linear_model_summary_shapes(device, max_depth): """Test that the model summary correctly computes the input- and output shapes.""" model = UnorderedModel().to(device) model.train() summary = summarize(model, max_depth=max_depth) assert summary.in_sizes == [[2, 10], [2, 7], [2, 3], [2, 7], UNKNOWN_SIZE] # layer 2 # combine # layer 1 # relu assert summary.out_sizes == [[2, 2], [2, 9], [2, 5], [2, 7], UNKNOWN_SIZE] # layer 2 # combine # layer 1 # relu assert model.training assert model.device == device
def test_model_size_precision(tmpdir): """Test model size for half and full precision.""" model = PreCalculatedModel() # fit model trainer = Trainer(default_root_dir=tmpdir, gpus=1, max_steps=1, max_epochs=1, precision=32) trainer.fit(model) summary = summarize(model) assert model.pre_calculated_model_size == summary.model_size
def test_summary_data_output(example_input): """Ensure all items are converted to strings when getting summary data.""" class TestModel(BoringModel): @property def example_input_array(self) -> Any: return example_input summary = summarize(TestModel()) summary_data = summary._get_summary_data() for column_name, entries in summary_data: assert all(isinstance(entry, str) for entry in entries)
def test_summary_parameter_count(mode): """Test that the summary counts the number of parameters in every submodule.""" model = UnorderedModel() summary = summarize(model, mode=mode) assert summary.param_nums == [ model.layer2.weight.numel() + model.layer2.bias.numel(), model.combine.weight.numel() + model.combine.bias.numel(), model.layer1.weight.numel() + model.layer1.bias.numel(), 0, # ReLU model.unused.weight.numel() + model.unused.bias.numel(), ]
def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: if not self._max_depth: return None model_summary = summarize(pl_module, max_depth=self._max_depth) summary_data = model_summary._get_summary_data() total_parameters = model_summary.total_parameters trainable_parameters = model_summary.trainable_parameters model_size = model_summary.model_size if trainer.is_global_zero: self.summarize(summary_data, total_parameters, trainable_parameters, model_size)
def test_rnn_summary_shapes(max_depth): """Test that the model summary works for RNNs.""" model = ParityModuleRNN() b = 3 t = 5 i = model.rnn.input_size h = model.rnn.hidden_size o = model.linear_out.out_features model.example_input_array = torch.zeros(b, t, 10) summary = summarize(model, max_depth=max_depth) assert summary.in_sizes == [[b, t, i], [b, t, h]] # rnn # linear assert summary.out_sizes == [[[b, t, h], [[1, b, h], [1, b, h]]], [b, t, o]] # rnn # linear
def test_example_input_array_types(example_input, expected_size, mode): """Test the types of example inputs supported for display in the summary.""" class DummyModule(nn.Module): def forward(self, *args, **kwargs): return None class DummyLightningModule(LightningModule): def __init__(self): super().__init__() self.layer = DummyModule() # this LightningModule and submodule accept any type of input def forward(self, *args, **kwargs): return self.layer(*args, **kwargs) model = DummyLightningModule() model.example_input_array = example_input summary = summarize(model, mode=mode) assert summary.in_sizes == [expected_size]
def test_rich_summary_tuples(mock_table_add_row, mock_console): """Ensure that tuples are converted into string, and print is called correctly.""" model_summary = RichModelSummary() class TestModel(BoringModel): @property def example_input_array(self) -> Any: return torch.randn(4, 32) model = TestModel() summary = summarize(model) summary_data = summary._get_summary_data() model_summary.summarize(summary_data=summary_data, total_parameters=1, trainable_parameters=1, model_size=1) # ensure that summary was logged + the breakdown of model parameters assert mock_console.call_count == 2 # assert that the input summary data was converted correctly args, kwargs = mock_table_add_row.call_args_list[0] assert args[1:] == ("0", "layer", "Linear", "66 ", "[4, 32]", "[4, 2]")
def summary(self, max_depth=-1): print(summarize(self.model, max_depth=max_depth))
def test_raise_invalid_max_depth_value(max_depth): with pytest.raises( ValueError, match=f"`max_depth` can be -1, 0 or > 0, got {max_depth}"): summarize(DeepNestedModel(), max_depth=max_depth)
def test_empty_model_size(mode): """Test empty model size is zero.""" model = EmptyModule() summary = summarize(model, mode=mode) assert 0.0 == summary.model_size
def test_model_size(mode): """Test model size is calculated correctly.""" model = PreCalculatedModel() summary = summarize(model, mode=mode) assert model.pre_calculated_model_size == summary.model_size
def test_summary_with_scripted_modules(mode): model = PartialScriptModel() summary = summarize(model, mode=mode) assert summary.layer_types == ["RecursiveScriptModule", "Linear"] assert summary.in_sizes == [UNKNOWN_SIZE, [2, 3]] assert summary.out_sizes == [UNKNOWN_SIZE, [2, 2]]
def test_mixed_dtype_model_summary(): """Test that the model summary works with models that have mixed input- and parameter dtypes.""" model = MixedDtypeModel() summary = summarize(model) assert summary.in_sizes == [[2, 3], [2, 3, 20]] # embed # reduce assert summary.out_sizes == [[2, 3, 20], [2, 3, 1]] # embed # reduce