如何基于1D张量取值实现2D张量的高效取整？优化循环方案

你可以直接利用PyTorch的高级索引特性，替换掉低效的双重循环，实现纯张量级的高效操作，代码修改非常简单：

核心优化点 #

• 确保rounded张量和输入predictions处于同一计算设备（CPU/GPU），避免设备不匹配报错
• 用rounded[indexes]直接生成目标张量，这行代码会自动输出与indexes形状完全一致的结果，无需逐元素循环

改进后的完整代码 #

def round_values(predictions):
    # Rounding values，同步到输入张量的设备
    rounded = torch.tensor([1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0]).to(predictions.device)

    # 计算与各个取整值的差值，找到最近值的索引
    dif = predictions.unsqueeze(2) - rounded.view(1, 1, -1)
    indexes = torch.argmin(torch.abs(dif), dim=2)

    # 直接通过索引生成结果，替换原有的双重循环
    rounded_predictions = rounded[indexes]
    
    return rounded_predictions 

a = 4*torch.rand((6,6)) + 1
print(a)
print(round_values(a))

原理说明 #

PyTorch允许使用任意形状的整数张量作为索引，去索引一维张量。当你用2维的indexes索引1维的rounded时，底层会自动完成广播匹配，每个位置的索引值会对应取出rounded中对应位置的数值，最终生成的张量形状和indexes完全一致。这种操作是基于优化后的张量运算实现的，效率远高于Python层面的双重循环，尤其是在处理大规模张量时，性能提升会非常显著。

额外优化建议 #

如果rounded的取值是固定不变的，可以把它定义成全局常量，避免每次调用函数时重复创建张量，进一步降低开销。

内容的提问来源于stack exchange，提问作者Alberto Martín