这是我的第394篇原创文章。
一、引言
CUDA out of memory问题通常发生在深度学习训练过程中，当GPU的显存不足以容纳模型、输入数据以及中间计算结果时就会触发。这个问题可能由几个因素引起：
-
模型和数据规模：深度学习模型尤其是大型模型，如Transformer或大型CNN，拥有大量的参数，这些参数在训练时需要被加载到GPU显存中。同时，如果批量大小（batch size）设置得过大，一次性处理的数据量也会增加，进一步加大显存的负担。
-
内存管理：深度学习框架在处理数据和模型时，可能会因为不当的内存管理而导致显存没有被及时释放，或者由于频繁的内存分配和释放造成内存碎片化，这会减少可用的连续内存块，即使总显存使用量没有达到100%，也可能出现内存不足的情况。
-
数据加载和预处理：如果数据加载和预处理步骤没有优化，比如加载了大量未被立即使用的数据，或者数据没有被适当地压缩或降维，也可能导致显存使用量激增。
多人共用服务器：如果服务器由多人共用，可能存在其他用户占用大量显存的情况。你可以尝试切换到其他显存较充足的GPU上运行你的任务，通过设置环境变量CUDA_VISIBLE_DEVICES来实现。
单人使用服务器：若服务器仅由你一人使用，检查是否有其他程序或任务正在同一GPU上运行。如果有，考虑停止或暂停这些任务，释放显存给当前需要的任务。
二、实现过程
2.1 Python脚本中
Python 中，可以通过 os.environ 读取、修改或删除环境变量
os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1, 3"这行代码设置了环境变量 CUDA_VISIBLE_DEVICES，它用于指定可以用于CUDA运算的GPU设备。这里指定了设备ID为0、1和3，这意味着在程序中只会看到这三块GPU，其他的GPU则不可见。这有助于管理多个GPU的使用，避免GPU之间的竞争。
device = torch.device("cuda:0" if torch.cuda.is_available() and not args.no_cuda else "cpu"):这行代码使用 torch.device 来确定当前的计算设备。如果CUDA可用且没有禁用CUDA（例如args.no_cuda为False），就使用第一个GPU(cuda:0)，否则使用CPU。args.no_cuda通常是来自命令行参数的选项，用于允许用户在运行程序时选择是否使用CUDA。
model = torch.nn.DataParallel(model, device_ids=[0, 1, 3])这行代码将模型包装在 torch.nn.DataParallel 中，以实现多GPU并行计算。device_ids 参数指定了要使用哪些GPU，这里是0、1和3（与前面的环境变量设置一致）。这样可以同时在多个GPU上训练模型，提高训练速度。
2.2 Python脚本外
如果遇上多个脚本需要同时使用GPU跑，可以按照下面的方式使用，在python前面添加指定的GPU编号。
CUDA_VISIBLE_DEVICES=1 python script1.py
CUDA_VISIBLE_DEVICES=2 python script2.py 或（Unix/Linux Shell）：
export CUDA_VISIBLE_DEVICES=1
python script1.py
export CUDA_VISIBLE_DEVICES=1
python script1.py或（windows cmd）：
set CUDA_VISIBLE_DEVICES=1
python script1.py
set CUDA_VISIBLE_DEVICES=1
python script1.py作者简介：
读研期间发表6篇SCI数据挖掘相关论文，现在某研究院从事数据算法相关科研工作，结合自身科研实践经历不定期分享关于Python、机器学习、深度学习、人工智能系列基础知识与应用案例。致力于只做原创，以最简单的方式理解和学习，关注我一起交流成长。需要数据集和源码的小伙伴可以关注底部公众号添加作者微信。