Add hide_recursive_layers option to ColumnSettings

[mert-kurttutan]

Addresses #65 and continues the work from #66

• Recursive Option
• Change recursion criterion for modules with no parameters
  I can also some test if this seems OK.
  Test code:

class DummyRNN(nn.Module):
    def __init__(self, max_length: int):
        super(DummyRNN, self).__init__()
        self.lstm = nn.LSTMCell(8, 4)
        self.activation = nn.Tanh()
        self.max_length = max_length


    def forward(self, token_embedding):
        for i in range(self.max_length):
            predict = self.lstm(token_embedding)
            predict = self.activation(predict[0])


model = DummyRNN(7)

batch_size = 2
data_shape = (8,)
random_data = torch.rand((batch_size, *data_shape))


recursive_summary = summary(
    model, 
    input_data=[random_data], 
    row_settings=('depth', 'var_names', 'no_recursive'),
    device='cpu',
)

Before commit:

==========================================================================================
Layer (type (var_name):depth-idx)        Output Shape              Param #
==========================================================================================
DummyRNN (DummyRNN)                      --                        --
├─LSTMCell (lstm): 1-1                   [2, 4]                    224
├─Tanh (activation): 1-2                 [2, 4]                    --
==========================================================================================
Total params: 224
Trainable params: 224
Non-trainable params: 0
Total mult-adds (M): 0.01
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
==========================================================================================
(pytorch-env) mert-kurttutan@EDELHCND1419HMY:~/projects$ python try.py 
==========================================================================================
Layer (type (var_name):depth-idx)        Output Shape              Param #
==========================================================================================
DummyRNN (DummyRNN)                      --                        --
├─LSTMCell (lstm): 1-1                   [2, 4]                    224
├─Tanh (activation): 1-2                 [2, 4]                    --
├─LSTMCell (lstm): 1-3                   [2, 4]                    (recursive)
├─Tanh (activation): 1-4                 [2, 4]                    --
├─LSTMCell (lstm): 1-5                   [2, 4]                    (recursive)
├─Tanh (activation): 1-6                 [2, 4]                    --
├─LSTMCell (lstm): 1-7                   [2, 4]                    (recursive)
├─Tanh (activation): 1-8                 [2, 4]                    --
├─LSTMCell (lstm): 1-9                   [2, 4]                    (recursive)
├─Tanh (activation): 1-10                [2, 4]                    --
├─LSTMCell (lstm): 1-11                  [2, 4]                    (recursive)
├─Tanh (activation): 1-12                [2, 4]                    --
├─LSTMCell (lstm): 1-13                  [2, 4]                    (recursive)
├─Tanh (activation): 1-14                [2, 4]                    --
==========================================================================================
Total params: 224
Trainable params: 224
Non-trainable params: 0
Total mult-adds (M): 0.01
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
==========================================================================================

After commit,

==========================================================================================
Layer (type (var_name):depth-idx)        Output Shape              Param #
==========================================================================================
DummyRNN (DummyRNN)                      --                        --
├─LSTMCell (lstm): 1-1                   [2, 4]                    224
├─Tanh (activation): 1-2                 [2, 4]                    --
==========================================================================================
Total params: 224
Trainable params: 224
Non-trainable params: 0
Total mult-adds (M): 0.01
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
==========================================================================================

[codecov]

Codecov Report

Merging #174 (aa92e71) into main (19ad6dd) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #174   +/-   ##
=======================================
  Coverage   97.36%   97.36%           
=======================================
  Files           6        6           
  Lines         569      570    +1     
=======================================
+ Hits          554      555    +1     
  Misses         15       15           

Impacted Files	Coverage Δ
torchinfo/enums.py	100.00% <100.00%> (ø)
torchinfo/formatting.py	100.00% <100.00%> (ø)
torchinfo/layer_info.py	95.87% <100.00%> (-0.03%)	⬇️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[TylerYep]

The change looks good -- I renamed the name of the setting to be more explicit.

Please add the test case you showed above, as well as another test case using the recursive layers in different orders.

Currently the simplified code looks really clean, but I'm wondering if we added some other layers at the end and then tried to use the same Tanh layer, would the Tanh mysteriously disappear?

[mert-kurttutan]

Just to be sure, you mean layers after tanh within the recursion, right?

[TylerYep]

Yes, something like:

class DummyRNN(nn.Module):
    def __init__(self, max_length: int):
        super().__init__()
        self.lstm = nn.LSTMCell(8, 4)
        self.activation = nn.Tanh()
        self.max_length = max_length


    def forward(self, token_embedding):
        for i in range(self.max_length):
            predict = self.lstm(token_embedding)
            predict = self.activation(predict[0])
        self.Linear(...)
        self.activation(...)

Does the last activation disappear too?

[mert-kurttutan]

==========================================================================================
Layer (type (var_name):depth-idx)        Output Shape              Param #
==========================================================================================
DummyRNN (DummyRNN)                      --                        --
├─LSTMCell (lstm): 1-1                   [2, 4]                    224
├─Tanh (activation): 1-2                 [2, 4]                    --
├─Linear (projection): 1-15              [2, 12]                   60
==========================================================================================
Total params: 284
Trainable params: 284
Non-trainable params: 0
Total mult-adds (M): 0.01
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
==========================================================================================

Yes, it does. I guess it is more suited for recursion due to the for-loop ( at least for the use-cases I can think of).

[TylerYep]

Since this feature is opt-in, I'm fine with this behavior, but let's add this as a separate test case so we'll know if the implementation changes in the future.

[TylerYep]

Thanks for the PR, this setting should be very useful for highly recurrent models. Will be released in v1.7.2.