pytorch 训练/测试模型时错误：RuntimeError: CUDA error: out of memory

Turbo_Come

于 2020-04-29 11:19:00 发布

阅读量2.4k

点赞数 3

CC 4.0 BY-SA版权

文章标签：神经网络

本文链接：https://blog.youkuaiyun.com/Turbo_Come/article/details/105835395

本文分享了三种有效降低GPU内存消耗的方法：1.通过调整batch-size大小；2.使用torch.no_grad()禁用梯度计算；3.将模型加载到CPU上运行，避免GPU内存溢出，适用于深度学习项目。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

方法1：batch-size设置多小

方法2：

with torch.no_grad():

       net = Net()

       out = net(imgs)

积累的梯度应该是会一直放在显存里的...用了这一行就会停止自动反向计算梯度

方法3：

设置cpu来加载模型：

model_path = 'path/to/model.pt' 

model = UNet(n_channels = 1, n_classes = 1) 

state_dict = torch.load(model_path,map_location='cpu') 

model.load_state_dict(state_dict) 

model.to(device)