如何精确实现PyTorch Linear层前向计算？消除自定义实现误差

问题根源 #

你的自定义实现误差来源于计算逻辑不一致和循环带来的浮点累加顺序差异：

• F.linear的本质是矩阵乘法：输入形状为(N, Cin)、权重为(Cout, Cin)时，输出等价于input @ weight.t()（输入矩阵与权重转置做矩阵乘法）。
• 你用循环逐样本做元素乘再求和，这种计算顺序和PyTorch底层调用的优化BLAS矩阵乘法（如CUDA上的cublas）的累加逻辑不同，会产生浮点精度差异，大张量场景下误差还会被放大。

无误差实现代码 #

直接用矩阵乘法替代循环，完全对齐F.linear的计算逻辑：

import torch
import torch.nn.functional as F

input = torch.randn(64,7,7, 120).to('cuda:0')
weight= torch.randn(360,120).to('cuda:0')
pytorch_output =  F.linear(input, weight)

# 自定义无误差实现
B,H,W,_ = input.shape
# 将输入flatten为(N, Cin)，其中N=B*H*W
input_flat = input.flatten(end_dim=-2)
# 直接执行矩阵乘法：(N, Cin) @ (Cin, Cout) = (N, Cout)
my_output = input_flat @ weight.t()
# 恢复原张量形状
my_output = my_output.reshape(B,H,W,-1)

# 验证误差
diff = my_output - pytorch_output
print(diff.max())  # 误差会降至浮点精度极限，如~1e-15

关键说明 #

• 该实现复用PyTorch底层优化的矩阵乘法逻辑，与F.linear的计算路径完全一致，可消除额外精度误差。
• 禁止用循环逐样本计算：循环不仅运行效率极低，还会破坏矩阵乘法的优化累加顺序，引发精度差异。
• 若需支持偏置项，只需在矩阵乘法结果后直接叠加偏置（偏置形状需为(Cout,)，与F.linear要求一致）。

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