PyTorch：如何保存和加载CNN图像处理模型

在使用PyTorch训练完CNN图像处理模型后，为了后续使用方便，您需要将训练好的模型保存到磁盘上。本教程将演示如何使用PyTorch提供的功能保存和加载模型。

保存模型

PyTorch 提供了 torch.save() 函数来保存模型。该函数可以保存模型的参数和其他相关信息到磁盘上，以便后续加载和使用。

以下是保存模型的示例代码：python# 假设训练好的模型为model# 假设保存路径为'checkpoint.pth'

保存模型参数torch.save(model.state_dict(), 'checkpoint.pth')

在上述代码中：

• model.state_dict() 函数返回模型的参数字典，包含了模型的所有学习参数。* torch.save() 函数将参数字典保存到指定的文件路径中（此处为 'checkpoint.pth'）。

如果您还希望保存模型的其他相关信息（如模型的结构、优化器的状态等），可以将这些信息组合为字典，一并保存到文件中。例如：python# 假设训练好的模型为model# 假设保存路径为'checkpoint.pth'# 假设优化器为optimizer

定义保存的字典checkpoint = { 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'epoch': epoch, # 其他需要保存的信息}

保存模型和相关信息torch.save(checkpoint, 'checkpoint.pth')

加载模型

保存模型后，您可以使用 torch.load() 函数加载模型的参数和其他相关信息。python# 加载模型checkpoint = torch.load('checkpoint.pth')

创建模型对象model = resnet101() # 使用与保存模型时相同的模型结构

加载模型参数model.load_state_dict(checkpoint['model_state_dict'])

加载优化器状态 (如果需要)optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

获取其他相关信息 (如果需要)epoch = checkpoint['epoch']

在加载模型时：

• 使用 torch.load() 函数加载保存的模型文件。* 创建一个与保存模型时结构相同的模型对象。* 使用 model.load_state_dict() 函数加载模型参数。* 如果保存了优化器状态等其他信息，可以使用相应的方法加载。

请确保保存路径为一个合适的位置，并根据需要指定保存的文件名和扩展名。