Skip to content

Commit

Permalink
Add the "groups" column. (#320)
Browse files Browse the repository at this point in the history
* Add the "groups" column.

Add test_groups().

Add groups field to single_input_all_cols.out

Add Groups field to parameter_list.out

* Move the GROUPS field so that it comes after KERNEL_SIZE.

This changes the order of the ColumnSettings enumeration.

---------

Co-authored-by: Andrew Lavin <[email protected]>
  • Loading branch information
andravin and AndrewLavin authored Nov 1, 2024
1 parent 6e0668a commit 29166cc
Show file tree
Hide file tree
Showing 9 changed files with 59 additions and 19 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,7 @@ Args:
"num_params",
"params_percent",
"kernel_size",
"groups",
"mult_adds",
"trainable",
)
Expand Down
15 changes: 15 additions & 0 deletions tests/test_output/groups.out
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
================================================================================================================================================================
Layer (type:depth-idx) Kernel Shape Groups Input Shape Output Shape Param # Mult-Adds
================================================================================================================================================================
Conv2d [3, 3] 4 [7, 16, 28, 28] [7, 32, 26, 26] 1,184 5,602,688
================================================================================================================================================================
Total params: 1,184
Trainable params: 1,184
Non-trainable params: 0
Total mult-adds (M): 5.60
================================================================================================================================================================
Input size (MB): 0.35
Forward/backward pass size (MB): 1.21
Params size (MB): 0.00
Estimated Total Size (MB): 1.57
================================================================================================================================================================
14 changes: 7 additions & 7 deletions tests/test_output/parameter_list.out
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
====================================================================================================================================================================================
Layer (type:depth-idx) Kernel Shape Input Shape Output Shape Param # Param % Mult-Adds Trainable
====================================================================================================================================================================================
ParameterListModel -- [100, 100] [100, 100] 30,000 100.00% -- True
====================================================================================================================================================================================
========================================================================================================================================================================================================
Layer (type:depth-idx) Kernel Shape Groups Input Shape Output Shape Param # Param % Mult-Adds Trainable
========================================================================================================================================================================================================
ParameterListModel -- -- [100, 100] [100, 100] 30,000 100.00% -- True
========================================================================================================================================================================================================
Total params: 30,000
Trainable params: 30,000
Non-trainable params: 0
Total mult-adds (M): 0
====================================================================================================================================================================================
========================================================================================================================================================================================================
Input size (MB): 0.04
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.04
====================================================================================================================================================================================
========================================================================================================================================================================================================
24 changes: 12 additions & 12 deletions tests/test_output/single_input_all_cols.out
Original file line number Diff line number Diff line change
@@ -1,20 +1,20 @@
====================================================================================================================================================================================
Layer (type:depth-idx) Kernel Shape Input Shape Output Shape Param # Param % Mult-Adds Trainable
====================================================================================================================================================================================
SingleInputNet -- [7, 1, 28, 28] [7, 10] -- -- -- True
├─Conv2d: 1-1 [5, 5] [7, 1, 28, 28] [7, 10, 24, 24] 260 1.19% 1,048,320 True
├─Conv2d: 1-2 [5, 5] [7, 10, 12, 12] [7, 20, 8, 8] 5,020 22.99% 2,248,960 True
├─Dropout2d: 1-3 -- [7, 20, 8, 8] [7, 20, 8, 8] -- -- -- --
├─Linear: 1-4 -- [7, 320] [7, 50] 16,050 73.49% 112,350 True
├─Linear: 1-5 -- [7, 50] [7, 10] 510 2.34% 3,570 True
====================================================================================================================================================================================
========================================================================================================================================================================================================
Layer (type:depth-idx) Kernel Shape Groups Input Shape Output Shape Param # Param % Mult-Adds Trainable
========================================================================================================================================================================================================
SingleInputNet -- -- [7, 1, 28, 28] [7, 10] -- -- -- True
├─Conv2d: 1-1 [5, 5] 1 [7, 1, 28, 28] [7, 10, 24, 24] 260 1.19% 1,048,320 True
├─Conv2d: 1-2 [5, 5] 1 [7, 10, 12, 12] [7, 20, 8, 8] 5,020 22.99% 2,248,960 True
├─Dropout2d: 1-3 -- -- [7, 20, 8, 8] [7, 20, 8, 8] -- -- -- --
├─Linear: 1-4 -- -- [7, 320] [7, 50] 16,050 73.49% 112,350 True
├─Linear: 1-5 -- -- [7, 50] [7, 10] 510 2.34% 3,570 True
========================================================================================================================================================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
Total mult-adds (M): 3.41
====================================================================================================================================================================================
========================================================================================================================================================================================================
Input size (MB): 0.02
Forward/backward pass size (MB): 0.40
Params size (MB): 0.09
Estimated Total Size (MB): 0.51
====================================================================================================================================================================================
========================================================================================================================================================================================================
13 changes: 13 additions & 0 deletions tests/torchinfo_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,19 @@ def test_single_input_all_cols() -> None:
)


def test_groups() -> None:
input_shape = (7, 16, 28, 28)
module = nn.Conv2d(16, 32, 3, groups=4)
col_names = ("kernel_size", "groups", "input_size", "output_size", "num_params", "mult_adds")
summary(
module,
input_data=torch.randn(*input_shape),
depth=1,
col_names=col_names,
col_width=20,
)


def test_single_input_batch_dim() -> None:
model = SingleInputNet()
col_names = ("kernel_size", "input_size", "output_size", "num_params", "mult_adds")
Expand Down
1 change: 1 addition & 0 deletions torchinfo/enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ class ColumnSettings(str, Enum):
__slots__ = ()

KERNEL_SIZE = "kernel_size"
GROUPS = "groups"
INPUT_SIZE = "input_size"
OUTPUT_SIZE = "output_size"
NUM_PARAMS = "num_params"
Expand Down
2 changes: 2 additions & 0 deletions torchinfo/formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

HEADER_TITLES = {
ColumnSettings.KERNEL_SIZE: "Kernel Shape",
ColumnSettings.GROUPS: "Groups",
ColumnSettings.INPUT_SIZE: "Input Shape",
ColumnSettings.OUTPUT_SIZE: "Output Shape",
ColumnSettings.NUM_PARAMS: "Param #",
Expand Down Expand Up @@ -107,6 +108,7 @@ def layer_info_to_row(
"""Convert layer_info to string representation of a row."""
values_for_row = {
ColumnSettings.KERNEL_SIZE: self.str_(layer_info.kernel_size),
ColumnSettings.GROUPS: self.str_(layer_info.groups),
ColumnSettings.INPUT_SIZE: self.str_(layer_info.input_size),
ColumnSettings.OUTPUT_SIZE: self.str_(layer_info.output_size),
ColumnSettings.NUM_PARAMS: layer_info.num_params_to_str(reached_max_depth),
Expand Down
7 changes: 7 additions & 0 deletions torchinfo/layer_info.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ def __init__(
self.input_size: list[int] = []
self.output_size: list[int] = []
self.kernel_size = self.get_kernel_size(module)
self.groups = self.get_groups(module)
self.trainable_params = 0
self.num_params = 0
self.param_bytes = 0
Expand Down Expand Up @@ -169,6 +170,12 @@ def get_kernel_size(module: nn.Module) -> int | list[int] | None:
return kernel_size
return None

@staticmethod
def get_groups(module: nn.Module) -> int | None:
if hasattr(module, "groups"):
return int(module.groups)
return None

def get_layer_name(self, show_var_name: bool, show_depth: bool) -> str:
layer_name = self.class_name
if show_var_name and self.var_name:
Expand Down
1 change: 1 addition & 0 deletions torchinfo/torchinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,7 @@ class name as the key. If the forward pass is an expensive operation,
"num_params",
"params_percent",
"kernel_size",
"groups",
"mult_adds",
"trainable",
)
Expand Down

0 comments on commit 29166cc

Please sign in to comment.