Best way to vectorize generating a batch of randomly rotated matrices

username: Yuerno
cooked: <p>Hello all! In my code right now, I need to generate a batch of randomly rotated matrices, which I do in a for loop. This is obviously inefficient, so I was wondering how/if it’s possible to vectorize this, such that I can generate a random batch of rotated matrices each time?</p>
<p>Right now, this is how I loop over a batch to generate a batch of rotated matrices one by one.</p>
<pre><code class="lang-auto">for view_i in range(batch_size):
            # Get rotated view grid points randomly
            idx = torch.randint(0, 8, (1,))
            pointsf = rotated_points[idx]
</code></pre>
<p>In the code below, I generate a pre-made set of random rotation matrices that get randomly selected from in a for-loop over the batch.</p>
<p>The <strong>make_3d_grid function</strong> generates a [grid_dim x grid_dim x grid_dim, 3] shape matrix (basically a 2D array of x, y, z coordinate points). The <strong>get_rotation_matrix function</strong> returns a 3 x 3 rotation matrix, where theta is used for rotation around the x-axis.</p>
<pre><code class="lang-auto">rotated_points = []
grid_dim = 16

pointsf = make_3d_grid((-1,)*3, (1,)*3, (grid_dim,)*3)

view_angles = torch.tensor([0, np.pi/4.0, np.pi/2.0, 3*np.pi/4.0, np.pi, 5*np.pi/4.0, 3*np.pi/2.0, 7*np.pi/4.0])

for i in range(len(view_angles)):
    theta = view_angles[i]

    rot = get_rotation_matrix(theta, torch.tensor(0.0), torch.tensor(0.0))

    pointsf_rot = torch.mm(pointsf, rot)

    rotated_points.append(pointsf_rot)
</code></pre>
<p>Any ideas on this would be appreciated! The main issue I’m struggling with right now is how to incorporate the randomness of the rotated grids into each batch.</p>

username: ptrblck
cooked: <p>I’m not sure to understand the posted code correctly, but based on your description I assume you would like to apply different rotation matrices on input tensors?<br>
I’ve written a manual rotation some time ago <a href="https://github.com/ptrblck/pytorch_misc/blob/master/image_rotation_with_matrix.py">here</a>. Would this be helpful as a starter or are you looking for something else?</p>

username: Yuerno
cooked: <p>Kind of! My original description was too complicated, hopefully this is a better explanation:</p>
<p>I’d like to generate batches of randomly rotated matrices based on an initial starting matrix (which has dimensions of, for example, [4096, 3]), where the rotation applied to each matrix in the batch is randomly chosen from a group of rotation matrices (in my code in the original post, I only want to randomly select from 8 possible rotation angles). Therefore, what I end up with is a tensor of shape [batch_size, 4096, 3].</p>
<p>My current approach is that I pre-make the possible rotated matrices (since I’m only dealing with 8 possible random rotations), and then use a for loop to generate the batch by randomly picking one of the eight pre-made rotated matrices for each item in the batch. This isn’t super efficient, so I was hoping to vectorize the whole process somehow.</p>
<p>Please let me know if that clarifies my question at all, or if I’m still being too confusing.</p>