当前位置：首页 > news >正文

pytorch checkpointing

news 来源：原创 2025/5/6 8:13:11

是一种在训练深度神经网络时通过增加计算代价来换取显存优化的技术。它的核心思想是：在反向传播过程中动态重新计算中间激活值（activations），而不是保存所有中间结果。这对于显存受限的场景（如训练大型模型）非常有用。

直接上代码：

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint# 1. 定义一个简单的 FFN 模型
class SimpleFFN(nn.Module):def __init__(self, input_dim=128, hidden_dim=256, output_dim=10):super().__init__()self.linear1 = nn.Linear(input_dim, hidden_dim)self.linear2 = nn.Linear(hidden_dim, hidden_dim)self.linear3 = nn.Linear(hidden_dim, output_dim)self.relu = nn.ReLU()def forward(self, x):# 2. 定义一个自定义的前向传播函数（用于 checkpoint）def custom_forward(x):x = self.linear1(x)x = self.relu(x)x = self.linear2(x)x = self.relu(x)x = self.linear3(x)return x# 3. 使用 checkpoint 包装前向传播return checkpoint(custom_forward, x)# 4. 初始化模型和数据
model = SimpleFFN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()# 模拟输入数据
input_data = torch.randn(64, 128)  # batch_size=64, input_dim=128
target = torch.randn(64, 10)       # 模拟目标输出# 5. 前向传播、损失计算和反向传播
output = model(input_data)
loss = criterion(output, target)
loss.backward()
optimizer.step()