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.