PyTorch保存模型用于测试和用于继续训练的区别详解

在PyTorch中，我们可以使用torch.save函数将训练好的模型保存到磁盘上，以便在以后的时间内进行测试或继续训练。但是，保存模型用于测试和用于继续训练有一些区别。本文将详细介绍这些区别，并提供两个示例说明。

保存模型用于测试

当我们将模型保存用于测试时，我们通常只需要保存模型的权重，而不需要保存优化器的状态。这是因为在测试时，我们只需要使用模型的权重来进行前向传播，而不需要进行反向传播或优化。

以下是一个保存模型用于测试的示例：

import torch
import torch.nn as nn

# 实例化模型
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 30),
    nn.ReLU(),
    nn.Linear(30, 1),
    nn.Sigmoid()
)

# 训练模型
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
criterion = nn.BCELoss()

for epoch in range(10):
    for i in range(100):
        x = torch.randn(32, 10)
        y = torch.randint(0, 2, (32, 1)).float()

        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()

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

在这个示例中，我们首先实例化了一个名为model的模型，并使用随机数据对其进行了训练。然后，我们使用torch.save函数将模型的权重保存到名为model_weights.pth的文件中。

保存模型用于继续训练

当我们将模型保存用于继续训练时，我们需要保存模型的权重和优化器的状态。这是因为在继续训练时，我们需要使用之前的优化器状态来更新模型的权重。

以下是一个保存模型用于继续训练的示例：

import torch
import torch.nn as nn

# 实例化模型
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 30),
    nn.ReLU(),
    nn.Linear(30, 1),
    nn.Sigmoid()
)

# 训练模型
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
criterion = nn.BCELoss()

for epoch in range(10):
    for i in range(100):
        x = torch.randn(32, 10)
        y = torch.randint(0, 2, (32, 1)).float()

        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()

    # 保存模型和优化器状态
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss
    }, 'model_checkpoint.pth')

在这个示例中，我们首先实例化了一个名为model的模型，并使用随机数据对其进行了训练。然后，我们使用torch.save函数将模型的权重和优化器的状态保存到名为model_checkpoint.pth的文件中。