TylerYep / torchinfo

View model summaries in PyTorch!
MIT License
2.56k stars 119 forks source link

nn.Parameter + modules #142

Closed Gregwar closed 2 years ago

Gregwar commented 2 years ago

Describe the bug If using both nn.Parameter and modules, the nn.Parameter doesn't show up. However, if not using the module, the nn.Parameter appears.

To Reproduce Example 1:

class Net(th.nn.Module):
    def __init__(self):
        super().__init__()
        self.w = th.nn.Parameter(th.zeros(1024))

"""
Outputs:
=================================================================
Layer (type:depth-idx)                   Param #
=================================================================
Net                                      1,024
=================================================================
Total params: 1,024
Trainable params: 1,024
Non-trainable params: 0
=================================================================
"""

Example 2:

class Net(th.nn.Module):
    def __init__(self):
        super().__init__()
        self.w = th.nn.Parameter(th.zeros(1024))
        self.conv = th.nn.Conv2d(3, 6, 3)

"""
Outputs:
=================================================================
Layer (type:depth-idx)                   Param #
=================================================================
Net                                      --
├─Conv2d: 1-1                            168
=================================================================
Total params: 168
Trainable params: 168pi
Non-trainable params: 0
=================================================================
"""

Expected behavior I'd expect the 1024 to appear in front of Net in the second case

Desktop (please complete the following information):

TylerYep commented 2 years ago

This has been fixed in d6e5dfa and will be released in torchinfo v1.7.0.

Thank you for reporting this issue!