I have a layer layer
in an nn.Module
and use it two or more times during a single forward
step. The output of this layer
is later inputted to the same layer
. Can pytorch's autograd
compute the grad of the weights of this layer correctly?
def forward(x):x = self.layer(x)x = self.layer(x)return x
Complete example:
import torch
import torch.nn as nn
import torch.nn.functional as Fclass net(nn.Module):def __init__(self,in_dim,out_dim):super(net,self).__init__()self.layer = nn.Linear(in_dim,out_dim,bias=False)def forward(self,x):x = self.layer(x)x = self.layer(x)return xinput_x = torch.tensor([10.])
label = torch.tensor([5.])
n = net(1,1)
loss_fn = nn.MSELoss()out = n(input_x)
loss = loss_fn(out,label)
n.zero_grad()
loss.backward()for param in n.parameters():w = param.item()g = param.gradprint('Input = %.4f; label = %.4f'%(input_x,label))
print('Weight = %.4f; output = %.4f'%(w,out))
print('Gradient w.r.t. the weight is %.4f'%(g))
print('And it should be %.4f'%(4*(w**2*input_x-label)*w*input_x))
Output:
Input = 10.0000; label = 5.0000
Weight = 0.9472; output = 8.9717
Gradient w.r.t. the weight is 150.4767
And it should be 150.4766
In this example, I have defined a module with only one linear layer (in_dim=out_dim=1
and no bias). w
is the weight of this layer; input_x
is the input value; label
is the desired value. Since the loss is chosen as MSE, the formula for the loss is
((w^2)*input_x-label)^2
Computing by hand, we have
dw/dx = 2*((w^2)*input_x-label)*(2*w*input_x)
The output of my example above shows that autograd
gives the same result as computed by hand, giving me a reason to believe that it can work in this case. But in a real application, the layer may have inputs and outputs of higher dimensions, a nonlinear activation function after it, and the neural network could have multiple layers.
What I want to ask is: can I trust autograd
to handle such situation, but a lot more complicated than that in my example? How does it work when a layer is called iteratively?