如果保存模型在 GPU 上,加载模型在 CPU 上,那么我们的保存与加载的代码应该如下:
# Save on GPU
device = torch.device("cuda")
model.to(device)
torch.save(model.state_dict(), PATH)
# Load on CPU
device = torch.device('cpu')
model = Model(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
如果保存模型在 GPU 上,加载模型在 GPU 上,那么我们的保存与加载的代码应该如下:
device = torch.device("cuda")
model.to(device)
torch.save(model.state_dict(), PATH)
model = Model(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
如果保存模型在 CPU 上,加载模型在 GPU 上,那么我们的保存与加载的代码应该如下:
torch.save(model.state_dict(), PATH)
device = torch.device("cuda")
model = Model(*args, **kwargs)
# 选择加载到哪一个GPU设备上
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))
model.to(device)
如果保存模型在 CPU 上,加载模型在 CPU 上,那么我们的保存与加载的代码应该如下:
torch.save(model.state_dict(), PATH)
model = Model(*args, **kwargs)
model.load_state_dict(torch.load(PATH))