1 问题描述

用Pytorch进行模型训练时出现以下OOM提示：

RuntimeError: CUDA out of memory. Tried to allocate 98.00 MiB (GPU 0; 12.00 GiB total capacity; 3.19 GiB already allocated; 6.40 GiB free; 9.60 GiB allowed; 3.33 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

CUDA尝试分配98MB显存，但失败了。GPU的显存为12G，其中已分配显存3.19G（允许分配上限为9.6G），仍有6.4G的显存空闲。

简言之，显存有一半以上的空间都是闲置的，却连区区98MB都挤不出来，这显然不对劲！

2 解决方案

先说结论：通过降低num_workers解决。 具体方法见2.3节。

根据网络上相关问题的经验贴，本人进行了一系列的尝试，现将相关解决方法记录如下。

2.1 确认是否因找不到GPU资源引起

参考：找不到GPU资源——显存充足，但是却出现CUDA error:out of memory错误_gpu out of memory-CSDN博客

尝试手工指定GPU资源：

1、通过python代码配置

import os
os.environ['CUDA_VISIBLE_DEVICES']='0'

2、通过设置环境配置

• Linux：export CUDA_VISIBLE_DEVICES=0
• Windows：setx CUDA_VISIBLE_DEVICES=0

进行如上配置后问题依旧存在。

2.2 调整max_split_size_mb参数

参考：CUDA oom 通过设置PYTORCH_CUDA_ALLOC_CONF中的max_split_size_mb解决Pytorch的显存碎片化导致的CUDA:Out Of Memory - 知乎 (zhihu.com)

pytoch的显存管理中，分配显存请求必须是连续的，max_split_size_mb设置的是可分割最大的空闲block，小于该值的空闲block可能由于被分割而无法连续使用，大于该值的空闲block将不会被分割。比如max_split_size_mb 设置为4000时，所有小于4000MB空闲block都可能被分割开，当需要连续4g的空间时，就不存在可分配的4g的连续空闲block，而报OOM错误。

最优设置策略：将max_split_size_mb设置为小于OOM发生时的显存请求大小最小值的最大整数值。

基于以上理论，本人尝试根据OOM提示将max_split_size_mb参数设置为98：

1、方法一

import os 
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:98"

2、方法二

• Linux：export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:98
• windows：setx PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:98

配置后问题依旧存在，也就是说目前显卡内存已经有能力切分出98MB大小的存储空间，却无法提供。

想了半天没搞明白究竟是怎么回事，于是想着检查下配置看看，打开配置文件，目光立即被那行num_workers=6吸引，顿时恍然大悟，请见下一节。

2.3 调整锁页配置

关于锁页的介绍可以看这篇文章：Pytorch DataLoader pin_memory 理解 - 知乎 (zhihu.com)。

简言之，锁页是一种允许GPU直接访问CPU存储空间的机制，通过提前锁定一部分存储空间来解决分页存储频繁换入换出的问题，从而提升访问速度。

之前为了提升模型的训练速度，开启了锁页机制，并设置num_workers=8。

合理猜测，显卡之所以无法提供模型需要的98MB空间，是因为空闲内存被锁定造成的。

通过缩小num_workers的数值，问题解决：

pin_memory = True
num_workers = 6