已解决。RuntimeError: CUDA out of memory.显存不够时使用预训练模型训练的方法。

近期使用预训练的Resnext101_32x32训练自己的数据集。众所周知,因为resnext离谱的权重参数,模型本身相对大,并且本身我处理的数据类间差距就比较小,所以输入图像的大小不能缩放太小,不然会严重影响精度。

各模型权重文件大小对比。

已解决。RuntimeError: CUDA out of memory.显存不够时使用预训练模型训练的方法。_第1张图片

然后就出现了以下情况:

我本身的GPU是3070,CPU是11代的i7。但是还是出现是显存不足。

我输入图像的大小设置的是400x400,batch_size是16,理论上虽然模型权重大,但是应该还是够用啊!然后在网上找解决方法,绝大部分都是让缩小batch_size或者减小输入图像的大小,虽然确实可以解决一部分的问题,但是在我的这个数据集上,resnet101采用224输入和400输入的准确率top1相差4%左右,top5差了6%。resnext因为显存不足报错了。

但是显然如果能保持400x400的输入,resnext的最终结果显然会更好。那到底要如何解决呢?

尝试了多种方法,最后发现,我的训练代码中选择了对全部的权重参数进行训练。然而实际情况中,预训练模型一些颜色,形状等特征识别的层早已经训练好了,完全可以将这一部分层冻结。

ct = 0
for param in model.parameters():
    ct += 1
    param.requires_grad = False
    if ct >= 50:
    #冻结前50层不训练
        continue
# 前面的backbone保持不变

随后运行!!!

终于开始训练了。

希望能对大家有所帮助,求点赞!!! 

你可能感兴趣的:(python,深度学习,pytorch,图像处理)