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

Question

常见的 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)

Answer 1

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

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

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))

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

CUDA out of memory error when reloading Pytorch model

memory

gpu

out-of-memory

pytorch