文章目录

  • 原文链接
  • 1.背景
  • 2.问题发现
  • 3.结论
  • 4.补充说明



原文链接

1.背景

记录一个在工作中偶然发现的情况。接触过深度学习的朋友的都知道,在使用GPU(显卡)训练模型时,影响模型训练的关键因素有两个:算力和显存。其中,算力影响模型训练的速度,显存决定了每个batch喂给GPU训练的数据量的上限。

因为工作内容的原因,经常需要将高分辨率的图像数据作为数据集,对模型进行训练。为了提高模型对高分辨率图像特征的学习能力,每个batch只喂一张图给GPU。经常在跑模型的时候都是把显存拉满了,偶尔还会报OOM(显存溢出)的错误,着实头痛,因为大显存的显卡都比较贵,经费有限,没有使用像RTXA6000这样的48GB显存以上的显卡(目前使用的最大显存为24GB,型号:RTX 4090)。有的朋友可能想到了多卡并联,但是多卡并联不能把一张图分给多个显卡同时训练(目前没有找到相关的解决方案,也许有)。

2.问题发现

在前一段时间更新了Nvidia的显卡驱动,版本从531升级到了536。之后训练模型的时候,突然就发现当显存溢出之后可以调用主板的内存了,发现问题的时候比较疑惑,并不知道是更新了驱动之后的影响(因为之前也经常更新驱动)。后面经过排除法,又用其他计算机做了测试,才发现是显卡驱动版本更新带来的优化。得知这个结果后还是比较高兴的,这样在显存偶然溢出的时候就不会报OOM的错误了,而且现在有ddr5的高频内存支持下,对训练速度要求不是特别高的情况下,可以暂时用共享GPU内存作为补充。

3.结论

后面又在不同的显卡和计算机上做了测试,目前的结果如下:
(1) 在Nvidia的10系、20系、30系、40系的显卡上更新驱动到536版本后均可实现共享GPU内存进行模型训练,536以上版本测试了随机测试了几个也可以(应该536版本以上驱动均支持该功能)。
(2) 在Nvidia的RTX A和quadro上的显卡上测试536的驱动不能实现调用共享GPU内存进行模型训练(主要测试了RTx A2000 、RTX A5000显卡等Ampere架构的显卡),而且该系列的多数显卡最新版本只有536的驱动。

关于536版本驱动的详细跟新说明链接如下,有兴趣的朋友可以去看看:https://us.download.nvidia.cn/Windows/536.99/536.99-win11-win10-release-notes.pdf

4.补充说明

(1) 以上测试均在win10平台中进行
(2) 深度学习的框架用的pytorch 1.11.0
(3) cuda版本11.1