使用torch.nn.DataParallel多GPU预测结果与单GPU/CPU不一致问题

核心原因 #

差异的主要来源是模型处于训练模式时，BatchNorm层的行为差异：

• 单GPU/CPU运行时，BatchNorm使用整个输入batch的均值和方差进行归一化；
• 多GPU（DataParallel）运行时，每个GPU仅处理batch的一部分（比如64个样本拆分为两个GPU各处理32个），训练模式下每个GPU会基于自己的子batch计算均值方差，导致归一化结果不同，最终输出出现显著差异；
• 此外，训练模式下的Dropout层会引入随机性，进一步放大结果差异。

解决方法 #

将所有模型切换到评估模式，此时BatchNorm会使用预初始化的running_mean和running_var（而非当前batch统计量），Dropout层也会关闭，从而保证多GPU与单GPU/CPU的输出一致（仅存在微小浮点数精度差异）。

修改后的代码 #

import torch
import torchvision

def load_data(num_gpus):
    transforms = torchvision.transforms.Compose([
        torchvision.transforms.Resize(256),
        torchvision.transforms.CenterCrop(224),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    dataset = torchvision.datasets.CIFAR10('datasets', train=True, transform=transforms, download=True)

    dataloader = torch.utils.data.DataLoader(
        dataset=dataset,
        batch_size=64,
        shuffle=False,
        num_workers=4*num_gpus
    )
    return dataloader

# 多GPU模型初始化
model = torchvision.models.resnet50(pretrained=False)
model = torch.nn.parallel.DataParallel(model, device_ids=[0,1])
model = model.cuda()
model.eval()  # 切换到评估模式

dataloader = load_data(2)

m_s = model.module.state_dict()

# 单CPU模型
model1 = torchvision.models.resnet50(pretrained=False)
model1.load_state_dict(m_s)
model1.eval()  # 切换到评估模式

# 单GPU模型
model2 = torchvision.models.resnet50(pretrained=False)
model2.load_state_dict(m_s)
model2 = model2.to('cuda:2')
model2.eval()  # 切换到评估模式

for images, labels in dataloader:
    break

# 关闭梯度计算，提升效率并避免不必要的计算
with torch.no_grad():
    multi_gpu = model(images.cuda())
    cpu_predict = model1(images)
    s_gpu = model2(images.to('cuda:2'))

a = multi_gpu.reshape(-1).cpu()
b = cpu_predict.reshape(-1)
c = s_gpu.reshape(-1).cpu()

for i,j,k in zip(a,b,c):
    print(i.item(), j.item(), k.item())

额外注意事项 #

• 加上torch.no_grad()上下文管理器，关闭梯度计算，既可以提升运行速度，也能避免因梯度计算带来的微小数值差异；
• DataParallel的dim参数默认就是0（对应batch维度），无需手动指定，你的代码中该设置是正确的；
• 若仍存在微小差异，属于GPU与CPU之间的浮点数计算精度差异（比如CUDA的float32计算与CPU的float32计算存在细微偏差），这是正常现象，可通过使用torch.float64类型缓解，但会增加计算开销。

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