当前位置：首页 > news >正文

pytorch框架保存和加载模型

news 2025/8/6 5:43:56

在 PyTorch 中，有几种常见的方法来保存和加载模型，具体方法取决于你想保存什么内容（例如，只保存模型的权重，还是保存整个模型）。下面我将介绍几种常见的保存和加载模型的方法。

1、保存和加载模型的 `state_dict`

state_dict 是一个简单的 Python 字典对象，它将每一层的参数映射到其对应的张量（tensor）。这是最推荐的保存方式，因为它只保存模型的权重和参数，而不依赖于模型的代码结构。
保存模型的 state_dict：

import torch
import torch.nn as nn# 假设我们有一个模型
model = nn.Linear(10, 2)# 保存模型的 state_dict
torch.save(model.state_dict(), 'model.pth')

加载模型的 state_dict：

import torch
import torch.nn as nn# 假设我们有相同结构的模型
model = nn.Linear(10, 2)# 加载 state_dict
model.load_state_dict(torch.load('model.pth'))# 使用模型
model.eval()

2、保存和加载整个模型

这种方法不仅保存了模型的权重和参数，还保存了模型的结构。因此，整个模型在加载时不需要再定义模型结构。尽管这提供了更大的便利性，但也有一些缺点，比如文件可能与代码过于耦合（会保存模型的类和结构，如果代码发生变化可能无法加载）。
保存整个模型：

import torch
import torch.nn as nn# 假设我们有一个模型
model = nn.Linear(10, 2)# 保存整个模型
torch.save(model, 'model_entire.pth')

加载整个模型：

import torch# 加载整个模型
model = torch.load('model_entire.pth')# 使用模型
model.eval()

3、保存和加载 checkpoint（保存更多的训练状态）

如果你希望保存更多的训练信息，例如优化器状态、当前的 epoch 以及训练损失等，以便后续继续训练，可以保存一个包含多个内容的字典。
保存 checkpoint：

import torch
import torch.nn as nn
import torch.optim as optim# 假设我们有一个模型和优化器
model = nn.Linear(10, 2)
optimizer = optim.SGD(model.parameters(), lr=0.001)# 保存 checkpoint
checkpoint = {'epoch': 10,'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'loss': 0.5,
}torch.save(checkpoint, 'checkpoint.pth')

加载 checkpoint：

import torch
import torch.nn as nn
import torch.optim as optim# 假设我们有相同结构的模型和优化器
model = nn.Linear(10, 2)
optimizer = optim.SGD(model.parameters(), lr=0.001)# 加载 checkpoint
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']# 恢复训练
model.train()

4、动态量化后的模型保存与加载

如果你对模型进行了量化（如动态量化），你可能需要保存量化后的模型。保存和加载方法与保存 state_dict 类似。
动态量化模型：

import torch
import torch.nn as nn# 假设我们有一个模型
model = nn.Linear(10, 2)# 进行动态量化
quantized_model = torch.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8
)# 保存量化后的模型
torch.save(quantized_model.state_dict(), 'quantized_model.pth')

加载量化模型：

import torch
import torch.nn as nn# 重新定义与量化前相同的模型结构
model = nn.Linear(10, 2)# 加载量化后的模型权重
model.load_state_dict(torch.load('quantized_model.pth'))
model.eval()