如何保存 PyTorch 模型？

2025 年 3 月 28 日 | 阅读 6 分钟

什么是 PyTorch？

PyTorch 是 Facebook 人工智能实验室 (FAIR) 开发的开源深度学习框架。它通常因其灵活性、易用性以及处理快速原型设计的强大功能而得到广泛使用。 PyTorch 提供了强大的计算图，这意味着该图是动态构建的，使其易于更改和调试。这与静态计算图有所不同，在静态计算图中，图一旦定义就无法修改，除非重新编译。

PyTorch 的主要特点包括

动态计算图： 考虑更自然、更灵活的模型结构。
易于学习和使用： PyTorch 的语法很基础，而且与标准结构非常相似，因此对新手来说是开放的。
CUDA 支持： PyTorch 可以通过 NVIDIA 的 CUDA 库利用 GPU 来加速注册。
全面的库： PyTorch 包含用于构建和准备神经网络的丰富库和工具的安排。
社区和生态系统： 强大的社区支持和广泛的环境，包括 torchvision 和 torchaudio 等库。

在 PyTorch 中保存和加载模型的逐步指南

在 PyTorch 中保存和加载模型是任何从事深度学习工作的人的基本技能。本指南将涵盖保存和加载 PyTorch 模型的各种策略，包括保存状态词典、保存整个模型以及处理这些任务的最佳实践。

第 1 步：安装 PyTorch

在开始之前，确保您已安装 PyTorch。您可以使用 pip 安装它

第 2 步：定义您的模型

首先，定义一个神经网络模型。这是一个使用完全关联网络的简单模型

import torch
import torch.nn as nn
import torch.optim as optim
# Define a simple neural network
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)
# Instantiate the model and define a loss function and optimizer
model = SimpleModel()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

第 3 步：训练您的模型

为演示目的使用假数据训练模型

# Dummy training loop
for epoch in range(100):
    # Generate dummy data
    inputs = torch.randn(10, 10)
    targets = torch.randn(10, 1)
    # Forward pass
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    # Backward pass and optimization
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
print("Model training complete.")

输出

Model training complete

保存和加载模型

在 PyTorch 中保存和加载模型有两种基本策略

仅保存模型边界（state_dict）。
保存整个模型，包括结构和边界。

方法 1：保存和加载模型参数 (state_dict)

由于其灵活性和效率，保存状态词典是最推荐的策略。

保存 state_dict

# Save the state_dict
torch.save(model.state_dict(), 'model_state_dict.pth')
print("Model state_dict saved.")

输出

Model state_dict saved

加载 state_dict

# Recreate the model architecture
model = SimpleModel()
# Load the state_dict
model.load_state_dict(torch.load('model_state_dict.pth'))
print("Model state_dict loaded.")
# Set the model to evaluation mode
model.eval()

输出

Model state_dict loaded

方法 2：保存和加载整个模型

保存整个模型包含结构和边界。对于一些用例，此技术不太灵活，但更直接。

保存整个模型

# Save the entire model
torch.save(model, 'entire_model.pth')
print("Entire model saved.")

输出

Entire model saved

加载整个模型

# Load the entire model
model = torch.load('entire_model.pth')
print("Entire model loaded.")
# Set the model to evaluation mode
model.eval()

输出

Entire model loaded

保存和加载模型的最佳实践

使用 state_dict 以获得灵活性和效率

通常，由于其灵活性和效率，使用 state_dict 是最佳实践。它允许您在不同的内容中重新生成模型设计，并加载边界，从而更直接地管理更改和更新。

如果您计划恢复训练，请保存优化器

假设您打算继续准备，请始终将优化器的 state_dict 与模型的 state_dict 一起保存。这可以确保增强器的边界（如学习率和力）得到正确恢复。

保存模型和优化器 state_dict

# Save the model and optimizer state_dict
torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
}, 'model_and_optimizer.pth')
print("Model and optimizer state_dict saved.")

输出

Model and optimizer state_dict saved

加载模型和优化器 state_dict

# Recreate the model and optimizer
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# Load the checkpoint
checkpoint = torch.load('model_and_optimizer.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
print("Model and optimizer state_dict loaded.")
# Set the model to evaluation mode
model.eval()

输出

Model and optimizer state_dict loaded

小心处理设备转换

在保存和加载模型时，请注意所用的设备（computer CPU 或 GPU）。这是处理设备转换的方式

在 GPU 上保存，在 CPU 上加载

# Save on GPU
torch.save(model.state_dict(), 'model_gpu.pth')
# Load on CPU
device = torch.device('cpu')
model.load_state_dict(torch.load('model_gpu.pth', map_location=device))
model.eval()
print("Model loaded on CPU.")

输出：

Model loaded on CPU

在 GPU 上保存，在 GPU 上加载

# Save on GPU
torch.save(model.state_dict(), 'model_gpu.pth')
# Load on GPU
device = torch.device('cuda')
model.load_state_dict(torch.load('model_gpu.pth'))
model.to(device)
model.eval()
print("Model loaded on GPU.")

输出

Model loaded on GPU

在 CPU 上保存，在 GPU 上加载

# Save on CPU
torch.save(model.state_dict(), 'model_cpu.pth')
# Load on GPU
device = torch.device('cuda')
model.load_state_dict(torch.load('model_cpu.pth'))
model.to(device)
model.eval()
print("Model loaded on GPU.")

输出

Model loaded on GPU

使用 eval() 进行推理

在执行归纳之前，请持续将模型设置为评估模式，使用 model.eval()。这可以确保 dropout 和分组标准化等层能够正常运行。

训练期间的检查点

对于长时间的准备过程，定期保存指定点是个好主意。通过这种方式，您可以在出现中断时从最后的指定点继续准备。

定期保存检查点

# Dummy training loop
for epoch in range(100):
    # Generate dummy data
    inputs = torch.randn(10, 10)
    targets = torch.randn(10, 1)    
    # Forward pass
    outputs = model(inputs)
    loss = criterion(outputs, targets)    
    # Backward pass and optimization
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # Periodically save checkpoints
    if epoch % 10 == 0:
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
        }, 'checkpoint.pth')
        print(f'Checkpoint saved at epoch {epoch}.')
print("Model training complete.")

输出

Checkpoint saved at epoch 0.
Checkpoint saved at epoch 10.
Checkpoint saved at epoch 20.
Checkpoint saved at epoch 30.
Checkpoint saved at epoch 40.
Checkpoint saved at epoch 50.
Checkpoint saved at epoch 60.
Checkpoint saved at epoch 70.
Checkpoint saved at epoch 80.
Checkpoint saved at epoch 90.
Model training complete.

从检查点恢复训练

# Load the checkpoint
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1
loss = checkpoint['loss']
print(f"Resuming training from epoch {start_epoch} with loss {loss.item()}.")
# Resume training
for epoch in range(start_epoch, 100):
    # Generate dummy data
    inputs = torch.randn(10, 10)
    targets = torch.randn(10, 1)
    # Forward pass
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    # Save the outputs for monitoring
    torch.save(outputs, f'outputs_epoch_{epoch}.pth')
    # Backward pass and optimization
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # Periodically save checkpoints
    if epoch % 10 == 0:
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
        }, 'checkpoint.pth')
        print(f'Checkpoint saved at epoch {epoch}.')
print("Model training complete.")

输出

Resuming training from epoch 91 with loss [loss value].
Checkpoint saved at epoch 100.
Model training complete.

高级用法

使用自定义项目保存和加载

如果您的模型包含自定义项目（例如自定义层），则在加载模型时，您需要保证这些项目已准确定义。

定义自定义模型

class CustomLayer(nn.Module):
    def __init__(self):
        super(CustomLayer, self).__init__()
        self.fc = nn.Linear(10, 1)
    def forward(self, x):
        return self.fc(x)
class CustomModel(nn.Module):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.custom_layer = CustomLayer()
    def forward(self, x):
        return self.custom_layer(x)
# Save the state_dict
model = CustomModel()
torch.save(model.state_dict(), 'custom_model_state_dict.pth')

加载自定义模型

# Recreate the model architecture
model = CustomModel()
# Load the state_dict
model.load_state_dict(torch.load('custom_model_state_dict.pth'))
print("Custom model state_dict loaded.")
# Set the model to evaluation mode
model.eval()

输出

Custom model state_dict loaded.

结论

总之，掌握在 PyTorch 中保存和加载模型的技巧对于有效的深度学习项目管理至关重要。本指南重点介绍了基本实践，例如由于其灵活性和生产力而使用 state_dict 进行边界功能，以及同时保存分析器以实现连续准备恢复。简单地保存和加载，但对于架构更改或迁移学习方案，灵活性较低。小心处理设备转换并保证模型在推断过程中设置为评估模式对于避免错误至关重要。执行定期的检查点改进工作流程的多功能性，从而能够持续进行准备和观察进度。这些实践与高级技术（如管理自定义项目和众多模型）相结合，使专家能够保持强大且可重现的 PyTorch 深度学习工作流程。通过结合这些方法，开发人员可以提高效率、协作能力及其人工智能应用的可靠性。

下一主题Torch-cuda-in-pytorch

← 上一主题下一主题 →