How to get nn.Parameter listed when net is printed?

1 
import torch
import torch.nn as nn

class MyNet(torch.nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.layer = nn.Linear(10, 10)
        self.parameter = torch.nn.Parameter(torch.zeros(10,10, requires_grad=True))

net = MyNet()
print(net)

Ouput is

MyNet(
  (layer): Linear(in_features=10, out_features=10, bias=True)
)
1 Like
2

You could print all parameters of the model via:

print(list(net.parameters()))
# or
print(dict(net.named_parameters()))
3

I want parameters to come in this command print(net)
This is more interpretable that others

4

In that case you could override the extra_repr method for the module.

5 
class MyNet(torch.nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.layer = nn.Linear(10, 10)
        self.parameter1 = torch.nn.Parameter(torch.zeros(10,10, requires_grad=True))
        self.parameter2 = torch.nn.Parameter(torch.zeros(10,10, requires_grad=True))
        
    def extra_repr(self) -> str:
        
        named_modules = set()
        for p in net.named_modules():
            named_modules.update([p[0]]  )    
        named_modules = list(named_modules)

        string_repr = ''
        for p in net.named_parameters():
            name = p[0].split('.')[0]
            if name not in name_modules:
                string_repr = string_repr + '('+ name +'): ' \
                    +'tensor(' + str(tuple(p[1].shape))+ ', requires_grad='+ str(p[1].requires_grad) +')\n' 
        
        return string_repr

net = MyNet()
print(net)

Output is

MyNet(
  (parameter1): tensor((10, 10), requires_grad=True)
  (parameter2): tensor((10, 10), requires_grad=True)
  
  (layer): Linear(in_features=10, out_features=10, bias=True)
)

Is there a better automated version that this?