找不到有效的pytorch广播

Question

我的实现中有以下代码片段。有一个嵌套的for循环，包含3个循环。在主代码中，原始系统的3D坐标被堆叠为点的组成堆叠的1D向量，对于具有坐标(x，y，z(的点，样本细胞将看起来像

Predictions =[...x,y,z,...]

而对于我的计算，我需要将预测向量重新整形为具有CCD_ 1的2D矩阵其中i是矩阵CCD_ 2中的任何采样行。以下代码显示逻辑

prediction_reshaped=torch.zeros([batch,num_node,dimesion])
for i in range(batch):
for j in range(num_node):
for k in range(dimesion):
prediction_reshaped[i][j][k]=prediction[i][3*j+k]

他们有没有有效的广播来避免这三个嵌套循环？它正在减慢我的代码。torch.reform不符合我的目的。该代码是使用pytorch实现的，所有矩阵都是pytorch张量，但任何numpy解决方案也会有所帮助。

Answer 1

这应该可以完成任务。

import torch
batch = 2
num_nodes = 4
x = torch.rand(batch, num_nodes * 3)
# tensor([[0.8076, 0.2572, 0.7100, 0.4180, 0.6420, 0.4668, 0.8915, 0.0366,  0.5704,
#          0.0834, 0.3313, 0.9080],
#         [0.2925, 0.7367, 0.8013, 0.4516, 0.5470, 0.5123, 0.1929, 0.4191,  0.1174,
#          0.0076, 0.2864, 0.9151]])
x = x.reshape(batch, num_nodes, 3)
# tensor([[[0.8076, 0.2572, 0.7100],
#         [0.4180, 0.6420, 0.4668],
#         [0.8915, 0.0366, 0.5704],
#         [0.0834, 0.3313, 0.9080]],
#
#        [[0.2925, 0.7367, 0.8013],
#         [0.4516, 0.5470, 0.5123],
#         [0.1929, 0.4191, 0.1174],
#         [0.0076, 0.2864, 0.9151]]])