Best way to define a scalar using nn.Parameter in Pytorch

Viewed 2294

In my CNN at some stage I want to multiply a feature map with some scalar which should be learnt by the network. Which of the following is the best way to do it or all are same? The scalar has to be initialised to 5.

# Method 1
def __init__(self):
    super(..., self).__init__()
    ...
    ...
    alpha = nn.Parameter(5)
    ...

def forward(self, x):
    ...
    x = x * alpha
    return x
# Method 2
def __init__(self):
    super(..., self).__init__()
    ...
    ...
    alpha = nn.Parameter(torch.tensor(5))
    ...

def forward(self, x):
    ...
    x = x * alpha
    return x
# Method 3
def __init__(self):
    super(..., self).__init__()
    ...
    ...
    alpha = nn.Parameter(torch.ones(1)*5)
    ...

def forward(self, x):
    ...
    x = x * alpha
    return x

If all are same I would prefer Method 1 and let CNN learn the appropriate multiplier alpha for the feature map x. I hope in all cases alpha will be a float32 tensor initialised to 5. I am using PyTorch 1.3.1

Yours sincerely,

Mohit

1 Answers

The third option will work since the parameter constructor needs a float. Parameters are updated with the optimizer, so they need to have gradients, apart from buffers.

Buffers are managed by yourself, not with the optimizer.

You may play with this experimental code.

BS=2
class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(BS, 2))
        self.bias = nn.Parameter(torch.zeros(BS))
        self.alpha = nn.Parameter(torch.tensor(5.))

    def forward(self, x):
        return x @ self.weights + self.bias

m=M()
m.parameters()
list(m.parameters())

Out:

[Parameter containing:
 tensor([[-0.5627,  2.3067],
         [ 1.3821, -0.1893]], requires_grad=True), Parameter containing:
 tensor([0., 0.], requires_grad=True), Parameter containing:
 tensor(5., requires_grad=True)]

Here I directly set the value 5. to the parameter alpha, and I added a few more parameters for fun.

You could also do as explained by Shai:

self.register_parameter(name='alpha', param=torch.nn.Parameter(torch.tensor(5.)))

You may asked why we have nn.Module.register_parameter, why don't we just use the nn.Parameter() approach?

nn.Module.register_parameter takes the name and tensor and first checks if the name is in the dictionary of the module. While nn.Parameter() doesn't have such a check.

Related