pytorch 拓展cuda语言 多维索引问题

四维矩阵索引公式: 四维矩阵[N, C, H, W], 当前四维索引为 n, c, h, w,   out_idx = n*C*H*W + c*H*W + h*W + w

pytorch中的 .permute()

a = torch.randn(5, 3, 4)    # a的size为[5, 3, 4]
b = a.permute(0, 2, 1)  # 此时b的size会变成[5, 4, 3],但是对于b的索引轴的顺序没有变,对应的指针顺序还是[5, 3, 4]
b = b.contiguous()     # 此时索引轴的顺序变为permute()操作之后的，对应的指针顺序也为[5, 4, 3]

在pytorch中需要底层处理时，会写cuda代码，其中大多数操作为指针，使用指针计算索引时可以参考上面的公式，为了计算方便，会在底层处理前对要处理的tensor进行维度转换，为了索引正确，务必在.permute()后使用.contiguous()。