pytorch前向计算显存占用太大，如何破？

东方耀

有一个pytorch版的人脸检测模型  模型本身只有7.4M


但是在进行前向计算时 显存的占用非常大 输入图片比较小还好 一旦大点的图片 马上报CUDA显存不足 退出啦
RuntimeError: CUDA out of memory


经测：
模型输入的x： torch.Size([1, 3, 416, 320])   显存占用1025M
模型输入的x： torch.Size([1, 3, 768, 1024])  显存占用1805M
模型输入的x： torch.Size([1, 3, 1152, 2048])  显存占用3819M

模型输入的x： torch.Size([1, 3, 1088, 1920])   显存占用3486M
模型输入的x： torch.Size([1, 3, 3648, 2752])   RuntimeError: CUDA out of memory





解决方法:
在模型预测的代码上包一下：
with torch.no_grad():
    predict = model(inputs)


with torch.no_grad():
    hm, box, landmark = model(torch_image)


加了之后经测：
模型输入的x： torch.Size([1, 3, 416, 320])   显存占用905M
模型输入的x： torch.Size([1, 3, 768, 1024])  显存占用1027M
模型输入的x： torch.Size([1, 3, 1152, 2048])  显存占用1343M
模型输入的x： torch.Size([1, 3, 1088, 1920])   显存占用1394M

模型输入的x： torch.Size([1, 3, 3648, 2752])   显存占用2843M