优化 PyTorch 模型中的内存使用
图片来源：编辑 | Midjourney
尽管超过40%的企业表示对人工智能感到满意，但许多企业对开箱即用的解决方案并不满意，这导致了对本地人工智能解决方案及其后续使用 PyTorch 进行调整的需求。
因此，高效的内存管理在使用 PyTorch 时至关重要，尤其是在资源受限的系统上训练和部署深度学习模型时。如果没有适当的优化，大型模型会很快耗尽可用内存，导致性能瓶颈甚至直接崩溃。
为了解决优化不足的问题，我们准备了这份指南。它深入探讨了 PyTorch 中优化内存使用的策略，涵盖了在保持模型性能的同时最大化效率的关键技术。
使用混合精度训练
混合精度训练同时利用16位和32位浮点计算来减少内存消耗并加速训练。PyTorch 的 torch.cuda.amp 模块使其实现变得非常简单。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() model = MyModel().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) for inputs, labels in dataloader: optimizer.zero_grad() with autocast(): outputs = model(inputs) loss = loss_fn(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() |
这种方法不仅减少了内存使用，而且通常还能在配备张量核心（tensor cores）的现代 GPU 上加快训练速度。
利用梯度检查点
梯度检查点（Gradient checkpointing）通过增加计算开销来节省内存。通过仅存储一部分中间激活并在反向传播期间重新计算它们，可以显著减少内存使用。
使用 torch.utils.checkpoint 模块在 PyTorch 中启用梯度检查点。
|
1 2 3 4 5 6 7 8 9 10 11 |
from torch.utils.checkpoint import checkpoint class CheckpointedModel(torch.nn.Module): def __init__(self, model): super().__init__() self.model = model def forward(self, x): return checkpoint(self.model, x) model = CheckpointedModel(MyModel()) |
这项技术对于层数众多的深度模型尤其有效。
使用 torch.utils.data 优化数据加载
内存效率低下的问题通常在数据加载期间出现。为最小化这些问题，请使用 DataLoader 类并进行以下优化：
- 在 GPU 训练中使用 pin_memory：确保主机和 GPU 内存之间的数据传输更快。
- 设置 prefetch_factor 和 num_workers：调整这些参数以启用并行数据加载，减少 GPU 空闲时间。
例如
|
1 2 3 4 5 6 7 8 9 |
from torch.utils.data import DataLoader data_loader = DataLoader( dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True ) |
采用模型量化
量化（Quantization）降低了模型参数和计算的数值精度，从而在不显著降低性能的情况下减少了内存使用。PyTorch 支持静态量化（部署前应用）和动态量化（推理期间应用）。
这是一个动态量化的例子：
|
1 2 3 4 5 6 |
from torch.quantization import quantize_dynamic model = MyModel() quantized_model = quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8 ) |
这种方法非常适合内存限制严格的部署场景。
减少冗余变量
过度使用临时变量会使内存使用量膨胀。请使用原地（in-place）操作，并在不再需要变量时显式地释放它们。
|
1 2 3 4 5 |
# 原地操作 x = x.add_(y) # 释放内存 del y |
在训练期间使用 torch.cuda.memory_allocated() 和 torch.cuda.memory_reserved() 来监控您的内存使用情况，以识别潜在的泄漏。
动态调整批量大小
批量大小（batch size）直接影响内存使用。对于内存受限的环境，可以采用动态批量大小策略：
- 从一个较大的批量大小开始。
- 如果发生 CUDA 内存不足错误，则逐步减小它。
以下是批量大小调整的示例：
|
1 2 3 4 5 6 7 8 9 10 11 |
batch_size = 64 while True: try: train(model, dataloader, batch_size) break except RuntimeError as e: if "out of memory" in str(e): batch_size //= 2 else: raise e |
剪枝不必要的参数
剪枝（Pruning）从模型中消除冗余权重，从而减少其内存占用。PyTorch 提供了 torch.nn.utils.prune 模块，用于结构化和非结构化剪枝。
|
1 2 |
import torch.nn.utils.prune as prune prune.l1_unstructured(model.layer, name='weight', amount=0.4) |
剪枝后，对模型进行微调以恢复性能。
使用分布式训练与内存分割
在训练大型模型时，将计算和内存分布到多个 GPU 上。PyTorch 的 torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel 为此提供了便利。
|
1 |
model = torch.nn.DataParallel(model) |
另外，对于内存密集型模型，可以考虑使用张量并行（tensor parallelism），将网络层分割到不同 GPU 上以减少内存瓶颈。
监控和分析内存使用情况
有效的内存优化始于了解模型的内存使用情况。PyTorch 的 torch.utils.bottleneck 以及 PyTorch Profiler 和 nvidia-smi 等第三方工具提供了详细的分析。
使用 torch.profiler 的示例：
|
1 2 3 4 5 6 |
import torch.profiler as profiler with profiler.profile(on_trace_ready=profiler.tensorboard_trace_handler('./log')) as prof: model(inputs) print(prof.key_averages().table()) |
诸如此类的分析工具可以帮助您精确定位内存瓶颈，从而实现有针对性的优化。例如，您可能会发现某些层或操作消耗了不成比例的内存，这会指导您去优化或替换它们。
采用高效的部署策略
对于部署，可以使用 TorchScript 或 ONNX 导出来优化模型序列化。这些格式在保持兼容性的同时减少了内存需求。
|
1 2 |
traced_model = torch.jit.trace(model, example_input) torch.jit.save(traced_model, "optimized_model.pt") |
利用 ONNX 与其他框架或推理引擎兼容，从而实现更广泛的部署选项。此类优化不仅节省内存，还能提高推理速度，使其对生产环境至关重要。
此外，可以考虑实施权重共享或张量压缩技术，以在部署期间进一步节省内存。像 LiteRT（以前称为 TensorFlow Lite）这样的压缩库或自定义的 PyTorch 压缩脚本有助于在不显著影响性能的情况下减小模型的整体大小。
对于大规模生产系统，利用基于云的 GPU 服务器托管解决方案可确保可扩展性和一致的性能，而不会受到本地硬件的限制。这有助于同时部署多个模型并减少运营开销。
更高效内存优化的技巧
使用稀疏矩阵
如果您的模型或数据集中包含大量零值，使用稀疏矩阵可以显著减少内存使用。PyTorch 提供了 torch.sparse 模块用于稀疏张量的操作。
|
1 |
sparse_tensor = torch.sparse_coo_tensor(indices, values, size) |
稀疏张量在自然语言处理等场景中特别有用，因为嵌入（embeddings）通常包含许多零值。
应用知识蒸馏
知识蒸馏（Knowledge distillation）涉及训练一个更小、更节省内存的模型来模仿一个更大模型的性能。这种方法可以在保持相似准确率水平的同时显著减少内存使用。
|
1 2 3 4 5 6 7 8 9 10 11 |
teacher_model = MyLargeModel() student_model = MySmallModel() # 训练学生模型以模仿教师模型的预测 for inputs in dataloader: teacher_outputs = teacher_model(inputs) student_outputs = student_model(inputs) loss = loss_fn(student_outputs, teacher_outputs) optimizer.zero_grad() loss.backward() optimizer.step() |
动态卸载层
在处理超大型模型时，可以考虑将某些层或计算卸载到 CPU 内存甚至磁盘存储中。像 Hugging Face 的 Accelerate 这样的库通过根据内存限制动态管理卸载，使这个过程变得无缝。
结论
在 PyTorch 中进行内存优化是一个多方面的过程，涉及训练、部署和基础设施层面的调整。混合精度训练、梯度检查点、量化和剪枝等技术能显著减少内存消耗。
将这些策略与有效的分析工具、稀疏矩阵操作和知识蒸馏相结合，可以确保在不牺牲性能的情况下高效利用资源。通过利用云资源和先进的部署策略，开发人员即使在资源最受限的环境中也能最大限度地发挥 PyTorch 的能力。
暂无评论。