pytorch如何从中断点继续训练-CFANZ编程社区

PyTorch 如何从中断点继续训练

在深度学习模型的训练过程中，由于各种原因（如系统重启、断电等），训练过程可能会中断。因此，能够从中断点继续训练是一个非常重要的功能。本文将详细介绍如何在PyTorch中实现这一功能，并展示具体的代码示例。

1. 方案概述

在PyTorch中实现从中断点恢复训练的基本步骤如下：

保存模型和优化器状态：在每个训练周期（epoch）或特定间隔保存模型的权重和优化器的状态。
加载模型和优化器状态：在恢复训练时，从最近一次保存的状态开始训练。
管理 epoch：记录当前 epoch，以确保从正确的地方继续。

2. 具体实现步骤

下面，我们将通过一个简单的示例实现从中断点继续训练的功能。

2.1 保存模型和优化器状态

在每个 epoch 结束后，可以使用 torch.save() 函数将模型状态和优化器状态保存到文件中。

import torch

def save_checkpoint(model, optimizer, epoch, loss, filename='checkpoint.pth'):
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
    }
    torch.save(checkpoint, filename)
    print(f'Checkpoint saved at {filename}')

2.2 加载模型和优化器状态

在恢复训练时，使用 torch.load() 函数加载先前保存的状态。

def load_checkpoint(model, optimizer, filename='checkpoint.pth'):
    checkpoint = torch.load(filename)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    loss = checkpoint['loss']
    print(f'Checkpoint loaded from {filename}')
    return epoch, loss

2.3 示例代码

下面的代码展示了一个完整的训练过程，包括保存和加载检查点的逻辑。

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)

# 训练过程
def train(model, optimizer, epochs, load_existing_checkpoint=False):
    start_epoch = 0
    loss_fn = nn.MSELoss()
    
    # 如果需要从中断点继续训练，尝试加载检查点
    if load_existing_checkpoint:
        start_epoch, _ = load_checkpoint(model, optimizer)

    for epoch in range(start_epoch, epochs):
        # 假设有输入数据和标签
        inputs = torch.randn(32, 10)
        labels = torch.randn(32, 1)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()

        # 保存检查点
        save_checkpoint(model, optimizer, epoch + 1, loss.item())

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

# 实例化模型与优化器
model = SimpleModel()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 启动训练
train(model, optimizer, epochs=10, load_existing_checkpoint=False)

3. 关系图

在上述代码实现中，模型状态与优化器状态的关系可以用以下 ER 图表示：

erDiagram
    MODEL {
        int id
        string name
        string architecture
    }
    OPTIMIZER {
        int id
        string name
        float learning_rate
    }
    CHECKPOINT {
        int id
        string filename
        int epoch
        float loss
    }

    MODEL ||--o| CHECKPOINT : saves
    OPTIMIZER ||--o| CHECKPOINT : saves

4. 类图

我们也可以用以下类图来表示模型和优化器之间的关系：

classDiagram
    class SimpleModel {
        +forward(inputs)
    }
    class Checkpoint {
        +save(model, optimizer, epoch, loss)
        +load(model, optimizer)
    }
    class Trainer {
        +train(model, optimizer, epochs)
    }

    SimpleModel --> Checkpoint : saves
    Checkpoint --> SimpleModel : loads
    Trainer --> SimpleModel : uses
    Trainer --> Checkpoint : manages