Constructing a matrix variable from other variables

kpar: 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 <SCODE>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
<ECODE>

fmassa: You can index the variables as you want <SCODE>rot_mat1 = Variable(torch.zeros(3, 3), requires_grad=False)
rot_mat1[1, 1] = a.cos()
...
<ECODE>

kpar: 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. <SCODE>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)
<ECODE>

AjayTalati: I’m might be doing it soon to?

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

AjayTalati: 

kpar: 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 <SCODE>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.
<ECODE> 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. <SCODE>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)
<ECODE> Edit: It looks like you need to re-create the intermediate matrix again every iteration. Moving <SCODE>diff_rot = rotation_tensor(diff_theta, diff_phi, diff_psi, n_diff_comps).cuda()
<ECODE> to the inner loop fixed the issue.

Zhipeng_Fan: Can we do it in a batch version?