model checkpoint代码

以下是使用 PyTorch 实现的模型检查点代码示例：

import torch

# 定义模型
model = ...

# 定义优化器
optimizer = ...

# 定义损失函数
criterion = ...

# 定义训练循环
def train_loop():
    for epoch in range(num_epochs):
        # 训练模型
        for batch in train_loader:
            # 前向传播
            output = model(batch)

            # 计算损失
            loss = criterion(output, batch.target)

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

        # 保存模型检查点
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss
        }
        torch.save(checkpoint, 'checkpoint.pth')

# 加载模型检查点
def load_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']
    return model, optimizer, epoch, loss

在训练循环中，我们可以通过 `torch.save()` 方法保存模型的状态字典、优化器的状态字典以及其他必要的信息。然后，我们可以使用 `torch.load()` 方法加载检查点，并将其应用于模型和优化器。