Add community contributions to the README

Author: TylerYep

.pre-commit-config.yaml

@@ -1,5 +1,5 @@
 ci:
-  skip: [mypy, pylint, pytest]
+  skip: [mypy, pytest]
 repos:
   - repo: https://github.com/asottile/pyupgrade
     rev: v2.29.1
@@ -34,6 +34,12 @@ repos:
             flake8-comprehensions,
           ]
 
+  - repo: https://github.com/PyCQA/pylint
+    rev: v2.12.2
+    hooks:
+      - id: pylint
+        args: ["--disable=import-error"]
+
   - repo: local
     hooks:
       - id: mypy
@@ -43,12 +49,6 @@ repos:
         types: [python]
         require_serial: true
 
-      - id: pylint
-        name: pylint
-        entry: pylint
-        language: python
-        types: [python]
-
       - id: pytest
         name: pytest
         entry: pytest --cov=torchinfo --cov-report=html --durations=0

README.md

@@ -7,7 +7,7 @@
 [![pre-commit.ci status](https://results.pre-commit.ci/badge/github/TylerYep/torchinfo/main.svg)](https://results.pre-commit.ci/latest/github/TylerYep/torchinfo/main)
 [![GitHub license](https://img.shields.io/github/license/TylerYep/torchinfo)](https://github.com/TylerYep/torchinfo/blob/main/LICENSE)
 [![codecov](https://codecov.io/gh/TylerYep/torchinfo/branch/main/graph/badge.svg)](https://codecov.io/gh/TylerYep/torchinfo)
-[![Downloads](https://pepy.tech/badge/torch-summary)](https://pepy.tech/project/torch-summary)
+[![Downloads](https://pepy.tech/badge/torchinfo)](https://pepy.tech/project/torchinfo)
 
 (formerly torch-summary)
 
@@ -69,7 +69,6 @@ See `tests/jupyter_test.ipynb` for examples.
 **This version now supports:**
 
 - RNNs, LSTMs, and other recursive layers
-- Sequentials & ModuleLists
 - Branching output used to explore model layers using specified depths
 - Returns ModelStatistics object containing all summary data fields
 - Configurable rows/columns
@@ -82,6 +81,13 @@ See `tests/jupyter_test.ipynb` for examples.
 - Customizable line widths and batch dimension
 - Comprehensive unit/output testing, linting, and code coverage testing
 
+**Community Contributions:**
+
+- Sequentials & ModuleLists (thanks to @roym899)
+- Improved Mult-Add calculations (thanks to @TE-StefanUhlich, @zmzhang2000)
+- Dict/Misc input data (thanks to @e-dorigatti)
+- Pruned layer support (thanks to @MajorCarrot)
+
 # Documentation
 
 ```python

tests/fixtures/models.py

@@ -1,4 +1,5 @@
 """ fixtures/models.py """
+# pylint: disable=too-few-public-methods
 from __future__ import annotations
 
 import math

torchinfo/layer_info.py

@@ -12,15 +12,12 @@
 ]
 
 
-def rgetattr(obj: torch.nn.Module, attr: str) -> torch.Tensor:
-    """Get the tensor submodule called attr from obj."""
+def rgetattr(module: torch.nn.Module, attr: str) -> torch.Tensor:
+    """Get the tensor submodule called attr from module."""
     for attr_i in attr.split("."):
-        obj = getattr(obj, attr_i)
-
-    if isinstance(obj, torch.Tensor):
-        return obj
-    else:
-        raise AttributeError(f"{attr} is not a tensor")
+        module = getattr(module, attr_i)
+    assert isinstance(module, torch.Tensor)
+    return module
 
 
 class LayerInfo:
@@ -127,18 +124,22 @@ def get_layer_name(self, show_var_name: bool, show_depth: bool) -> str:
                 layer_name += f"-{self.depth_index}"
         return layer_name
 
-    def __get_cur_params(self, name: str, param: torch.Tensor) -> tuple[int, str]:
+    def get_param_count(self, name: str, param: torch.Tensor) -> tuple[int, str]:
         """
-        Get count of number of params, accounting for mask
+        Get count of number of params, accounting for mask.
+
+        Masked models save parameters with the suffix "_orig" added.
+        They have a buffer ending with "_mask" which has only 0s and 1s.
+        If a mask exists, the sum of 1s in mask is number of params.
         """
-        # Masked models save the parameter with the name "_orig" added
-        # They have a buffer ending with "_mask" which has only 0s and 1s
-        if name[-4:] == "orig":
-            # If a mask exists, the sum of 1s in mask is number of params
-            # Remove "_orig" for better readability and integration
-            return int(torch.sum(rgetattr(self.module, f"{name[:-4]}mask"))), name[:-5]
-        else:
-            return param.nelement(), name
+        if name.endswith("orig"):
+            # Remove "_orig" suffix for better readability and integration
+            without_suffix = name[:-5]
+            parameter_count = int(
+                torch.sum(rgetattr(self.module, f"{without_suffix}_mask"))
+            )
+            return parameter_count, without_suffix
+        return param.nelement(), name
 
     def calculate_num_params(self) -> None:
         """
@@ -147,7 +148,7 @@ def calculate_num_params(self) -> None:
         """
         name = ""
         for name, param in self.module.named_parameters():
-            cur_params, name = self.__get_cur_params(name, param)
+            cur_params, name = self.get_param_count(name, param)
 
             self.num_params += cur_params
             if param.requires_grad:
@@ -179,7 +180,7 @@ def calculate_macs(self) -> None:
         i.e., taking the batch-dimension into account.
         """
         for name, param in self.module.named_parameters():
-            cur_params, name = self.__get_cur_params(name, param)
+            cur_params, name = self.get_param_count(name, param)
             if name in ("weight", "bias"):
                 # ignore C when calculating Mult-Adds in ConvNd
                 if "Conv" in self.class_name: