[PyTorch 學習筆記] 7.3 使用 GPU 訓練模型

本章代碼：

• https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson7/cuda_use.py
• https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson7/multi_gpu.py

這篇文章主要介紹了 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：結果輸出設備

需要注意的是：使用 DataParallel 時，device 要指定某個 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

這是因為，使用多 GPU 需要有一個主 GPU，來把每個 batch 的數據分發到每個 GPU，並從每個 GPU 收集計算好的結果。如果不指定主 GPU，那么數據就直接分發到每個 GPU，會造成有些數據在某個 GPU，而另一部分數據在其他 GPU，計算出錯。

詳情請參考 [[RuntimeError: module must have its parameters and buffers on device cuda:1 (device_ids0]) but found one of them on device: cuda:2]([RuntimeError: module must have its parameters and buffers on device cuda:1 (device_ids0]) but found one of them on device: cuda:2)

下面的代碼設置兩個可見 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)
    # 這里注意，需要指定一個 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()))

輸出如下：

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 的利用率

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 Usage 和 Volatile GPU-Util 太低，表示並沒有充分利用 GPU。

因此，使用 GPU 訓練模型，需要盡量提高 GPU 的 Memory Usage 和 Volatile GPU-Util 這兩個指標，可以更進一步加速你的訓練過程。

下面談談如何提高這兩個指標。

Memory Usage

這個指標是由數據量主要是由模型大小，以及數據量的大小決定的。

模型大小是由網絡的參數和網絡結構決定的，模型越大，訓練反而越慢。

我們主要調整的是每個 batch 訓練的數據量的大小，也就是 batch_size。

在模型結構固定的情況下，盡量將batch size設置得比較大，充分利用 GPU 的內存。

Volatile GPU-Util

上面設置比較大的 batch size可以提高 GPU 的內存使用率，卻不一定能提高 GPU 運算單元的使用率。

從前面可以看到，我們的數據首先讀取到 CPU 中的，並在循環訓練的時候，通過tensor.to()方法從 CPU 加載到 CPU 中，如下代碼所示。

# 遍歷 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得比較大，那么在 Dataset和 DataLoader ，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),會出現下面的報錯：

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

然后再把參數加載到模型中。

參考資料

• 深度之眼 PyTorch 框架班
• 深度學習模型 GPU 利用率低的原因：https://blog.csdn.net/qq_32998593/article/details/92849585
• Pytorch 數據加載的分析：https://zhuanlan.zhihu.com/p/100762487

如果你覺得這篇文章對你有幫助，不妨點個贊，讓我有更多動力寫出好文章。