Best way to vectorize generating a batch of randomly rotated matrices in Numpy/PyTorch?

Question

I’d like to generate batches of randomly rotated matrices based on an initial starting matrix (which has a shape 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.

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 * grid_dim * grid_dim, 3) shaped matrix (basically a 2D array of x, y, z coordinate points). The get_rotation_matrix function returns a (3, 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 help in vectorizing this would be greatly appreciated! If code for this can be done in Numpy that works fine too, since I can convert it to PyTorch myself.

Mad Physicist Mad Physicist · Accepted Answer · 2020-02-03T15:50:44

You can pre-generate your rotation matrices as a (batch_size, 3, 3) array, and then multiply by your (N, 3) points array broadcasted to (batch_size, N, 3).

rotated_points = np.dot(pointsf, rots)

np.dot will sum-product over the last axis of pointsf and the second-to-last axis of rots, putting the dimensions of pointsf first. This means that your result will be of shape (N, batch_size, 3) rather than (batch_size, N, 3). You can of course fix this with a simple axis swap:

rotated_points = np.dot(pointsf, rots).transpose(1, 0, 2)

OR

rotated_points = np.swapaxes(np.dot(pointsf, rots), 0, 1)

I would suggest, however, that you make rots be the inverse (transposed) rotation matrices from what you had before. In that case, you can just compute:

rotated_points = np.dot(transposed_rots, pointsf.T)

You should be able to convert np.dot to torch.mm fairly trivially.

Best way to vectorize generating a batch of randomly rotated matrices in Numpy/PyTorch?

1 Answers