Colab中PyTorch使用GPU训练模型耗时高于CPU的问题排查

1. 逐样本串行计算完全浪费GPU并行优势 #

• 训练流程采用逐样本遍历的原生SGD实现：每个迭代轮次里，通过内层Python循环逐个取出单条784维MNIST样本，触发一次极小规模的矩阵运算送入GPU计算。
• GPU的设计目标是并行处理大规模矩阵任务，面对单次仅几KB数据量的小运算，固定的CUDA Kernel启动开销、CPU与GPU间的设备通信开销，远大于实际矩阵计算的耗时。这种场景下CPU没有跨设备通信的额外成本，小向量运算延迟更低，自然运行速度更快。

2. 高频小算子调用放大GPU固定开销 #

• 全流程的计算逻辑由Python层的for循环逐次触发PyTorch算子，哪怕是计算10维向量的差值这种简单操作，都要走「Python解释器调用→CUDA初始化→计算结果回传」的完整流程。
• 按MNIST 6万训练集、1万测试集规模计算，单轮迭代光是训练环节就要触发超12万次小型CUDA算子调用，加上误差统计环节的逐样本循环，累计数十万次的CUDA启动开销在CPU运行时完全不存在。

3. 频繁的跨设备同步打断GPU计算流水线 #

• 逐样本统计训练/测试误差时，每次触发Python层的判断语句（if predict != value），都需要PyTorch强制等待GPU完成当前计算、把结果同步回CPU才能执行判断逻辑。每算一个样本就触发一次同步，完全打乱GPU的批量计算流水线，进一步放大额外耗时。

• 将逐样本计算改为批量张量运算：跳过内层单样本循环，直接对整版训练数据做矩阵乘法（例如用 w1 @ train_data.T 一次完成全量样本的隐藏层计算），一次性完成所有样本的前向、反向传播流程，把数十万次的小算子调用压缩为数次大规模矩阵运算，充分释放GPU并行算力。
• 误差统计逻辑改为批量张量计算：一次性算出全量样本的预测标签、通过张量逐元素比较得到总错误数，把跨设备同步操作集中到统计完成的最后一步，避免频繁打断GPU计算。
• 梯度清零环节做就地优化：复用已创建的梯度张量，通过 zero_() 方法就地清零，替代每轮迭代重新创建全零梯度张量的逻辑，减少不必要的GPU内存申请开销。

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