Dive-into-DL-PyTorch项目解析：PyTorch多GPU计算实战指南

赖旦轩

于 2025-06-02 09:04:24 发布

阅读量234

点赞数 5

本文链接：https://blog.csdn.net/gitblog_00804/article/details/148375531

版权

Dive-into-DL-PyTorch项目解析：PyTorch多GPU计算实战指南

Dive-into-DL-PyTorch 本项目将《动手学深度学习》(Dive into Deep Learning)原书中的MXNet实现改为PyTorch实现。项目地址: https://gitcode.com/gh_mirrors/di/Dive-into-DL-PyTorch

多GPU计算概述

在现代深度学习实践中，使用多GPU进行模型训练已经成为加速训练过程的标配方案。PyTorch作为当前主流的深度学习框架，提供了简洁高效的多GPU支持。本文将详细介绍如何在PyTorch环境下利用多块GPU进行模型训练和推理。

检查GPU设备

在开始多GPU计算前，我们需要确认当前系统中的GPU资源情况。在Linux系统中，可以通过nvidia-smi命令查看：

nvidia-smi

该命令会输出类似如下的信息：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 390.48                 Driver Version: 390.48                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  TITAN X (Pascal)    Off  | 00000000:02:00.0 Off |                  N/A |
| 46%   76C    P2    87W / 250W |  10995MiB / 12196MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  TITAN X (Pascal)    Off  | 00000000:04:00.0 Off |                  N/A |
| 53%   84C    P2   143W / 250W |  11671MiB / 12196MiB |      4%      Default |
+-------------------------------+----------------------+----------------------+

输出信息包含了每块GPU的型号、温度、功耗、显存使用情况和利用率等重要信息，帮助我们了解当前GPU的工作状态。

使用DataParallel实现多GPU计算

PyTorch提供了torch.nn.DataParallel这个简洁的接口来实现多GPU计算。其基本原理是将输入数据在batch维度上进行分割，然后将分割后的数据分发到不同的GPU上进行并行计算。

基本使用方法

import torch

# 定义一个简单的线性模型并移动到GPU上
net = torch.nn.Linear(10, 1).cuda()

# 使用DataParallel包装模型
net = torch.nn.DataParallel(net)

经过DataParallel包装后，模型的前向传播会自动将计算分布到所有可用的GPU上。

指定特定GPU

如果系统中有多块GPU但只想使用其中部分设备，可以通过device_ids参数指定：

net = torch.nn.DataParallel(net, device_ids=[0, 3])

这表示只使用第0和第3号GPU进行计算。

多GPU模型的保存与加载

使用多GPU训练模型后，模型的保存和加载需要特别注意，因为DataParallel会在模型外层添加一个包装层。

常见错误方式

# 保存模型
torch.save(net.state_dict(), "model.pt")

# 加载模型
new_net = torch.nn.Linear(10, 1)
new_net.load_state_dict(torch.load("model.pt"))  # 这里会报错

这种直接保存和加载的方式会报错，因为保存的state_dict包含了DataParallel添加的"module."前缀。

正确保存方式

方法一：保存内部模型

# 保存时访问module属性
torch.save(net.module.state_dict(), "model.pt")

# 加载时直接加载
new_net = torch.nn.Linear(10, 1)
new_net.load_state_dict(torch.load("model.pt"))  # 成功加载

方法二：保持DataParallel结构

# 保存完整模型
torch.save(net.state_dict(), "model.pt")

# 加载时先创建DataParallel结构
new_net = torch.nn.Linear(10, 1)
new_net = torch.nn.DataParallel(new_net)
new_net.load_state_dict(torch.load("model.pt"))  # 成功加载

推荐使用方法一，因为它更灵活，可以在不需要多GPU的环境下直接加载使用模型。

多GPU计算的注意事项

数据并行与模型并行：DataParallel实现的是数据并行，即同一模型在不同数据上并行计算。PyTorch还支持模型并行，即将模型的不同部分放在不同GPU上。
batch size设置：使用多GPU时，总batch size是各GPU上batch size的总和。需要合理设置以确保每个GPU上的batch size仍然有效。
GPU间通信开销：DataParallel需要在每次前向传播后同步梯度，这会带来一定的通信开销。对于小模型，可能反而会降低训练速度。
显存使用：多GPU可以分摊显存使用，使得可以训练更大的模型或使用更大的batch size。
随机种子：如果需要可重复的结果，需要确保各GPU上的随机种子一致。