Pytorch怎样保存训练好的模型

在PyTorch中，我们可以使用torch.save()函数来保存训练好的模型。该函数接受两个参数：要保存的模型和保存模型的文件路径。我们还可以使用torch.load()函数来加载已保存的模型。

以下是详细的攻略：

保存模型

我们可以使用torch.save()函数来保存训练好的模型。该函数接受两个参数：要保存的模型和保存模型的文件路径。我们可以使用.pt或.pth作为文件扩展名。

例如，我们可以使用以下代码来保存训练好的模型：

import torch

# 定义模型
model = torch.nn.Linear(10, 1)

# 训练模型...

# 保存模型
torch.save(model.state_dict(), 'model.pt')

在这个示例中，我们首先定义了一个包含10个输入和1个输出的线性模型。然后，我们训练了模型，并使用torch.save()函数来保存训练好的模型。我们将模型的状态字典作为第一个参数传递给torch.save()函数，并将文件路径作为第二个参数传递给该函数。

加载模型

我们可以使用torch.load()函数来加载已保存的模型。该函数接受一个文件路径作为参数，并返回一个包含模型状态字典的字典对象。我们可以使用model.load_state_dict()函数来将状态字典加载到模型中。

例如，我们可以使用以下代码来加载已保存的模型：

import torch

# 定义模型
model = torch.nn.Linear(10, 1)

# 加载模型
model.load_state_dict(torch.load('model.pt'))

在这个示例中，我们首先定义了一个包含10个输入和1个输出的线性模型。然后，我们使用torch.load()函数来加载已保存的模型，并使用model.load_state_dict()函数将状态字典加载到模型中。

示例1：保存和加载线性回归模型

在这个示例中，我们将使用Tensor来实现一个简单的线性回归模型。我们将使用随机生成的数据来训练模型，并使用梯度下降算法来更新模型参数。在训练完成后，我们将使用torch.save()函数来保存训练好的模型，并使用torch.load()函数来加载已保存的模型。

import torch

# 生成随机数据
x = torch.randn(1000, 1000)
y = torch.randn(1000, 1)

# 定义模型
model = torch.nn.Linear(1000, 1)

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

# 进行模型训练
for epoch in range(1000):
    # 前向传
    y_pred = model(x)

    # 计算损失
    loss = criterion(y_pred, y)

    # 反向传播
    loss.backward()

    # 更新模型参数
    optimizer.step()
    optimizer.zero_grad()

    # 打印训练日志
    if (epoch + 1) % 100 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, 1000, loss.item()))

# 保存模型
torch.save(model.state_dict(), 'model.pt')

# 加载模型
model.load_state_dict(torch.load('model.pt'))

在这个示例中，我们首先生成了1000个随机数据点，并使用这数据点来训练模型。然后，我们定义了一个包含1000个权重和一个偏置的线性模型。我们使用均方误差损失函数和随机梯度下降优化器来训练模型。最后，我们使用torch.save()函数来保存训练好的模型，并使用torch.load()函数来加载已保存的模型。

示例2：保存和加载卷积神经网络模型

在这个示例中，我们将使用Tensor来实现一个简单的卷积神经网络模型。我们将使用MNIST数据集来训练模型，并使用Adam优化器来更新模型参数。在训练完成后，我们将使用torch.save()函数来保存训练好的模型，并使用torch.load()函数来加载已保存的模型。

import torch
import torchvision
import torchvision.transforms as transforms

# 定义数据预处理
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))])

# 加载MNIST数据集
trainset = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64,
                                          shuffle=True, num_workers=2)

# 定义卷积神经网络模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 6, 5)
        self.pool = torch.nn.MaxPool2d(2, 2)
        self.conv2 = torch.nn.Conv2d(6, 16, 5)
        self.fc1 = torch.nn.Linear(16 * 4 * 4, 120)
        self.fc2 = torch.nn.Linear(120, 84)
        self.fc3 = torch.nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(torch.nn.functional.relu(self.conv1(x)))
        x = self.pool(torch.nn.functional.relu(self.conv2(x)))
        x = x.view(-1, 16 * 4 * 4)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)

# 进行模型训练
for epoch in range(5):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data

        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 100 == 99:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0

# 保存模型
torch.save(net.state_dict(), 'model.pt')

# 加载模型
net.load_state_dict(torch.load('model.pt'))

在这个示例中，我们首先加载了MNIST数据集，并定义了一个包含两个卷积层和三个全连接层的卷积神经网络模型。我们使用交叉熵损失函数和Adam优化器来训练模型。最后，我们使用torch.save()函数来保存训练好的模型，并使用torch.load()函数来加载已保存的模型。