PyTorch 模型保存、加载与结构查看入门指南

PyTorch 保存和加载模型、查看模型结构的方法（入门级，不包括保存优化器、只加载部分参数等进阶方法）
在 PyTorch 中，模型的保存和加载是一项重要的任务，它可以帮助我们在训练过程中保存模型的状态，以便在需要时进行预测或继续训练。此外，查看模型结构可以帮助我们理解模型的组成和结构，以便更好地调整和优化模型。下面我们将详细介绍这些基本操作。
一、保存模型
在 PyTorch 中，可以使用 torch.save() 函数来保存模型。通常，我们只需要将模型对象作为参数传递给该函数即可。下面是一个简单的示例：

# 假设我们有一个已经训练好的模型 model
# 训练完成后，我们可以这样保存模型
torch.save(model.state_dict(), 'model.pth')

在这个例子中，model.state_dict() 返回一个包含模型所有参数的字典对象，'model.pth' 是保存模型的文件名。使用 state_dict() 方法可以只保存模型的参数，而不包括其他信息（如模型的类定义等）。
如果你希望保存整个模型（包括其类定义），可以直接将模型对象本身传递给 torch.save() 函数：

torch.save(model, 'model.pth')

二、加载模型
要加载模型，可以使用 PyTorch 的 torch.load() 函数。如果之前使用 state_dict() 方法保存了模型，那么加载时需要首先实例化模型的结构，然后再加载参数。下面是一个示例：

# 假设我们有一个与之前保存的模型结构相同的模型 new_model
# 首先，我们加载模型的参数
params = torch.load('model.pth')
# 然后，我们将参数加载到新的模型中
new_model.load_state_dict(params)

如果之前使用 torch.save() 方法保存了整个模型（包括其类定义），那么可以直接加载模型：

model = torch.load('model.pth')

三、查看模型结构
在 PyTorch 中，可以使用 torch.nn.Module 类的 summary() 方法来查看模型的结构。这个方法会输出每个层的名称、输入和输出的特征数量等详细信息。下面是一个示例：

print(model.summary())

以上就是 PyTorch 中保存和加载模型、查看模型结构的基本方法。需要注意的是，这些方法都是基于基本的保存和加载操作，不包括一些更高级的功能，如保存优化器、只加载部分参数等。对于这些高级功能，可以参考 PyTorch 的官方文档或其他相关教程进行学习。