PyTorch/PyTorch GPU 使用

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

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

• 从 CPU 到 GPU

  device = torch.device("cuda")
  tensor = tensor.to(device)
  module.to(device)

• 从 GPU 到 CPU

  device = torch.device(cpu)
  tensor = tensor.to("cpu")
  module.to("cpu")

  tensor和module的 to()方法的区别是：tensor.to()执行的不是 inplace 操作，因此需要赋值；module.to()执行的是 inplace 操作。

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

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

tensor.to() 和 module.to()

首先导入库，获取 GPU 的 device

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

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

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)))

输出如下：

x_cpu:
device: cpu is_cuda: False id: 1415020820304
x_gpu:
device: cpu is_cuda: True id: 2700061800153

可以看到Tensor的to()方法不是 inplace 操作，x_cpu和x_gpu的内存地址不一样。

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

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))

输出如下：

id:2325748158192 is_cuda: False
id:1756341802643 is_cuda: True

可以看到Module的to()方法是 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 的分发并行

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：结果输出设备

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

# 设置 2 个可见 GPU
gpu_list = [0,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")

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()))

输出如下：

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

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

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 行，如下：

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所在的行，也就是提取剩余内存的信息，如下：

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()默认以空格分割，结果是：

['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>]表示从start到stop取出元素，间隔为step，step=-1表示从stop到start取出元素。start默认为第一个元素的位置，stop默认为最后一个元素的位置。

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

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

GPU 相关的报错

1.

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

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中参数的名称对应不上。

Missing key(s) in state_dict: xxxxxxxxxx

Unexpected key(s) in state_dict:xxxxxxxxxx

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

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

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