Parameters added to pre-trained model have none grads

username: MStypulkowski
cooked: <p>Hi everybody,</p>
<p>I’ve trained a model on 3D data and now I’d like to add learnable rotation matrix and apply it just after taking new batch (with frozen model parameters).</p>
<p>This is my Rotation_Matrix class:</p>
<pre><code class="lang-auto">class Rotation_Matrix(nn.Module):
    def __init__(self):
        super(Rotation_Matrix, self).__init__()
        self.a = nn.Parameter(torch.tensor(1.)) # init value of Rx cosine
        self.b = nn.Parameter(torch.tensor(1.)) # init value of Ry cosine
        self.c = nn.Parameter(torch.tensor(1.)) # init value of Rz cosine
        
    def get_rotation_matrix(self):
        Rx = torch.tensor([[1., 0., 0.],
                          [0., self.a, -torch.sqrt(1 - torch.pow(self.a,2))],
                          [0., torch.sqrt(1 - torch.pow(self.a,2)), self.a]])

        Ry = torch.tensor([[self.b, 0., torch.sqrt(1 - torch.pow(self.b,2))],
                          [0., 1., 0.],
                          [-torch.sqrt(1 - torch.pow(self.b,2)), 0, self.b]])

        Rz = torch.tensor([[self.c, -torch.sqrt(1 - torch.pow(self.c,2)), 0.],
                          [torch.sqrt(1 - torch.pow(self.c,2)), self.c, 0.],
                          [0., 0., 1.]])

        return torch.mm(Rx, torch.mm(Ry, Rz))
    
    def forward(self, x):
        self.matrix = self.get_rotation_matrix()
        return torch.mm(x, self.matrix)
</code></pre>
<p>Then initialisation and training of rotation matrix. I haven’t taken care of forcing rotation’s parameters to stay in [-1,1] interval just yet.</p>
<pre><code class="lang-auto"># load pre-trained model
model = model.load_state_dict(...)
model.eval()

rotation_matrix = Rotation_Matrix()
optimizer = torch.optim.Adam(rotation_matrix.parameters(), lr=l_rate)

rotation_matrix.train()

for i in range(n_epochs):
    for j, x in enumerate(dataloader):
        x = rotation_matrix(x)
        x = model(x)
        
        loss = loss_fun(x)
        
        optimizer.zero_grad()
        rotation_matrix.zero_grad()

        loss.backward(retain_graph=True)
        
</code></pre>
<p>Printing grad of any parameter from rotation_matrix gives none value.<br>
I read some related topics but I didn’t find solution. I tried using ‘retain_grad()’ on parameters before calling ‘loss.backward()’ and I played with ‘autograd.grad()’ function but with no results. What am I missing here? In addition, do I use ‘with torch.no_grad():’ on the ‘x = model(x)’ or it wont allow to backpropagate rotation’s parameters?</p>
<p>Thank you for all answers,<br>
MS.</p>

username: Diego
cooked: <p>Hi, So you are missing one important part which is calling <code>optimizer.step()</code> after <code>loss.backward()</code>. The <code>.backward</code> function merely computes the gradients of each parameter with respect to the loss, it is the <code>.step()</code> function that actually applies the updates to the parameters.  Also doing <code>optimizer.zero_grad()</code> will set the gradients to 0, there is no need to call <code>rotation_matrix.zero_grad()</code>  after it.</p>

username: MStypulkowski
cooked: <p>Hi Diego,</p>
<p>Of course, I forgot this line. However, still any gradients aren’t computed. As I mentioned in the first post, the problem is with None values of parameters’ grads after  calling <code>.backward</code> function, not actually with applying the updates with correctly computed gradients (which is done by <code>optimizer.step()</code>).</p>

username: ptrblck
cooked: <p><code>self.matrix</code> will be a plain <code>tensor</code>, not an <code>nn.Parameter</code>, if you wrap just single <code>nn.Parameters</code> in another tensor.<br>
You can check if by calling <code>print(rotation_matrix.matrix)</code>.</p>
<p>If you create the whole <code>matrix</code> as an <code>nn.Parameter</code>, the code should work:</p>
<pre><code class="lang-python">class Rotation_Matrix(nn.Module):
    def __init__(self):
        super(Rotation_Matrix, self).__init__()
        self.a = torch.tensor(1.) # init value of Rx cosine
        self.b = torch.tensor(1.) # init value of Ry cosine
        self.c = torch.tensor(1.) # init value of Rz cosine
        self.matrix = self.get_rotation_matrix()
        
    def get_rotation_matrix(self):
        Rx = torch.tensor([[1., 0., 0.],
                          [0., self.a, -torch.sqrt(1 - torch.pow(self.a,2))],
                          [0., torch.sqrt(1 - torch.pow(self.a,2)), self.a]])

        Ry = torch.tensor([[self.b, 0., torch.sqrt(1 - torch.pow(self.b,2))],
                          [0., 1., 0.],
                          [-torch.sqrt(1 - torch.pow(self.b,2)), 0, self.b]])

        Rz = torch.tensor([[self.c, -torch.sqrt(1 - torch.pow(self.c,2)), 0.],
                          [torch.sqrt(1 - torch.pow(self.c,2)), self.c, 0.],
                          [0., 0., 1.]])

        return nn.Parameter(torch.mm(Rx, torch.mm(Ry, Rz)))
    
    def forward(self, x):
        return torch.mm(x, self.matrix)
    
model = nn.Linear(3, 1)
model.eval()

rotation_matrix = Rotation_Matrix()
optimizer = torch.optim.Adam(rotation_matrix.parameters(), lr=1.0)

rotation_matrix.train()

x = torch.randn(1, 3)
target = torch.randn(1, 1)
criterion = nn.MSELoss()

for i in range(10):
    optimizer.zero_grad()

    output = rotation_matrix(x)
    output = model(output)    
    loss = criterion(output, target)

    loss.backward()
    print(rotation_matrix.matrix.grad)
    optimizer.step()
</code></pre>

username: MStypulkowski
cooked: <p>Thanks for the answer.</p>
<p>After changes you suggested, gradients are computed correctly for every element in the matrix. This cause another problem. After update of weights the matrix is no longer a rotation matrix (orthogonal and with determinant equal to 1). I want to remain its structure so that after each iteration it would be still a rotation matrix, just like in constructor. So I’d like to update only parameters <code>a, b, c</code>.</p>
<p>In addition, next question came up. Where do I restrict <code>a, b, c</code> to stay in the [-1,1] segment to force these parameters to act like cosine functions? I thought of something like:</p>
<pre><code class="lang-auto">rotation_matrix.a = torch.max(torch.tensor(-1.), torch.min(a, torch.tensor(1.)))
</code></pre>
<p>but have no idea what place in the code is good for it.</p>