Constructing a matrix variable from other variables

1

I’m trying to construct a rotation matrix from the Euler angles. As simple as I feel this should be, I’m not sure what the best approach is in PyTorch. Is there a way to do this without spamming torch.cat? Essentially I just to want to stack some existing tensor variables into a matrix for convenience.

Something like

angles = Variable(torch.zeros(3))
a, b, c = angles[0], angles[1], angles[2]
rot_mat1 = Variable([
   (1, 0, 0),
   (0, a.cos(), a.sin()),
   (0, -a.sin(), a.cos()),
])
rot_mat2 = ...
rot_mat3 = ...
rot_mat = rot_mat1 @ rot_mat2 @ rot_mat3
2 Likes
2

You can index the variables as you want

rot_mat1 = Variable(torch.zeros(3, 3), requires_grad=False)
rot_mat1[1, 1] = a.cos()
...

I’d initialize the matrix with all the fixed values (using a tensor), and I’d then assign the variable elements after the Variable was created.

2 Likes
3

Okay that seems like a better solution that what I came up with. Thanks for the solution. Glad to know assigning indices like that works. It seemed like it wouldn’t since I’ve had errors from setting values with masks etc.

Here’s the messy torch.cat solution in case anyone is interested.

def rotation_tensor(theta, phi, psi, n_comps):
    one = Variable(torch.ones(n_comps, 1, 1))
    zero = Variable(torch.zeros(n_comps, 1, 1))
    rot_x = torch.cat((
        torch.cat((one, zero, zero), 1),
        torch.cat((zero, theta.cos(), theta.sin()), 1),
        torch.cat((zero, -theta.sin(), theta.cos()), 1),
    ), 2)
    rot_y = torch.cat((
        torch.cat((phi.cos(), zero, -phi.sin()), 1),
        torch.cat((zero, one, zero), 1),
        torch.cat((phi.sin(), zero, phi.cos()), 1),
    ), 2)
    rot_z = torch.cat((
        torch.cat((psi.cos(), -psi.sin(), zero), 1),
        torch.cat((psi.sin(), psi.cos(), zero), 1),
        torch.cat((zero, zero, one), 1)
    ), 2)
    return torch.bmm(rot_z, torch.bmm(rot_y, rot_x))

n_comps = 5
theta = Variable(torch.zeros(n_comps, 1, 1), requires_grad=True)
phi = Variable(torch.zeros(n_comps, 1, 1), requires_grad=True)
psi = Variable(torch.zeros(n_comps, 1, 1), requires_grad=True)

rotation_mat = rotation_tensor(theta, phi, psi, n_comps)
1 Like
4

@kpar, are you implementing this?

I’m might be doing it soon to?

5

Nothing so fancy, I’m just trying to optimize for a color transformation and I’m parameterizing part it as a vector rotation.

1 Like
6

Sounds fun :slight_smile: :slight_smile: Thanks for posting your rotation tensor, it’s cool !

The unitary RNN is a bit complicated :frowning: - but still doable

7

I’m not sure if it’s a different issue, but using a matrix constructed like this to optimize the base variables is causing a second backwards pass to fail. I’ve removed the unrelated code. The optimization works without the rotation matrix but gives

RuntimeError: Trying to backward through the graph second time, but the buffers have already been freed. Please specify retain_variables=True when calling backward for the first time.

when run with the rotation matrix. Any insight as to why this might be? Here’s the cruft of the optimization that’s causing the issue.

def rotation_tensor(theta, phi, psi, n_comps):
    rot_x = Variable(torch.zeros(n_comps, 3, 3).cuda(), requires_grad=False)
    rot_y = Variable(torch.zeros(n_comps, 3, 3).cuda(), requires_grad=False)
    rot_z = Variable(torch.zeros(n_comps, 3, 3).cuda(), requires_grad=False)
    rot_x[:, 0, 0] = 1
    rot_x[:, 0, 1] = 0
    rot_x[:, 0, 2] = 0
    rot_x[:, 1, 0] = 0
    rot_x[:, 1, 1] = theta.cos()
    rot_x[:, 1, 2] = theta.sin()
    rot_x[:, 2, 0] = 0
    rot_x[:, 2, 1] = -theta.sin()
    rot_x[:, 2, 2] = theta.cos()
    
    rot_y[:, 0, 0] = phi.cos()
    rot_y[:, 0, 1] = 0
    rot_y[:, 0, 2] = -phi.sin()
    rot_y[:, 1, 0] = 0
    rot_y[:, 1, 1] = 1
    rot_y[:, 1, 2] = 0
    rot_y[:, 2, 0] = phi.sin()
    rot_y[:, 2, 1] = 0
    rot_y[:, 2, 2] = phi.cos()
    
    rot_z[:, 0, 0] = psi.cos()
    rot_z[:, 0, 1] = -psi.sin()
    rot_z[:, 0, 2] = 0
    rot_z[:, 1, 0] = psi.sin()
    rot_z[:, 1, 1] = psi.cos()
    rot_z[:, 1, 2] = 0
    rot_z[:, 2, 0] = 0
    rot_z[:, 2, 1] = 0
    rot_z[:, 2, 2] = 1
    return torch.bmm(rot_z, torch.bmm(rot_y, rot_x))

optimizer = optim.LBFGS([diff_theta, diff_phi, diff_psi], lr=lr)
diff_theta = Variable(torch.zeros(n_diff_comps, 1, 1).cuda(), requires_grad=True)
diff_phi = Variable(torch.zeros(n_diff_comps, 1, 1).cuda(), requires_grad=True)
diff_psi = Variable(torch.zeros(n_diff_comps, 1, 1).cuda(), requires_grad=True)

diff_rot = rotation_tensor(diff_theta, diff_phi, diff_psi, n_diff_comps).cuda()

for i in range(3000):
    def closure():
        optimizer.zero_grad()
        new_diff = Variable(torch.zeros(*alb_diff_batch_init.size()).cuda())
        for comp in range(n_diff_comps):
            new_diff += (diff_rot[comp] @ diff_base_vars[comp].view(3, -1)).view(*diff_base_vars[comp].size())
        loss = new_diff.sum()
        loss.backward()
        return loss
    optimizer.step(closure)

Edit: It looks like you need to re-create the intermediate matrix again every iteration. Moving

diff_rot = rotation_tensor(diff_theta, diff_phi, diff_psi, n_diff_comps).cuda()

to the inner loop fixed the issue.

1 Like
8

Can we do it in a batch version?