PyTorch实现Cifar10分类耗时远高于TensorFlow问题求助

导致速度差距的核心原因是PyTorch默认的数据加载策略效率远低于TensorFlow的tf.data默认配置，按下面的步骤修改即可拉平甚至超过TensorFlow的运行速度：

必改核心配置 #

• 给DataLoader开启多进程加载和锁页内存
  PyTorch默认num_workers=0，仅使用单进程加载数据，训练过程中大部分时间都在等CPU加载数据，GPU利用率极低。同时开启pin_memory=True可以减少数据从内存转移到GPU的耗时，修改代码如下：

dataloaders = {phase: torch.utils.data.DataLoader(
    data[phase], 
    batch_size=32, 
    shuffle=True,
    num_workers=4, # Colab环境下设置2~4即可，数值过大会增加CPU调度开销
    pin_memory=True
) for phase in ['train', 'test']}

• 修正准确率和损失计算逻辑（不影响速度，但会修复你结果里错误的准确率数值）
  你当前用len(dataloaders[phase])（总batch数）做分母，计算出来的数值不是准确率，应该改为除以总样本数：

epoch_loss = running_loss / len(dataloaders[phase].dataset)
epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

可选提速配置 #

• 开启CuDNN自动调优
  在代码开头加入以下配置，CuDNN会提前适配最优的卷积计算方案，适合输入尺寸固定的场景：

import torch.backends.cudnn as cudnn
cudnn.benchmark = True

• 降低进度条刷新开销
  tqdm默认每步刷新进度条会占用少量CPU资源，设置最小刷新间隔可减少开销：

for inputs, labels in tqdm(iter(dataloaders[phase]), mininterval=1):

• 启用混合精度训练
  使用PyTorch自带的自动混合精度工具，可将训练速度提升30%~50%，同时降低显存占用。

修改完成后训练耗时会降到和TensorFlow相近的水平。

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