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
