memory - 重新加载 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(恰好被占用)。

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