Constructing a matrix variable from other variables

kpar: <p>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.</p>
<p>Something like</p>
<pre><code class="lang-python">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
</code></pre>

fmassa: <p>You can index the variables as you want</p>
<pre><code class="lang-python">rot_mat1 = Variable(torch.zeros(3, 3), requires_grad=False)
rot_mat1[1, 1] = a.cos()
...
</code></pre>
<p>I’d initialize the matrix with all the fixed values (using a tensor), and I’d then assign the variable elements after the <code>Variable</code> was created.</p>

kpar: <p>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.</p>
<p>Here’s the messy torch.cat solution in case anyone is interested.</p>
<pre><code class="lang-python">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)
</code></pre>

AjayTalati: <p><a class="mention" href="/u/kpar">@kpar</a>, are you implementing this?</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/jingli9111/EUNN-tensorflow/blob/master/EUNN.py#L99" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/jingli9111/EUNN-tensorflow/blob/master/EUNN.py#L99" target="_blank" rel="nofollow noopener">jingli9111/EUNN-tensorflow/blob/master/EUNN.py#L99</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="89" style="counter-reset: li-counter 88 ;">
<li>capacity_a = capacity - capacity_b</li>
<li>
</li>
<li>hidden_size_a = hidden_size//2</li>
<li>hidden_size_b = (hidden_size-1)//2</li>
<li>
</li>
<li>params_theta_0 = vs.get_variable("theta_0", [capacity_a, hidden_size_a], initializer=theta_phi_initializer)</li>
<li>cos_theta_0 = array_ops.reshape(math_ops.cos(params_theta_0), [capacity_a, -1, 1])</li>
<li>sin_theta_0 = array_ops.reshape(math_ops.sin(params_theta_0), [capacity_a, -1, 1])</li>
<li>
</li>
<li>params_theta_1 = vs.get_variable("theta_1", [capacity_b, hidden_size_b], initializer=theta_phi_initializer)</li>
<li class="selected">cos_theta_1 = array_ops.reshape(math_ops.cos(params_theta_1), [capacity_b, -1, 1])</li>
<li>sin_theta_1 = array_ops.reshape(math_ops.sin(params_theta_1), [capacity_b, -1, 1])</li>
<li>
</li>
<li>if comp:</li>
<li>    params_phi_0 = vs.get_variable("phi_0", [capacity_a, hidden_size_a], initializer=theta_phi_initializer)</li>
<li>    cos_phi_0 = array_ops.reshape(math_ops.cos(params_phi_0), [capacity_a, -1, 1])</li>
<li>    sin_phi_0 = array_ops.reshape(math_ops.sin(params_phi_0), [capacity_a, -1, 1])</li>
<li>
</li>
<li>    cos_list_0_re = array_ops.reshape(array_ops.concat([cos_theta_0, math_ops.multiply(cos_theta_0, cos_phi_0)], 2), [capacity_a, -1])</li>
<li>    cos_list_0_im = array_ops.reshape(array_ops.concat([array_ops.zeros_like(cos_theta_0), math_ops.multiply(cos_theta_0, sin_phi_0)], 2), [capacity_a, -1])</li>
<li>    if hidden_size_a*2 != hidden_size:</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>I’m might be doing it soon to?</p>

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

AjayTalati: <p>Sounds fun <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:">  <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"> Thanks for posting your rotation tensor, it’s cool !</p>
<p>The unitary RNN is a bit complicated <img src="//discuss.pytorch.org/images/emoji/apple/frowning.png?v=5" title=":frowning:" class="emoji" alt=":frowning:"> - but still doable</p>

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

Zhipeng_Fan: <p>Can we do it in a batch version?</p>