重新加载 Pytorch 模型时出现 CUDA 内存不足错误

【问题标题】：CUDA out of memory error when reloading Pytorch model重新加载 Pytorch 模型时出现 CUDA 内存不足错误
【发布时间】：2021-10-10 17:10:24
【问题描述】：

这里是常见的 pytorch 错误，但我在特殊情况下看到它：重新加载模型时，我收到 CUDA: Out of Memory 错误，即使我还没有将模型放在 GPU 上。

model = model.load_state_dict(torch.load(model_file_path))
optimizer = optimizer.load_state_dict(torch.load(optimizer_file_path))
# Error happens here ^, before I send the model to the device.
model = model.to(device_id)

【问题讨论】：

标签： memory pytorch gpu out-of-memory

【解决方案1】：

问题是我试图加载到新的 GPU (cuda:2) 但最初是从不同的 GPU (cuda:0) 保存模型和优化器。所以即使我没有明确告诉它重新加载到之前的 GPU，默认行为是重新加载到原来的 GPU（恰好被占用）。

将map_location=device_id 添加到每个torch.load 调用解决了问题：

model.to(device_id)
model = model.load_state_dict(torch.load(model_file_path, map_location=device_id))
optimizer = optimizer.load_state_dict(torch.load(optimizer_file_path, map_location=device_id))

【讨论】：