pytorch跨设备保存和加载模型（变量类型(cpu/gpu)不匹配原因之一）

最新推荐文章于 2025-06-17 09:00:04 发布

转载最新推荐文章于 2025-06-17 09:00:04 发布 · 2k 阅读

深度学习专栏收录该内容

12 篇文章

订阅专栏

本文详细介绍了如何在不同设备（CPU和GPU）间保存和加载PyTorch模型，包括模型状态字典的保存与加载过程，以及在不同场景下如何正确设置设备和参数。

部署运行你感兴趣的模型镜像

原官网教程地址

跨设备保存和加载模型

保存在GPU上，在CPU上加载

保存：

torch.save(model.state_dict(), PATH)

加载：

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

在使用GPU训练的CPU上加载模型时，请传递 torch.device('cpu')给map_location函数中的 torch.load()参数。在这种情况下，使用map_location参数将张量下面的存储器动态地重新映射到CPU设备。

保存在GPU上，在GPU上加载

保存：

torch.save(model.state_dict(), PATH)

加载：

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

在GPU上训练并保存在GPU上的模型时，只需将初始化model模型转换为CUDA优化模型即可model.to(torch.device('cuda'))。此外，请务必.to(torch.device('cuda'))在所有模型输入上使用该功能来准备模型的数据。请注意，调用my_tensor.to(device) 返回my_tensorGPU上的新副本。它不会覆盖 my_tensor。因此，请记住手动覆盖张量：。my_tensor = my_tensor.to(torch.device('cuda'))

保存在CPU，在GPU上加载

保存：

torch.save(model.state_dict(), PATH)

加载：

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

在已训练并保存在CPU上的GPU上加载模型时，请将map_location函数中的参数设置 torch.load()为 cuda：device_id。这会将模型加载到给定的GPU设备。接下来，请务必调用model.to(torch.device('cuda'))将模型的参数张量转换为CUDA张量。最后，确保.to(torch.device('cuda'))在所有模型输入上使用该函数来为CUDA优化模型准备数据。请注意，调用 my_tensor.to(device)返回my_tensorGPU上的新副本。它不会覆盖my_tensor。因此，请记住手动覆盖张量：。my_tensor = my_tensor.to(torch.device('cuda'))