From 6040c7ac6fd47edeb9b3530e0e13ac1542069378 Mon Sep 17 00:00:00 2001 From: Tyler Yep Date: Sat, 12 Jun 2021 19:00:08 -0700 Subject: [PATCH] v0.1.5 --- .pre-commit-config.yaml | 2 +- README.md | 31 +++++++++++++++++-------------- torchinfo/layer_info.py | 2 +- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e52753e..62aa81e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -13,7 +13,7 @@ repos: - id: isort - repo: https://github.com/psf/black - rev: 21.5b2 + rev: 21.6b0 hooks: - id: black args: [-C] diff --git a/README.md b/README.md index eb4322a..e615ed6 100644 --- a/README.md +++ b/README.md @@ -218,27 +218,29 @@ summary( verbose=2, col_width=16, col_names=["kernel_size", "output_size", "num_params", "mult_adds"], + row_settings=("var_names"), ) ``` ``` ======================================================================================================================== -Layer (type:depth-idx) Kernel Shape Output Shape Param # Mult-Adds +Layer (type (var_name)) Kernel Shape Output Shape Param # Mult-Adds ======================================================================================================================== LSTMNet -- -- -- -- -├─embedding.weight [20, 300] -├─encoder.weight_ih_l0 [2048, 300] -├─encoder.weight_hh_l0 [2048, 512] -├─encoder.weight_ih_l1 [2048, 512] -├─encoder.weight_hh_l1 [2048, 512] -├─decoder.weight [20, 512] -├─Embedding: 1-1 [300, 20] [1, 100, 300] 6,000 6,000 -├─LSTM: 1-2 -- [1, 100, 512] 3,768,320 376,832,000 -│ └─weight_ih_l0 [2048, 300] -│ └─weight_hh_l0 [2048, 512] -│ └─weight_ih_l1 [2048, 512] -│ └─weight_hh_l1 [2048, 512] -├─Linear: 1-3 [512, 20] [1, 100, 20] 10,260 10,260 +├─Embedding (embedding) [300, 20] [1, 100, 300] 6,000 6,000 +│ └─weight [300, 20] 6,000 +├─LSTM (encoder) -- [1, 100, 512] 3,768,320 376,832,000 +│ └─weight_ih_l0 [2048, 300] 614,400 +│ └─weight_hh_l0 [2048, 512] 1,048,576 +│ └─bias_ih_l0 [2048] 2,048 +│ └─bias_hh_l0 [2048] 2,048 +│ └─weight_ih_l1 [2048, 512] 1,048,576 +│ └─weight_hh_l1 [2048, 512] 1,048,576 +│ └─bias_ih_l1 [2048] 2,048 +│ └─bias_hh_l1 [2048] 2,048 +├─Linear (decoder) [512, 20] [1, 100, 20] 10,260 10,260 +│ └─weight [512, 20] 10,240 +│ └─bias [20] 20 ======================================================================================================================== Total params: 3,784,580 Trainable params: 3,784,580 @@ -250,6 +252,7 @@ Forward/backward pass size (MB): 0.67 Params size (MB): 15.14 Estimated Total Size (MB): 15.80 ======================================================================================================================== + ``` diff --git a/torchinfo/layer_info.py b/torchinfo/layer_info.py index 32b8e8c..fce23e8 100644 --- a/torchinfo/layer_info.py +++ b/torchinfo/layer_info.py @@ -90,7 +90,7 @@ def nested_list_size(inputs: Sequence[Any]) -> List[int]: elif isinstance(inputs, torch.Tensor): size = list(inputs.size()) - if batch_dim is not None and len(size) > batch_dim: + if batch_dim is not None and batch_dim < len(size): size[batch_dim] = 1 elif isinstance(inputs, (list, tuple)):