pytorch 多GPU使用_optimizer和nn.dataparallel-优快云博客

本文链接：https://blog.youkuaiyun.com/iridescent_mian/article/details/88058978

本文详细介绍了PyTorch中使用多GPU进行并行计算的方法，包括设置环境变量、模型并行化、优化器并行化及预训练模型加载等关键步骤。解析了DataParallel接口的工作原理，强调了输入数据batch数量需不小于GPU数量的要求。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

单GPU：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

多GPU：

device_ids = [0,1,2,3]

model = model.cuda(device_ids[0])
model = nn.DataParallel(model, device_ids=device_ids)

optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.001)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

optimizer.module.step()

for param_lr in optimizer.module.param_groups: # 同样是要加module
# param_lr['lr'] = param_lr['lr'] * 0.999

加载多GPU预训练模型

model = ft_net()
pretained_model = torch.load('./model/all/8_model.pkl')
pretained_dict = pretained_model.module.state_dict()
model = ft_net()
model.load_state_dict(pretained_dict)

解释：

Pytorch 的多 GPU 处理接口是 torch.nn.DataParallel(module, device_ids)，其中 module 参数是所要执行的模型，而 device_ids 则是指定并行的 GPU id 列表。

而其并行处理机制是，首先将模型加载到主 GPU 上，然后再将模型复制到各个指定的从 GPU 中，然后将输入数据按 batch 维度进行划分，具体来说就是每个 GPU 分配到的数据 batch 数量是总输入数据的 batch 除以指定 GPU 个数。每个 GPU 将针对各自的输入数据独立进行 forward 计算，最后将各个 GPU 的 loss 进行求和，再用反向传播更新单个 GPU 上的模型参数，再将更新后的模型参数复制到剩余指定的 GPU 中，这样就完成了一次迭代计算。所以该接口还要求输入数据的 batch 数量要不小于所指定的 GPU 数量。