【报错解决】深度学习模型训练时cuda内存足够但测试时反而报错cuda out of memory

报错描述

报错的代码如下：

model = reader(config=args, encoder=encoder)#初始化模型
model.to('cuda')#把模型放到gpu上
model.load_state_dict(torch.load(join(args.checkpoint_path, 'best_ckpt_model1.pkl')))#加载模型参数
model = torch.nn.DataParallel(model)#并行化处理

具体的报错解决如下所示：

Traceback (most recent call last):
  File "run_cail.py", line 314, in <module>
    best_join_f1=train_epoch(Loader, model1, model2,best_join_f1=best_join_f1)
  File "run_cail.py", line 132, in train_epoch
    train_bh(model1,model2, batch)
  File "run_cail.py", line 164, in train_bh
    type_logits, qc_out, sp_logits, att = model1(batch)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 159, in forward
    return self.module(*inputs[0], **kwargs[0])
  File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/EBF_Reader-master/EBF_Reader-master/EBF Reader/model/modeling.py", line 104, in forward
    qc_out=add_word_emb_into_char_emb(qc_out,word_ids,word_embedding)
  File "/root/EBF_Reader-master/EBF_Reader-master/EBF Reader/model/modeling.py", line 24, in add_word_emb_into_char_emb
    word_id=word_ids[id][j].item()
RuntimeError: CUDA error: device-side assert triggered
terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: device-side assert triggered
Exception raised from create_event_internal at /pytorch/c10/cuda/CUDACachingAllocator.cpp:687 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x42 (0x7f8c3e3298b2 in /root/miniconda3/lib/python3.8/site-packages/torch/lib/libc10.so)
frame #1: c10::cuda::CUDACachingAllocator::raw_delete(void*) + 0xad2 (0x7f8c3e57b952 in /root/miniconda3/lib/python3.8/site-packages/torch/lib/libc10_cuda.so)
frame #2: c10::TensorImpl::release_resources() + 0x4d (0x7f8c3e314b7d in /root/miniconda3/lib/python3.8/site-packages/torch/lib/libc10.so)
frame #3: <unknown function> + 0x5fa0a2 (0x7f8cd21fd0a2 in /root/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_python.so)
frame #4: <unknown function> + 0x5fa156 (0x7f8cd21fd156 in /root/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>
frame #21: __libc_start_main + 0xe7 (0x7f8cfab0eb97 in /lib/x86_64-linux-gnu/libc.so.6)

Fatal Python error: Aborted

Current thread 0x00007f8cfbaae2c0 (most recent call first):
<no Python frame>
run_cail2020.sh: line 19:   891 Aborted                 (core dumped) python run_cail.py --name train_v1 --bert_model $bert_dir --data_dir data --val_dir data_2020 --batch_size 4 --eval_batch_size 4 --lr 3e-5 --seed 1 --k_v_dim 256 --att_back_lambda 3.0 --gradient_accumulation_steps 1 --pre_each_epc 5 --epc_start_pre 7 --epochs 10

报错原因

报错的原理，简单来说，就是：

        初始化模型后将模型直接放到gpu上，然后先将模型参数加载到gpu上然后再加载带参数的模型，这就导致gpu上加载了两次模型参数。而深度学习模型的参数量一般是比较庞大的，因此，这就可能会造成gpu内存不够，进而在测试时报cuda out of memory的错。

详细说，就是：

        torch.load(model_path,map_location)函数的第一个参数model_path是你保存的模型所在路径，第二个参数是用于控制将模型加载到什么位置，默认是加载到cuda:0即第一块gpu上，可以通过改变该参数的值来控制模型加载到哪里，例如，令map_location='cpu'即可令模型加载到cpu上。

        回来接着说原因，我上面的报错代码中只有第一个参数model_path，而没有设置map_location的值，所以其直接默认加载到gpu上。就是说，torch.load函数将模型参数加载到了gpu上，这是gpu上的第一份模型参数。而函数model.load_state_dict函数又将模型加载了一次，也就是gpu上出现了第二份模型参数，这就造成了重复，有可能会造成gpu内存不够而报错。

这下应该明白为什么会训练的时候cuda内存足够，而测试的时候cuda内存不够报错了吧，就是因为重复加载模型参数了。

解决办法

我几乎尝试了网上能找到的所有解决办法，都没有用，但是最后还是让我找到解决办法了哈哈，这里来记录一下，顺便也希望能帮助到遇到同样问题的友友们：）

将报错代码改成下面这个样子就可以成功运行了：

model = reader(config=args, encoder=encoder)
model.load_state_dict(torch.load(join(args.checkpoint_path,'best_ckpt_model1.pkl'),map_location='cpu'))
model.to('cuda')
model = torch.nn.DataParallel(model)

下面来详细解释一下原因，就是为什么这样做cuda内存就够了，不会报错了。

第一行代码初始化模型，在cpu上进行，第二行代码，首先torch.load函数将模型参数加载到cpu上，然后model.load_state_dict函数加载模型到cpu上，第三行代码直接将整个模型放到gpu上，避免了重复加载模型参数，节省了gpu内存，就不会再造成cuda内存不够而报错了，第四行代码是将模型进行并行化处理，也就是说，如果你有多块gpu的话，这行代码可以将模型复制到多块gpu上进行并行化计算，当然如果你只有一块gpu或只有cpu，加了这行代码也没什么影响，不会报错。

看到这里，你应该明白为什么会出现深度学习模型训练时cuda内存够但测试时cuda内存不够的问题了吧，也知道怎么解决了吧。如果有用的话，希望可以不吝啬地点赞收藏夹关注啊~

这里是希望你能越来越好的 小白冲鸭 ~~~