Как разбить трехмерную сетку значений координат на блоки в PyTorch?

Question

Если у меня есть партия однородных трехмерных сеток координатных местоположений с формой (например) [1, 32, 32, 32, 3], каков наилучший способ для меня разделить это на несколько четных кусков так что я мог бы закончить с чем-то вроде [1, 4096, 2, 2, 2, 3]? Другими словами, я делю один большой куб размером 32 x 32 x 32, где каждая точка является координатой x, y, z, на 4096 меньших 2 x 2 x 2 куба? Имеет ли смысл простая операция просмотра или она отбросит значения координат? Я смотрел на такие операции, как torch.chunk, но они требуют определенного размера c, чтобы разделить их, что я не уверен, применимо здесь.

Мой вариант использования для этого заключается в том, что у меня меньше [ 1, 16, 16, 16, 3] куба, поэтому я пытаюсь сопоставить точки из этой меньшей фигуры с соответствующими кубами в форме с повышенной дискретизацией [1, 32, 32, 32, 3] (начиная с одной координатной точки в форме 16 ^ 3 соответствует 8 точкам в форме 32 ^ 3).

Для дополнительного контекста, я сейчас создаю свою трехмерную сетку:

pxs = torch.linspace(-1, 1, 32)
pys = torch.linspace(-1, 1, 32)
pzs = torch.linspace(-1, 1, 32)

pxs = pxs.view(-1, 1, 1).expand(*shape).contiguous().view(size)
pys = pys.view(1, -1, 1).expand(*shape).contiguous().view(size)
pzs = pzs.view(1, 1, -1).expand(*shape).contiguous().view(size)
points = torch.stack([pxs, pys, pzs], dim=1)

grid_3d = torch.reshape(points, (32, 32, 32, 3))

Answer 1

Просто reshape это с размерами, которые вы хотите:

In [29]: lin = np.linspace(-1, 1, 32)
    ...: cube = np.stack(np.meshgrid(lin, lin, lin), axis=-1)
    ...: cube.shape
Out[29]: (32, 32, 32, 3)

In [30]: new_cube = cube.reshape((1, 4096, 2, 2, 2, 3))
    ...: new_cube.shape
Out[30]: (1, 4096, 2, 2, 2, 3)