CUDA内存不足问题：如何指定空闲GPU运行PyTorch模型

问题详情 #

运行模型测试时触发CUDA error: out of memory错误，错误栈如下：

(val_env) jovyan@jupyter-:~/use-cases/Scene_understanding/Code_Woodscape/deeplab_v3+$ python test_shap.py 
Traceback (most recent call last):
  File "/home/jovyan/use-cases/Scene_understanding/Code_Woodscape/deeplab_v3+/test_shap.py", line 33, in <module>
    model.load_state_dict(torch.load(PATH))
  File "/home/jovyan/.local/lib/python3.9/site-packages/torch/serialization.py", line 712, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/jovyan/.local/lib/python3.9/site-packages/torch/serialization.py", line 1049, in _load
    result = unpickler.load()
  File "/home/jovyan/.local/lib/python3.9/site-packages/torch/serialization.py", line 1019, in persistent_load
    load_tensor(dtype, nbytes, key, _maybe_decode_ascii(location))
  File "/home/jovyan/.local/lib/python3.9/site-packages/torch/serialization.py", line 1001, in load_tensor
    wrap_storage=restore_location(storage, location),
  File "/home/jovyan/.local/lib/python3.9/site-packages/torch/serialization.py", line 175, in default_restore_location
    result = fn(storage, location)
  File "/home/jovyan/.local/lib/python3.9/site-packages/torch/serialization.py", line 157, in _cuda_deserialize
    return obj.cuda(device)
  File "/home/jovyan/.local/lib/python3.9/site-packages/torch/_utils.py", line 78, in _cuda
    return torch._UntypedStorage(self.size(), device=torch.device('cuda')).copy_(self, non_blocking)
RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

当前GPU状态显示cuda:0已被完全占用，cuda:1、2、3有空闲内存：

(val_env) jovyan@jupyter-:~/use-cases/Scene_understanding/Code_Woodscape/deeplab_v3+$ nvidia-smi
Mon Sep  5 14:47:40 2022       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01    Driver Version: 515.65.01    CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100S-PCI...  On   | 00000000:00:06.0 Off |                    0 |
| N/A   42C    P0    38W / 250W |  32504MiB / 32768MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100S-PCI...  On   | 00000000:00:07.0 Off |                    0 |
| N/A   43C    P0    39W / 250W |   7811MiB / 32768MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  Tesla V100S-PCI...  On   | 00000000:00:08.0 Off |                    0 |
| N/A   42C    P0    42W / 250W |    926MiB / 32768MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  Tesla V100S-PCI...  On   | 00000000:00:0B.0 Off |                    0 |
| N/A   43C    P0    36W / 250W |   3437MiB / 32768MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+                                                                                
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

当前运行代码如下：

import shap
import torch
import glob
from deeplab import deeplab_v3plus

PATH = "/home/shared/EC1/Scene_understanding/Dataset_Woodscape/Models/deeplab_v3+/Model_Woodscape.pth"

class Arguments:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.seg_file = "/home/shared/EC1/Scene_understanding/Dataset_Woodscape/Models/deeplab_v3+/Model_Woodscape.pth"
        self.img_folder = glob.glob("test_img/*.png")
        self.mean = [0.485, 0.456, 0.406]
        self.std = [0.229, 0.224, 0.225]
        self.h, self.w = 483, 640
        self.nclass = 10
        self.cmap = {
            1: [128, 64, 128],  # "road",
            2: [69, 76, 11],    # "lanemarks",
            3: [0, 255, 0],     # "curb",
            4: [220, 20, 60],   # "person",
            5: [255, 0, 0],     # "rider",
            6: [0, 0, 142],     # "vehicles",
            7: [119, 11, 32],   # "bicycle",
            8: [0, 0, 230],     # "motorcycle",
            9: [220, 220, 0],   # "traffic_sign",
            0: [0, 0, 0]        # "void"
        }

args = Arguments()

model = deeplab_v3plus('resnet101', num_classes=args.nclass, output_stride=16, pretrained_backbone=True)
model.load_state_dict(torch.load(PATH))

explainer = shap.Explainer(model)
shap_values = explainer(X_frame)

解决方案：指定空闲GPU运行 #

以空闲内存最多的cuda:2为例，通过4处修改实现指定GPU：

1. 修改Arguments类中的device指定：将默认的cuda改为具体GPU编号，确保后续操作统一使用该设备
2. 加载模型时指定map_location：避免torch默认将模型加载到cuda:0，直接映射到目标GPU
3. 将模型移动到指定设备：确保模型参数全部迁移到目标GPU
4. 将输入数据移动到指定设备：保证模型和输入数据在同一设备上

修改后的完整代码：

import shap
import torch
import glob
from deeplab import deeplab_v3plus

PATH = "/home/shared/EC1/Scene_understanding/Dataset_Woodscape/Models/deeplab_v3+/Model_Woodscape.pth"

class Arguments:
    def __init__(self):
        # 指定使用cuda:2，若GPU不可用则 fallback 到CPU
        self.device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
        self.seg_file = "/home/shared/EC1/Scene_understanding/Dataset_Woodscape/Models/deeplab_v3+/Model_Woodscape.pth"
        self.img_folder = glob.glob("test_img/*.png")
        self.mean = [0.485, 0.456, 0.406]
        self.std = [0.229, 0.224, 0.225]
        self.h, self.w = 483, 640
        self.nclass = 10
        self.cmap = {
            1: [128, 64, 128],  # "road",
            2: [69, 76, 11],    # "lanemarks",
            3: [0, 255, 0],     # "curb",
            4: [220, 20, 60],   # "person",
            5: [255, 0, 0],     # "rider",
            6: [0, 0, 142],     # "vehicles",
            7: [119, 11, 32],   # "bicycle",
            8: [0, 0, 230],     # "motorcycle",
            9: [220, 220, 0],   # "traffic_sign",
            0: [0, 0, 0]        # "void"
        }

args = Arguments()

model = deeplab_v3plus('resnet101', num_classes=args.nclass, output_stride=16, pretrained_backbone=True)
# 加载模型时指定map_location到目标设备，避免cuda:0内存不足
model.load_state_dict(torch.load(PATH, map_location=args.device))
# 将模型移动到指定设备
model = model.to(args.device)

# 确保输入数据X_frame也移动到指定设备
X_frame = X_frame.to(args.device)

explainer = shap.Explainer(model)
shap_values = explainer(X_frame)

额外说明 #

• 若cuda:2仍出现内存不足，可替换为cuda:1或cuda:3，根据nvidia-smi的内存占用情况选择最空闲的设备
• 若内存问题持续，可尝试开启模型半精度模式：model = model.half()，同时输入数据也转为半精度：X_frame = X_frame.half()

内容的提问来源于stack exchange，提问作者Revolucion for Monica