pytorch保存加载模型的两种方式

pytorch保存模型的两种方式

开始：定义一个简单的模型

import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

# 创建模型实例
model = SimpleModel()

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

# 模拟训练过程
for epoch in range(10):
    # 生成一些随机数据
    inputs = torch.randn(1, 10)
    target = torch.randn(1, 1)

    # 前向传播
    output = model(inputs)
    loss = criterion(output, target)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    print(f'Epoch [{epoch+1}/10], Loss: {loss.item():.4f}')

方式一：保存和加载模型的状态字典

# 保存模型的状态字典
torch.save(model.state_dict(), 'model_state_dict1.pth')

# 加载模型的状态字典
model = SimpleModel()
model.load_state_dict(torch.load('model_state_dict1.pth'))
model.eval()  # 设置模型为评估模式


######################在GPU上保存的模型在cpu上使用
# 保存模型的状态字典
torch.save(model.state_dict(), 'model_state_dict1.pth')

# 加载模型的状态字典，指定 map_location='cpu'
model = SimpleModel()
model.load_state_dict(torch.load('model_state_dict1.pth', map_location=torch.device('cpu')))
model.eval()  # 设置模型为评估模式

方式二：2保存和加载整个模型

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

# 加载整个模型
model = torch.load('entire_model2.pth')
model.eval()  # 设置模型为评估模式




######################在GPU上保存的模型在cpu上使用

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

# 加载整个模型，指定 map_location='cpu'
model = torch.load('entire_model2.pth', map_location=torch.device('cpu'))
model.eval()  # 设置模型为评估模式