如何使用checkpoint()函数保存PyTorch模型的训练进度

在PyTorch中，可以使用checkpoint()函数保存模型的训练进度。checkpoint()函数可以保存模型的参数、优化器状态、训练状态等信息，并可以指定保存的路径和文件名。

下面是一个使用checkpoint()函数保存训练进度的例子：

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

# 假设我们有一个简单的神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 2)
    
    def forward(self, x):
        x = self.fc(x)
        return x

# 创建模型、损失函数和优化器
model = Net()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练过程
num_epochs = 10
for epoch in range(num_epochs):
    # 假设每个epoch都有一些训练数据
    input_data = torch.randn(16, 10)
    target = torch.randn(16, 2)
    
    # 前向传播
    output = model(input_data)
    loss = criterion(output, target)
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # 保存训练进度
    checkpoint = {
        'epoch': epoch + 1,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'loss': loss
    }
    
    # 保存路径和文件名
    save_path = 'checkpoint.pth'
    
    # 使用torch.save()函数保存checkpoint
    torch.save(checkpoint, save_path)

上述代码中，我们首先创建了一个简单的神经网络模型Net，并定义了损失函数criterion和优化器optimizer。然后，进入训练循环，并在每个epoch结束后保存训练进度。

在每个epoch结束后，我们创建了一个字典checkpoint，其中包含了当前epoch的索引、模型的状态字典、优化器的状态字典以及当前epoch的损失值。接着，我们指定了保存路径和文件名，并使用torch.save()函数保存了checkpoint。

通过以上操作，我们可以在训练过程中保存模型的训练进度，方便以后恢复训练或进行预测。在实际应用中，可以根据需要自定义保存的信息和文件名，以满足不同的需求。