Debian PyTorch如何进行模型保存与加载

在Debian系统上使用PyTorch进行模型的保存与加载，可以按照以下步骤操作：

模型保存

1. 训练模型： 在训练过程中，你会得到一个训练好的模型对象。

2. 保存模型： 使用torch.save()函数来保存整个模型或者模型的状态字典。

   import torch
   import torchvision.models as models
   
   # 假设我们使用的是ResNet18模型
   model = models.resnet18(pretrained=True)
   
   # 训练模型...
   
   # 保存整个模型
   torch.save(model, 'model.pth')
   
   # 或者只保存模型的状态字典
   torch.save(model.state_dict(), 'model_state_dict.pth')
   

模型加载

1. 加载模型： 使用torch.load()函数来加载模型。如果之前保存的是整个模型，则直接加载；如果是状态字典，则需要先创建一个模型实例，然后加载状态字典。

   # 加载整个模型
   model = torch.load('model.pth')
   
   # 或者加载模型的状态字典
   model = models.resnet18(pretrained=False)  # 创建一个新的模型实例
   model.load_state_dict(torch.load('model_state_dict.pth'))
   

2. 将模型移动到GPU（如果可用）： 如果你的系统有GPU，并且你想在GPU上运行模型，可以使用to()方法。

   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   model.to(device)
   

注意事项

• 版本兼容性：确保保存和加载模型时使用的PyTorch版本相同或兼容。
• 设备兼容性：如果你在不同的设备上保存和加载模型（例如，从CPU到GPU），可能需要调整模型的设备。
• 自定义模型：如果你使用的是自定义模型，确保在加载状态字典之前已经定义了相同的模型结构。

示例代码

以下是一个完整的示例，展示了如何在Debian系统上保存和加载一个ResNet18模型：

import torch
import torchvision.models as models

# 定义并训练模型（这里省略了训练过程）
model = models.resnet18(pretrained=True)

# 保存整个模型
torch.save(model, 'resnet18_model.pth')

# 加载整个模型
loaded_model = torch.load('resnet18_model.pth')

# 或者加载模型的状态字典
model = models.resnet18(pretrained=False)
model.load_state_dict(torch.load('resnet18_model.pth'))

# 将模型移动到GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

通过以上步骤，你可以在Debian系统上轻松地进行PyTorch模型的保存与加载。