Search code examples
parametersmoduleneural-networkpytorchcustomization

Defining named parameters for a customized NN module in Pytorch


This question is about how to appropriately define the parameters of a customized layer in Pytorch. I am wondering how one can make the parameter to end up being named parameter?

class My_layer(torch.nn.Module):
    def __init__(self):
        self.mu = torch.nn.Parameter(torch.tensor([[0.0],[1.0]]))

So that when I print the parameters as below, the p.name is not empty.

my_layer = My_Layer()
for p in my_layer.parameters():
     print(p.name, p.data, p.requires_grad)

Solution

  • You are registering your parameter properly, but you should use nn.Module.named_parameters rather than nn.Module.parameters to access the names. Currently you are attempting to access Parameter.name, which is probably not what you want. The name attribute of Parameter and Tensor do not appear to be documented, but as far as I can tell, they refer to an internal attribute of the _C._TensorBase class and cannot be modified externally.

    Every time you assign a Parameter to an attribute of your module it is registered with a name (this occurs in nn.Module.__setattr__ here). The parameter always takes the same name as the attribute itself, so "mu" in this case. To iterate over all the parameters and their associated names use nn.Module.named_parameters. For example,

    my_layer = My_Layer()
    for n, p in my_layer.named_parameters():
        print('Parameter name:', n)
        print(p.data)
        print('requires_grad:', p.requires_grad)
    

    which prints

    Parameter name: mu
    tensor([[0.],
            [1.]])
    requires_grad: True