pytorch 如何解决RuntimeError: CUDA out of memory

出现的原因

这个提醒的意思就是运行过程中所需的内存超过CUDA的内存

解决办法

1.减少batch_size
最直接的解决办法是减小batch_size(常用),或者减小模型参数和输入大小(一般很少这样做);
降低batch_size的同时，如果不想影响训练效果（太低的batch_size可能会降低精度），可以迭代一定次数后再更新参数，迭代的过程其会自动累加梯度，不用每次迭代都更新参数，如下图的代码所示：
比如你的GPU显存只能装的下small_batch，然后迭代large_batch后再更新参数，相当于batch_size=small_batch*large_batch;

train_data = DataLoader(train_set, batch_size=small_batch, ...)

opt.zero_grad()  # this signifies the start of a large_batch
for i, (x, y) in train_data:
  pred = model(x)
  loss = criterion(pred, y)
  loss.backward()  # gradeints computed for small_batch
  if (i+1) % k == 0 or (i+1) == len(train_data):
    opt.step()  # update the weights only after accumulating k small batches
    opt.zero_grad()  # reset gradients for accumulation for the next large_batch

2.with torch.no_grad():
如果测试过程遇到这种情况,加上
with torch.no_grad():
内存就不会分配参数梯度的空间
3.torch.cuda.empty_cache()
如果在训练过程遇到这种情况,可以尝试在训练前先释放CUDA内存
nvidia-smi查看GPU使用率,如果使用率不高,就使用torch.cuda.empty_cache()释放内存
官网对于torch.cuda.empty_cache()的说明:
Releases all unoccupied cached memory currently held by the caching allocator so that those can be used in other GPU application and visible in nvidia-smi.
释放当前由缓存分配器保留的所有未占用的缓存内存，以便这些内存可在其他GPU应用程序中使用，并在nvidia-smi中可见 。
该语句是释放没有被使用的缓存,所以可以放心大胆地使用.
4.killall python
观察 报错中是不是大部分内存都 reserved by pytorch,如果这样，可能CUDA内存被其他python应用占用或者锁住了；
直接执行 killall python
或者

ps -elf | grep python 

or

ps -aux | grep python 

找到该进程，然后强制杀死他

kill -9 PID

5.分批次加载数据集
如果一次性加载的数据集过大，也会占用大量的cuda内存；
我们可以分批次的加载数据集
SubsetRandomSampler就是无放回地抽样加载数据集，是pytorch自带的函数；

    trainDataset = KFDataset(config,"ubuntu_new_dataset_image.txt","ubuntu_new_dataset_label.txt")
    train_dataset_size = len(trainDataset)
    train_indices = list(range(train_dataset_size))
    random_seed = 42
    np.random.seed(random_seed)
    np.random.shuffle(train_indices)
    train_sampler = torch.utils.data.SubsetRandomSampler(train_indices)
    trainDataLoader = DataLoader(trainDataset, batch_size=config['batch_size'],
                                 sampler=train_sampler)
    for epoch in range(config['start_epoch'],config['epoch_num']+config['start_epoch']):
        for i, (inputs, heatmaps_targets, gts) in enumerate(trainDataLoader):
             #执行训练部分

6.Dataloader中的pin_memory=True
如果您将样本加载Dataset到CPU上，并希望在训练过程中将其推送到GPU，则可以通过启用加快主机到设备的传输速度pin_memory。

7.降精度训练或者混合精度训练
可以降为fp16的精度训练；
Nvidia AMP - 自动混合精度的惊人结果，这是 fp16 与 fp32 优势的完美结合。

8.负载均衡的方式
当我们在单机多卡上训练模型的时候，会出现GPU负载不均衡的情况，主要原因是所有batch的loss计算都会放在０卡上，所以０卡上的显存远远大于其他卡，那么如何解决呢？
有以下三个办法：

• 用DistributedDataParallel代替DataParallel
  具体使用方式可以参考mmdet等开源项目的实现。虽然之前看到好多人说DistributedDataParallel也会导致GPU的负载不均衡，不过反正我没遇到过。

• 重写DataParallel
  因为GPU负载不均衡大概率是0卡占的显存比其他卡多很多，原因主要是因为反传的时候梯度默认都在0卡上算了。所以你可以重写下DataParallel，来减少0卡上的batch_size，增大其他卡上的batch_size，从而实现训练时候的GPU负载均衡。例如总的batch_size是8，你可以在0卡上放2张图，1卡和2卡上放3张图，来实现GPU的负载均衡。这个方案可以参考：transformer-xl的实现。

• 分布化损失函数
  原理大概是让梯度在各个 GPU 上单独计算和反向传播，从而避免GPU的负载不均衡（和2一样，我们假设负载不均衡产生的原因是反传的时候梯度默认都在0卡上算了），这个方案可以参考：Pytorch-Encoding的实现。如果嫌麻烦的话，楼主可以直接在每个module里算loss，然后把返回的loss求mean做反传，注意训练模式下返回的时候别把预测结果一起返回，即只把loss给return出来就行了，不然负载还是可能很不均衡的。