Trying to find an angle (use of trigonometric functions); optimizer not updating

Ingenieur: <p>Hi everybody,</p>
<p>I’m currently trying to figure out how to use PyTorch to optimize an angle representing the angular part of an axis/angle rotation of one vector set into another. That is, I have two sets of vectors of, let’s say, shape (100, 3). One is the input, the other is the target. The target is equal to the input, rotated by the given axis/angle rotation. The angle for each sample in the sets includes some unknown gaussian noise. Due to that unknown noise a simple matrix operation to solve the equation system is impossible.</p>
<p>What I have so far is an nn.Module that executes the rotation with the current angle on the input. Code looks as follows:</p>
<pre><code class="lang-auto">class AngleModel(torch.nn.Module):
    def __init__(self):
        super(AngleModel, self).__init__()
        self.angle = nn.Parameter(torch.tensor(0.0, requires_grad = True))
        self.qw = torch.cos(self.angle / 2.)
        self.qx = torch.zeros(1)
        self.qy = torch.zeros(1)
        self.qz = torch.sin(self.angle / 2.)

    def forward(self, input):
        m11 = 1. - 2. * torch.pow(self.qy, 2) - 2. * torch.pow(self.qz, 2).requires_grad_()
        m22 = 1. - 2. * torch.pow(self.qx, 2) - 2. * torch.pow(self.qz, 2).requires_grad_()
        m33 = 1. - 2. * torch.pow(self.qx, 2) - 2. * torch.pow(self.qy, 2).requires_grad_()

        m21 = 2. * self.qx * self.qy - 2. * self.qz * self.qw
        m12 = 2. * self.qx * self.qy + 2. * self.qz * self.qw

        m31 = 2. * self.qx * self.qz + 2 * self.qy * self.qw
        m13 = 2. * self.qx * self.qz - 2 * self.qy * self.qw

        m32 = 2. * self.qy * self.qz - 2. * self.qx * self.qw
        m23 = 2. * self.qy * self.qz + 2. * self.qx * self.qw

        matrix = torch.Tensor([
            [m11, m21, m31],
            [m12, m22, m32],
            [m13, m23, m33],
        ])

        output = torch.matmul(input.float(), matrix.float())
        return output
</code></pre>
<p>I initialize the model and optimizer like this:</p>
<pre><code class="lang-auto">        model = AngleModel()
        crit = torch.nn.MSELoss()
        l_rate = 0.01
        optim = torch.optim.SGD(model.parameters(), lr = l_rate)
        epochs = 100
</code></pre>
<p>and execute the training like this:</p>
<pre><code class="lang-auto">        for epoch in range(epochs):
            _x = torch.tensor(input, requires_grad = True)
            _y = torch.tensor(target)

            optim.zero_grad()
            outputs = model.forward(_x)
            loss = crit(outputs, _y)
            loss.backward()
            optim.step()

            print("loss %05.3f; %s" % (loss.item(), model.angle.item()))
</code></pre>
<p>I suspect the trigonometric and <code>pow</code> functions to be the culprit of why <code>angle</code> isn’t getting updated. Am I right? <code>angle.grad</code> is None after <code>loss.backward()</code>.</p>
<p>How would I go on about this? What do I need to do in order to get this working?</p>
<p>Thank you very much!</p>

tom: <p>You need to declare your parameters in the <code>__init__</code>, but put all computation (except for constants) in the forward pass.<br>
You shouldn’t have to do requires_grad_ in the forward.<br>
You should not use tensor or Tensor within your calculation. This will break the graph.<br>
It isn’t entirely clear to me whether you want qx and qy to be learnable in some way, too (otherwise you could just write them as <code>0</code> rather than using <code>torch.zeros</code>. You could leave those in the <code>__init__</code> if they’re constant.</p>
<pre><code class="lang-auto">class AngleModel(torch.nn.Module):
    def __init__(self):
        super(AngleModel, self).__init__()
        self.angle = nn.Parameter(torch.tensor(0.0)) # parameter will have requires_grad by default, so no need to pass in a tensor requiring grad.

    def forward(self, input):
        # make qx...qz locals
        qx = torch.zeros(1)
        qy = torch.zeros(1)
        # qw and qz are dependent on the angle, so we need to recompute them here
        qw = torch.cos(self.angle / 2.)
        qz = torch.sin(self.angle / 2.)
        # there isn't anything wrong with using torch.pow, but personally, I like **
        m11 = 1. - 2. * qy**2 - 2. * qz**2
        m22 = 1. - 2. * qx**2 - 2. * qz**2
        m33 = 1. - 2. * qx**2 - 2. * qy**2

        m21 = 2. * self.qx * self.qy - 2. * self.qz * self.qw
        m12 = 2. * self.qx * self.qy + 2. * self.qz * self.qw

        m31 = 2. * self.qx * self.qz + 2 * self.qy * self.qw
        m13 = 2. * self.qx * self.qz - 2 * self.qy * self.qw

        m32 = 2. * self.qy * self.qz - 2. * self.qx * self.qw
        m23 = 2. * self.qy * self.qz + 2. * self.qx * self.qw

        # noone uses Tensor anymore, and tensor isn't right here either, so you have to cat your way to the matrix or do something differently elsewhere
        matrix = torch.stack([
            torch.cat([m11, m21, m31], dim=0),
            torch.cat([[m12, m22, m32], dim=0),
            torch.cat([[m13, m23, m33], dim=0),
        ], dim=0)

        output = torch.matmul(input, matrix)
        return output
</code></pre>
<p>or something similar (I didn’t run your code because your example isn’t completely self-contained nor did I check the maths) should work.</p>
<p>Best regards</p>
<p>Thomas</p>

Ingenieur: <p>Works like a charm! Thank you very much! That’s what I figured that it had gotten something to do with a broken graph…</p>
<p>Since I know the axis of the rotation and it is one of the three main axes of the coordinate system, the structure of the corresponding quaternion is already given. That’s why I don’t need to learn qx and qy, hence they can both be 0.</p>
<p>For the sake of completeness, the working final code of the nn.Module looks like this now:</p>
<pre><code class="lang-auto">class AngleModel(nn.Module):
    def __init__(self):
        super(AngleModel, self).__init__()
        self.angle = nn.Parameter(torch.tensor(0.0))

    def forward(self, input):
        qw = torch.cos(self.angle / 2.)
        qx = 0.0
        qy = 0.0
        qz = torch.sin(self.angle / 2.)

        matrix = torch.zeros(3, 3)

        matrix[0, 0] = 1. - 2. * qy ** 2 - 2. * qz ** 2
        matrix[1, 1] = 1. - 2. * qx ** 2 - 2. * qz ** 2
        matrix[2, 2] = 1. - 2. * qx ** 2 - 2. * qy ** 2

        matrix[0, 1] = 2. * qx * qy - 2. * qz * qw
        matrix[1, 0] = 2. * qx * qy + 2. * qz * qw

        matrix[0, 2] = 2. * qx * qz + 2 * qy * qw
        matrix[2, 0] = 2. * qx * qz - 2 * qy * qw

        matrix[1, 2] = 2. * qy * qz - 2. * qx * qw
        matrix[2, 1] = 2. * qy * qz + 2. * qx * qw

        output = torch.matmul(input, matrix)
        return output
</code></pre>

aksg87: <p><a class="mention" href="/u/ingenieur">@Ingenieur</a></p>
<p>Can you provide sample inputs where there is a learnable function that maps <code>input</code> to <code>target</code>?</p>
<p>I created an example notebook for running this code:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1CwiBqnSXk7VjTdWEavDExXkgZrKVvErl?usp=sharing">
  <header class="source">
      <img src="https://colab.research.google.com/img/favicon.ico?vrz=colab-20210920-060054-RC00_397697772" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1CwiBqnSXk7VjTdWEavDExXkgZrKVvErl?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar" width="260" height="260">

<h3><a href="https://colab.research.google.com/drive/1CwiBqnSXk7VjTdWEavDExXkgZrKVvErl?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>