如何使用PyTorch张量算子替换代码中的for循环迭代

问题描述 #

需要将以下包含三层Python for循环的PyTorch代码替换为无循环的原生张量运算，消除Python层迭代开销：

B=10
L=20
H=5

mat_A=torch.randn(B,L,L,H)
mat_B=torch.randn(L,B,B,H)
tmp_B=torch.zeros_like(mat_B)

for x in range(L):
    for y in range(B):
        for z in range(B):
            tmp_B[:,y,z,:]+=mat_B[x,y,z,:]*mat_A[z,x,:,:]

实现思路 #

完全可以通过维度对齐+广播+规约求和实现无循环替换，计算结果和原逻辑完全一致：

1. 先梳理原循环的维度对应关系：
   • 输出tmp_B形状为(L, B, B, H)，第一个L维度对应原代码中mat_A[z,x,:,:]的第二个L维度
   • 循环变量x是累加维度，遍历范围0~L-1
   • 循环变量y、z是输出的自由维度，不需要累加
   • 最后一个H维度是逐元素计算的公共维度，不需要归约
2. 对两个输入张量做维度扩展/重排，让所有维度位置对齐后直接广播相乘，再沿累加维度x求和即可得到最终结果。

无循环实现代码 #

import torch

B=10
L=20
H=5

mat_A = torch.randn(B, L, L, H)
mat_B = torch.randn(L, B, B, H)

# 维度对齐
# mat_B原维度顺序(x, y, z, h)，扩维后形状(1, L, B, B, H)，新增第0维对应输出的第一个L维度
mat_B_align = mat_B.unsqueeze(0)
# mat_A原维度顺序(z, x, i, h)，重排+扩维后形状(L, L, 1, B, H)，维度顺序对应(i, x, 占位y, z, h)
mat_A_align = mat_A.permute(2, 1, 0, 3).unsqueeze(2)
# 广播相乘后沿x维度（第1维）求和，得到输出形状(L, B, B, H)
tmp_B = (mat_B_align * mat_A_align).sum(dim=1)

正确性验证 #

可以通过和原循环结果逐元素对比验证一致性，运行以下代码会输出True：

# 原循环实现
tmp_B_origin = torch.zeros_like(mat_B)
for x in range(L):
    for y in range(B):
        for z in range(B):
            tmp_B_origin[:,y,z,:] += mat_B[x,y,z,:] * mat_A[z,x,:,:]

print(torch.allclose(tmp_B, tmp_B_origin))

这种实现完全运行在PyTorch底层C++/CUDA算子上，没有Python层循环开销，在GPU上运行时速度提升会非常显著。

内容的提问来源于stack exchange，提问作者Newbie1111