PyTorch/[PyTorch 学习笔记] 7.3 使用 GPU 训练模型

PyTorch/[PyTorch 学习笔记] 7.3 使用 GPU 训练模型

本章代码:

这篇文章主要介绍了 GPU 的使用。

在数据运算时,两个数据进行运算,那么它们必须同时存放在同一个设备,要么同时是 CPU,要么同时是 GPU。而且数据和模型都要在同一个设备上。数据和模型可以使用to()方法从一个设备转移到另一个设备。而数据的to()方法还可以转换数据类型。

  • 从 CPU 到 GPU

    1
    2
    3
    device = torch.device("cuda")
    tensor = tensor.to(device)
    module.to(device)
  • 从 GPU 到 CPU

    1
    2
    3
    device = torch.device(cpu)
    tensor = tensor.to("cpu")
    module.to("cpu")

    tensormoduleto()方法的区别是:tensor.to()执行的不是 inplace 操作,因此需要赋值;module.to()执行的是 inplace 操作。

下面的代码是转换数据类型

1
2
x = torch.ones((3,3))
x = x.to(torch.float64)

tensor.to()module.to()

首先导入库,获取 GPU 的 device

1
2
3
import torch
import torch.nn as nn
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

下面的代码是执行Tensorto()方法

1
2
3
4
5
x_cpu = torch.ones((3, 3))
print("x_cpu:\ndevice: {} is_cuda: {} id: {}".format(x_cpu.device, x_cpu.is_cuda, id(x_cpu)))

x_gpu = x_cpu.to(device)
print("x_gpu:\ndevice: {} is_cuda: {} id: {}".format(x_gpu.device, x_gpu.is_cuda, id(x_gpu)))

输出如下:

1
2
3
4
x_cpu:
device: cpu is_cuda: False id: 1415020820304
x_gpu:
device: cpu is_cuda: True id: 2700061800153

可以看到Tensorto()方法不是 inplace 操作,x_cpux_gpu的内存地址不一样。

下面代码执行的是Moduleto()方法

1
2
3
4
5
6
net = nn.Sequential(nn.Linear(3, 3))

print("\nid:{} is_cuda: {}".format(id(net), next(net.parameters()).is_cuda))

net.to(device)
print("\nid:{} is_cuda: {}".format(id(net), next(net.parameters()).is_cuda))

输出如下:

1
2
id:2325748158192 is_cuda: False
id:2325748158192 is_cuda: True

可以看到Moduleto()方法是 inplace 操作,内存地址一样。

torch.cuda常用方法

  • torch.cuda.device_count():返回当前可见可用的 GPU 数量
  • torch.cuda.get_device_name():获取 GPU 名称
  • torch.cuda.manual_seed():为当前 GPU 设置随机种子
  • torch.cuda.manual_seed_all():为所有可见 GPU 设置随机种子
  • torch.cuda.set_device():设置主 GPU 为哪一个物理 GPU,此方法不推荐使用
  • os.environ.setdefault("CUDA_VISIBLE_DEVICES", "2", "3"):设置可见 GPU

在 PyTorch 中,有物理 GPU 可以逻辑 GPU 之分,可以设置它们之间的对应关系。


在上图中,如果执行了os.environ.setdefault("CUDA_VISIBLE_DEVICES", "2", "3"),那么可见 GPU 数量只有 2 个。对应关系如下:

逻辑 GPU 物理 GPU
gpu0 gpu2
gpu1 gpu3

如果执行了os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0", "3", "2"),那么可见 GPU 数量只有 3 个。对应关系如下:

逻辑 GPU 物理 GPU
gpu0 gpu0
gpu1 gpu3
gpu2 gpu2

设置的原因是可能系统中有很多用户和任务在使用 GPU,设置 GPU 编号,可以合理分配 GPU。通常默认gpu0为主 GPU。主 GPU 的概念与多 GPU 的分发并行机制有关。

多 GPU 的分发并行

1
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

功能:包装模型,实现分发并行机制。可以把数据平均分发到各个 GPU 上,每个 GPU 实际的数据量为 \(\frac{batch_size}{GPU 数量}\),实现并行计算。

主要参数:

  • module:需要包装分发的模型
  • device_ids:可分发的 GPU,默认分发到所有可见可用的 GPU
  • output_device:结果输出设备

需要注意的是:使用 DataParallel 时,device 要指定某个 GPU 为 主 GPU,否则会报错:

1
RuntimeError: module must have its parameters and buffers on device cuda:1 (device_ids[0]) but found one of them on device: cuda:2

这是因为,使用多 GPU 需要有一个主 GPU,来把每个 batch 的数据分发到每个 GPU,并从每个 GPU 收集计算好的结果。如果不指定主 GPU,那么数据就直接分发到每个 GPU,会造成有些数据在某个 GPU,而另一部分数据在其他 GPU,计算出错。

详情请参考 [RuntimeError: module must have its parameters and buffers on device cuda:1 (device_ids[0]) but found one of them on device: cuda:2](https://stackoverflow.com/questions/59249563/runtimeerror-module-must-have-its-parameters-and-buffers-on-device-cuda1-devi)

下面的代码设置两个可见 GPU,batch_size 为 2,那么每个 GPU 每个 batch 拿到的数据数量为 8,在模型的前向传播中打印数据的数量。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
# 设置 2 个可见 GPU
gpu_list = [0,1]
gpu_list_str = ','.join(map(str, gpu_list))
os.environ.setdefault("CUDA_VISIBLE_DEVICES", gpu_list_str)
# 这里注意,需要指定一个 GPU 作为主 GPU。
# 否则会报错:module must have its parameters and buffers on device cuda:1 (device_ids[0]) but found one of them on device: cuda:2
# 参考:https://stackoverflow.com/questions/59249563/runtimeerror-module-must-have-its-parameters-and-buffers-on-device-cuda1-devi
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

batch_size = 16

# data
inputs = torch.randn(batch_size, 3)
labels = torch.randn(batch_size, 3)

inputs, labels = inputs.to(device), labels.to(device)

# model
net = FooNet(neural_num=3, layers=3)
net = nn.DataParallel(net)
net.to(device)

# training
for epoch in range(1):

outputs = net(inputs)

print("model outputs.size: {}".format(outputs.size()))

print("CUDA_VISIBLE_DEVICES :{}".format(os.environ["CUDA_VISIBLE_DEVICES"]))
print("device_count :{}".format(torch.cuda.device_count()))

输出如下:

1
2
3
4
batch size in forward: 8
model outputs.size: torch.Size([16, 3])
CUDA_VISIBLE_DEVICES :0,1
device_count :2

下面的代码是根据 GPU 剩余内存来排序。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def get_gpu_memory():
import platform
if 'Windows' != platform.system():
import os
os.system('nvidia-smi -q -d Memory | grep -A4 GPU | grep Free > tmp.txt')
memory_gpu = [int(x.split()[2]) for x in open('tmp.txt', 'r').readlines()]
os.system('rm tmp.txt')
else:
memory_gpu = False
print("显存计算功能暂不支持windows操作系统")
return memory_gpu


gpu_memory = get_gpu_memory()
if not gpu_memory:
print("\ngpu free memory: {}".format(gpu_memory))
gpu_list = np.argsort(gpu_memory)[::-1]

gpu_list_str = ','.join(map(str, gpu_list))
os.environ.setdefault("CUDA_VISIBLE_DEVICES", gpu_list_str)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

其中nvidia-smi -q -d Memory是查询所有 GPU 的内存信息,-q表示查询,-d是指定查询的内容。

nvidia-smi -q -d Memory | grep -A4 GPU是截取 GPU 开始的 4 行,如下:

1
2
3
4
5
6
7
8
9
10
11
12
Attached GPUs                       : 2
GPU 00000000:1A:00.0
FB Memory Usage
Total : 24220 MiB
Used : 845 MiB
Free : 23375 MiB
--
GPU 00000000:68:00.0
FB Memory Usage
Total : 24217 MiB
Used : 50 MiB
Free : 24167 MiB

nvidia-smi -q -d Memory | grep -A4 GPU | grep Free是提取Free所在的行,也就是提取剩余内存的信息,如下:

1
2
Free                        : 23375 MiB
Free : 24167 MiB

nvidia-smi -q -d Memory | grep -A4 GPU | grep Free > tmp.txt是把剩余内存的信息保存到tmp.txt中。

[int(x.split()[2]) for x in open('tmp.txt', 'r').readlines()]是用列表表达式对每行进行处理。

假设x=" Free : 23375 MiB",那么x.split()默认以空格分割,结果是:

1
['Free', ':', '23375', 'MiB']

x.split()[2]的结果是23375

假设gpu_memory=['5','9','3']np.argsort(gpu_memory)的结果是array([2, 0, 1], dtype=int64),是从小到大取排好序后的索引。np.argsort(gpu_memory)[::-1]的结果是array([1, 0, 2], dtype=int64),也就是把元素的顺序反过来。

在 Python 中,list[<start>:<stop>:<step>]表示从startstop取出元素,间隔为stepstep=-1表示从stopstart取出元素。start默认为第一个元素的位置,stop默认为最后一个元素的位置。

','.join(map(str, gpu_list))的结果是'1,0,2'

最后os.environ.setdefault("CUDA_VISIBLE_DEVICES", gpu_list_str)就是根据 GPU 剩余内存从大到小设置对应关系,这样默认最大剩余内存的 GPU 为主 GPU。

提高 GPU 的利用率

nvidia-smi命令查看可以 GPU 的利用率,如下图所示。


上面的截图中,有两张显卡(GPU),其中上半部分显示的是显卡的信息下半部分显示的是每张显卡运行的进程。可以看到编号为 0 的 GPU 运行的是 PID 为 14383 进程。Memory Usage表示显存的使用率,编号为 0 的 GPU 使用了 16555 MB 显存,显存的利用率大概是70% 左右。Volatile GPU-Util表示计算 GPU 实际运算能力的利用率,编号为 0 的 GPU 只有 27% 的使用率。

虽然使用 GPU 可以加速训练模型,但是如果GPU 的 Memory UsageVolatile GPU-Util 太低,表示并没有充分利用 GPU。

因此,使用 GPU 训练模型,需要尽量提高 GPU 的 Memory UsageVolatile GPU-Util 这两个指标,可以更进一步加速你的训练过程。

下面谈谈如何提高这两个指标。

Memory Usage

这个指标是由数据量主要是由模型大小,以及数据量的大小决定的。

模型大小是由网络的参数和网络结构决定的,模型越大,训练反而越慢。

我们主要调整的是每个 batch 训练的数据量的大小,也就是 batch_size

在模型结构固定的情况下,尽量将batch size设置得比较大,充分利用 GPU 的内存。

Volatile GPU-Util

上面设置比较大的 batch size可以提高 GPU 的内存使用率,却不一定能提高 GPU 运算单元的使用率。

从前面可以看到,我们的数据首先读取到 CPU 中的,并在循环训练的时候,通过tensor.to()方法从 CPU 加载到 CPU 中,如下代码所示。

1
2
3
4
5
6
7
8
# 遍历 train_loader 取数据
for i, data in enumerate(train_loader):
inputs, labels = data
inputs = inputs.to(device) # 把数据从 CPU 加载到 GPU
labels = labels.to(device) # 把数据从 CPU 加载到 GPU
.
.
.

如果batch size得比较大,那么在 DatasetDataLoader ,CPU 处理一个 batch 的数据就会很慢,这时你会发现Volatile GPU-Util的值会在 0%,20%,70%,95%,0% 之间不断变化。

nvidia-smi命令查看可以 GPU 的利用率,但不能动态刷新显示。如果你想每隔一秒刷新显示 GPU 信息,可以使用watch -n 1 nvidia-smi

其实这是因为 GPU 处理数据非常快,而 CPU 处理数据较慢。GPU 每接收到一个 batch 的数据,使用率就跳到逐渐升高,处理完这个 batch 的数据后,使用率又逐渐降低,等到 CPU 把下一个 batch 的数据传过来。

解决方法是:设置 Dataloader的两个参数:

  • num_workers:默认只使用一个 CPU 读取和处理数据。可以设置为 4、8、16 等参数。但线程数并不是越大越好。因为,多核处理需要把数据分发到每个 CPU,处理完成后需要从多个 CPU 收集数据,这个过程也是需要时间的。如果设置num_workers过大,分发和收集数据等操作占用了太多时间,反而会降低效率。
  • pin_memory:如果内存较大,建议设置为 True
    • 设置为 True,表示把数据直接映射到 GPU 的相关内存块上,省掉了一点数据传输时间。
    • 设置为 False,表示从 CPU 传入到缓存 RAM 里面,再给传输到 GPU 上。

GPU 相关的报错

1.

如果模型是在 GPU 上保存的,在无 GPU 设备上加载模型时torch.load(path_state_dict),会出现下面的报错:

1
RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.

可能的原因:gpu训练的模型保存后,在无gpu设备上无法直接加载。解决方法是设置map_location="cpu"torch.load(path_state_dict, map_location="cpu")

2.

如果模型经过net = nn.DataParallel(net)包装后,那么所有网络层的名称前面都会加上mmodule.。保存模型后再次加载时没有使用nn.DataParallel()包装,就会加载失败,因为state_dict中参数的名称对应不上。

1
2
3
Missing key(s) in state_dict: xxxxxxxxxx

Unexpected key(s) in state_dict:xxxxxxxxxx

解决方法是加载参数后,遍历 state_dict 的参数,如果名字是以module.开头,则去掉module.。代码如下:

1
2
3
4
5
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
namekey = k[7:] if k.startswith('module.') else k
new_state_dict[namekey] = v

然后再把参数加载到模型中。

参考资料

  • 深度之眼 PyTorch 框架班
  • 深度学习模型 GPU 利用率低的原因:https://blog.csdn.net/qq_32998593/article/details/92849585
  • Pytorch数据加载的分析:https://zhuanlan.zhihu.com/p/100762487


如果你觉得这篇文章对你有帮助,不妨点个赞,让我有更多动力写出好文章。

我的文章会首发在公众号上,欢迎扫码关注我的公众号张贤同学


评论