重新加载 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))