Hi everybody,
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.
What I have so far is an nn.Module that executes the rotation with the current angle on the input. Code looks as follows:
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

I initialize the model and optimizer like this:
        model = AngleModel()
        crit = torch.nn.MSELoss()
        l_rate = 0.01
        optim = torch.optim.SGD(model.parameters(), lr = l_rate)
        epochs = 100

and execute the training like this:
        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()))

I suspect the trigonometric and pow functions to be the culprit of why angle isn’t getting updated. Am I right? angle.grad is None after loss.backward().
How would I go on about this? What do I need to do in order to get this working?
Thank you very much!