Best way to vectorize generating a batch of randomly rotated matrices

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?

Right now, this is how I loop over a batch to generate a batch of rotated matrices one by one.

for view_i in range(batch_size):
            # Get rotated view grid points randomly
            idx = torch.randint(0, 8, (1,))
            pointsf = rotated_points[idx]

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.

The make_3d_grid function generates a [grid_dim x grid_dim x grid_dim, 3] shape matrix (basically a 2D array of x, y, z coordinate points). The get_rotation_matrix function returns a 3 x 3 rotation matrix, where theta is used for rotation around the x-axis.

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)

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.

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?
I’ve written a manual rotation some time ago here. Would this be helpful as a starter or are you looking for something else?

1 Like

Kind of! My original description was too complicated, hopefully this is a better explanation:

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].

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.

Please let me know if that clarifies my question at all, or if I’m still being too confusing.