在Numpy/PyTorch中矢量化生成一批随机旋转矩阵的最佳方法？

Question

我想基于初始起始矩阵（其形状为，例如，(4096, 3)）生成随机旋转矩阵的批次，其中应用于批次中每个矩阵的旋转是从一组旋转矩阵中随机选择的（在我的原始帖子代码中，我只想从8个可能的旋转角度中随机选择）。因此，我最终得到的是一个(batch_size, 4096, 3)形状的张量

我目前的方法是预先生成可能的旋转矩阵（因为我只处理8个可能的随机旋转），然后使用for循环为批中的每个项目随机选择八个预先制作的旋转矩阵中的一个来生成批。这不是非常高效，所以我希望以某种方式将整个过程矢量化

现在，我就是这样循环一批，一个接一个地生成一批旋转矩阵的：

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

在下面的代码中，我生成了一组预先制作的随机旋转矩阵，这些矩阵是在批处理的for循环中随机选择的

make_3d_grid函数生成(grid_dim * grid_dim * grid_dim, 3)形矩阵（基本上是x、y、z坐标点的二维数组）。函数get_rotation_matrix返回(3, 3)旋转矩阵，其中θ用于绕x轴旋转

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)

如果这方面的代码可以在Numpy中完成，那么也可以很好地工作，因为我可以自己将其转换为PyTorch，我将不胜感激

Answer 1

您可以将旋转矩阵预生成为(batch_size, 3, 3)数组，然后乘以广播到(batch_size, N, 3)的(N, 3)点数组

rotated_points = np.dot(pointsf, rots)

^{}将pointsf的最后一个轴和rots的第二个到最后一个轴的乘积相加，将pointsf的维数放在第一位。这意味着您的结果将是(N, batch_size, 3)形状，而不是(batch_size, N, 3)。当然，您可以通过简单的轴交换来解决此问题：

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

或

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

然而，我建议你将rots作为你之前的旋转矩阵的逆（转置）。在这种情况下，您只需计算：

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

您应该能够非常轻松地将np.dot转换为torch.mm