Repository: Lightning-AI/pytorch-lightning
Pull request: Adding non-layer param count to summary
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://github.com/Lightning-AI/lightning/issues/12736)
-->
This PR adds the count of parameters not associated to any layer to the model summary. The new model summary should add an extra row to the summary table that specifies the count of params that have not been reported by any layer summary.

Fixes #12736

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [ ] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


Problem statement:
ModelSummary ignores nn.Parameter
## 🐛 Bug

I don't know if this is a bug of a feature, but currently ModelSummary only reports member `nn.Module`s and ignores `nn.Parameter`s that are directly created. I believe these parameters are an important part (and sometimes, a big part in terms of the number of weights) of the model so that they should be reported.

cc @borda @kaushikb11 @awaelchli @rohitgr7

Hints from issue discussion:
Hey @ZhaofengWu 

Happy to see you find the model summary useful! 
I looked at this very briefly now. With the following code I get:

```py
import torch

from pytorch_lightning import LightningModule
from pytorch_lightning.utilities.model_summary import ModelSummary

class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.param = torch.nn.Parameter(torch.rand(2, 2))  # 2 * 2 = 4
        self.layer = torch.nn.Linear(2, 2)  # 2 * 2 + 2 = 6


if __name__ == "__main__":
    model = BoringModel()
    summary = ModelSummary(model)
    print(summary)
```

```
  | Name  | Type   | Params
---------------------------------
0 | layer | Linear | 6     
---------------------------------
10        Trainable params
0         Non-trainable params
10        Total params
0.000     Total estimated model params size (MB)

```

We can see that the correct amount of parameters are getting reported (10 trainable parameters). Perhaps the size in MB is wrong. 

Are you suggesting to include the nn.Parameter as a separate row in the table? 

> Perhaps the size in MB is wrong.

That's probably just a rounding thing?

> Are you suggesting to include the nn.Parameter as a separate row in the table?

Yes. It's a bit confusing that the individual param counts don't add up to the total number as well.
But definitely not a huge issue and probably has low priority.
This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

Commenting for activity
I believe this problem still exists in pl version 1.6.4

```python
import torch
import torch.nn as nn
import pytorch_lightning as pl
from pytorch_lightning.utilities.model_summary import ModelSummary

class Example(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.lin1 = nn.Linear(32,64)
        self.lin2 = nn.Parameter(torch.rand([32,64]))
model = Example()
print(ModelSummary(model,max_depth=-1))
print(pl.__version__)
```

shows only lin1 but not lin2

```
  | Name | Type   | Params
--------------------------------
0 | lin1 | Linear | 2.1 K 
--------------------------------
4.2 K     Trainable params
0         Non-trainable params
4.2 K     Total params
0.017     Total estimated model params size (MB)
1.6.4
```
This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

Commenting for activity
Hi! @ZhaofengWu [ZhaofengWu](https://github.com/ZhaofengWu) 
I want to be assigned to this issue. It will be my first contribution.

And to make sure, I have a question.
When the code is below, is it the proper result?

``` python
import torch

from pytorch_lightning import LightningModule
from pytorch_lightning.utilities.model_summary import ModelSummary

class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.param = torch.nn.Parameter(torch.rand(2, 2))  # 2 * 2 = 4
        self.layer = torch.nn.Linear(2, 2)  # 2 * 2 + 2 = 6


if __name__ == "__main__":
    model = BoringModel()
    summary = ModelSummary(model)
    print(summary)
```

```
  | Name  | Type   | Params
---------------------------------
0 | param | Parameter | 4     
1 | layer | Linear    | 6     
---------------------------------
10        Trainable params
0         Non-trainable params
10        Total params
0.000     Total estimated model params size (MB)
```
Yes, I consider this behavior to be expected.

Commit message: remove commented code

Patch to review:
diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary.py b/src/lightning/pytorch/utilities/model_summary/model_summary.py
index 15437ec913d3b..9f960594a114e 100644
--- a/src/lightning/pytorch/utilities/model_summary/model_summary.py
+++ b/src/lightning/pytorch/utilities/model_summary/model_summary.py
@@ -32,6 +32,8 @@
 
 PARAMETER_NUM_UNITS = [" ", "K", "M", "B", "T"]
 UNKNOWN_SIZE = "?"
+LEFTOVER_PARAMS_NAME = "other params"
+NOT_APPLICABLE = "n/a"
 
 
 class LayerSummary:
@@ -141,6 +143,9 @@ class ModelSummary:
     intermediate input- and output shapes of all layers. Supported are tensors and
     nested lists and tuples of tensors. All other types of inputs will be skipped and show as `?`
     in the summary table. The summary will also display `?` for layers not used in the forward pass.
+    If there are parameters not associated with any layers or modules, the count of those parameters
+    will be displayed in the table under `other params`. The summary will display `n/a` for module type,
+    in size, and out size.
 
     Example::
 
@@ -235,6 +240,10 @@ def trainable_parameters(self) -> int:
             p.numel() if not _is_lazy_weight_tensor(p) else 0 for p in self._model.parameters() if p.requires_grad
         )
 
+    @property
+    def total_layer_params(self) -> int:
+        return sum(self.param_nums)
+
     @property
     def model_size(self) -> float:
         return self.total_parameters * self._precision_megabytes
@@ -292,8 +301,24 @@ def _get_summary_data(self) -> List[Tuple[str, List[str]]]:
             arrays.append(("In sizes", [str(x) for x in self.in_sizes]))
             arrays.append(("Out sizes", [str(x) for x in self.out_sizes]))
 
+        total_leftover_params = self.total_parameters - self.total_layer_params
+        if total_leftover_params > 0:
+            self._add_leftover_params_to_summary(arrays, total_leftover_params)
+
         return arrays
 
+    def _add_leftover_params_to_summary(self, arrays: List[Tuple[str, List[str]]], total_leftover_params: int):
+        """Add summary of params not associated with module or layer to model summary."""
+        layer_summaries = dict(arrays)
+        layer_summaries[" "].append(str(len(self._layer_summary) + 1))
+        layer_summaries["Name"].append(LEFTOVER_PARAMS_NAME)
+        layer_summaries["Type"].append(NOT_APPLICABLE)
+        layer_summaries["Params"].append(get_human_readable_count(total_leftover_params))
+        if "In sizes" in layer_summaries:
+            layer_summaries["In sizes"].append(NOT_APPLICABLE)
+        if "Out sizes" in layer_summaries:
+            layer_summaries["Out sizes"].append(NOT_APPLICABLE)
+
     def __str__(self) -> str:
         arrays = self._get_summary_data()
 
diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py b/src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py
index aac2c91ea4d1d..e70ae59c55d7e 100644
--- a/src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py
+++ b/src/lightning/pytorch/utilities/model_summary/model_summary_deepspeed.py
@@ -25,6 +25,7 @@
     get_human_readable_count,
     LayerSummary,
     ModelSummary,
+    NOT_APPLICABLE,
 )
 
 
@@ -96,4 +97,14 @@ def _get_summary_data(self) -> List[Tuple[str, List[str]]]:
             arrays.append(("In sizes", [str(x) for x in self.in_sizes]))
             arrays.append(("Out sizes", [str(x) for x in self.out_sizes]))
 
+        total_leftover_params = self.total_parameters - self.total_layer_params
+        if total_leftover_params > 0:
+            self._add_leftover_params_to_summary(arrays, total_leftover_params)
+
         return arrays
+
+    def _add_leftover_params_to_summary(self, arrays: List[Tuple[str, List[str]]], total_leftover_params: int):
+        """Add summary of params not associated with module or layer to model summary."""
+        super()._add_leftover_params_to_summary(arrays, total_leftover_params)
+        layer_summaries = dict(arrays)
+        layer_summaries["Params per Device"].append(NOT_APPLICABLE)
diff --git a/tests/tests_pytorch/utilities/test_model_summary.py b/tests/tests_pytorch/utilities/test_model_summary.py
index 7bb4ccf6258be..da201c497c18b 100644
--- a/tests/tests_pytorch/utilities/test_model_summary.py
+++ b/tests/tests_pytorch/utilities/test_model_summary.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from collections import OrderedDict
 from typing import Any
 
 import pytest
@@ -19,7 +20,13 @@
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.demos.boring_classes import BoringModel
-from lightning.pytorch.utilities.model_summary.model_summary import ModelSummary, summarize, UNKNOWN_SIZE
+from lightning.pytorch.utilities.model_summary.model_summary import (
+    LEFTOVER_PARAMS_NAME,
+    ModelSummary,
+    NOT_APPLICABLE,
+    summarize,
+    UNKNOWN_SIZE,
+)
 from tests_pytorch.helpers.advanced_models import ParityModuleRNN
 from tests_pytorch.helpers.runif import RunIf
 
@@ -137,6 +144,18 @@ def forward(self, inp):
         return self.head(self.branch1(inp), self.branch2(inp))
 
 
+class NonLayerParamsModel(LightningModule):
+    """A model with parameters not associated with pytorch layer."""
+
+    def __init__(self):
+        super().__init__()
+        self.param = torch.nn.Parameter(torch.ones(2, 2))
+        self.layer = torch.nn.Linear(2, 2)
+
+    def forward(self, inp):
+        self.layer(self.param @ inp)
+
+
 def test_invalid_max_depth():
     """Test that invalid value for max_depth raises an error."""
     model = LightningModule()
@@ -358,3 +377,24 @@ def example_input_array(self) -> Any:
     summary_data = summary._get_summary_data()
     for column_name, entries in summary_data:
         assert all(isinstance(entry, str) for entry in entries)
+
+
+@pytest.mark.parametrize("example_input", [None, torch.ones(2, 2)])
+def test_summary_data_with_non_layer_params(example_input):
+    model = NonLayerParamsModel()
+    model.example_input_array = example_input
+
+    summary = summarize(model)
+    summary_data = OrderedDict(summary._get_summary_data())
+    assert summary_data["Name"][-1] == LEFTOVER_PARAMS_NAME
+    assert summary_data["Type"][-1] == NOT_APPLICABLE
+    assert int(summary_data["Params"][-1]) == 4
+    if example_input is not None:
+        assert summary_data["In sizes"][-1] == NOT_APPLICABLE
+        assert summary_data["Out sizes"][-1] == NOT_APPLICABLE
+
+
+def test_summary_data_with_no_non_layer_params():
+    summary = summarize(PreCalculatedModel())
+    summary_data = OrderedDict(summary._get_summary_data())
+    assert summary_data["Name"][-1] != LEFTOVER_PARAMS_NAME


Code review comments:
diff_hunk: @@ -292,8 +301,24 @@ def _get_summary_data(self) -> List[Tuple[str, List[str]]]:
             arrays.append(("In sizes", [str(x) for x in self.in_sizes]))
             arrays.append(("Out sizes", [str(x) for x in self.out_sizes]))
 
+        total_leftover_params = self.total_parameters - self.total_layer_params
+        if total_leftover_params > 0:
+            self._add_leftover_params_to_summary(arrays, total_leftover_params)
+
         return arrays
 
+    def _add_leftover_params_to_summary(self, arrays: List[Tuple[str, List[str]]], total_leftover_params: int):
+        """Add summary of params not associated with module or layer to model summary."""
+        layer_summaries = dict(arrays)
+        layer_summaries[" "].append(str(len(self._layer_summary) + 1))
path: src/lightning/pytorch/utilities/model_summary/model_summary.py
text: @user1:
Pretty sure this shouldn't be +1 here. I'm getting non-consecutive numbering here:


```
  | Name         | Type   | Params
----------------------------------------
0 | layer        | Linear | 1.1 K 
1 | layer1       | Linear | 1.1 K 
2 | layer2       | Linear | 66    
4 | other params | n/a    | 3     
----------------------------------------
```

On another note, I would not even put a number in the first place. The layers appear there in the order they were recorded during use in forward. But the "other params" is really not something we can record in order so they are just there at the end as a separate row unrelated to the ordering of others. What do you think? 

@author:
Sounds good. I agree that `other params` should not have an index in the summary.