Лучший способ векторизации генерации партии случайно повернутых матриц в Numpy / PyTorch?

Question

Я хотел бы генерировать партии случайно повернутых матриц на основе начальной стартовой матрицы (которая имеет форму, например, (4096, 3)), где вращение, применяемое к каждой матрице в пакете, случайным образом выбирается из группа матриц вращения (в моем коде в оригинальном посте я хочу только случайным образом выбрать один из 8 возможных углов поворота). Поэтому я получаю тензор формы (batch_size, 4096, 3).

. Мой текущий подход заключается в том, что я заранее делаю возможные повернутые матрицы (поскольку я имею дело только с 8 возможными случайными вращениями), и затем используйте для l oop, чтобы сгенерировать партию путем случайного выбора одной из восьми предварительно сделанных повернутых матриц для каждого элемента в партии. Это не очень эффективно, поэтому я надеялся как-то векторизовать весь процесс.

Прямо сейчас, вот как я oop перебираю партию, чтобы генерировать партию повернутых матриц одна за другой:

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

В приведенном ниже коде я генерирую заранее подготовленный набор матриц случайного вращения, которые выбираются случайным образом в течение -1 oop для пакета.

Функция 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), а затем умножить его на массив (N, 3) точек, транслируемый в (batch_size, N, 3).

rotated_points = np.dot(pointsf, rots)

np.dot будет суммировать произведение по последней оси 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 довольно тривиально.