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

Chapter5.4 Loading and saving model weights in PyTorch

article 2025/1/17 5:28:57

5 Pretraining on Unlabeled Data

5.4 Loading and saving model weights in PyTorch

训练LLM的计算成本很高，因此能够保存和加载LLM的权重至关重要。

在PyTorch中，推荐的方式是通过将torch.save函数应用于.state_dict()方法来保存模型权重，即所谓的state_dict：

torch.save(model.state_dict(),"model.pth")

我们可以将模型权重加载到新的 GPTModel 模型实例中

model = GPTModel(GPT_CONFIG_124M)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.load_state_dict(torch.load("model.pth", map_location=device, weights_only=True))
model.eval();

自适应优化器（如 AdamW）为每个模型权重存储额外的参数。AdamW 使用历史数据动态调整每个模型参数的学习率。如果没有这些参数，优化器会重置，模型可能会学习效果不佳，甚至无法正确收敛，这意味着模型将失去生成连贯文本的能力。使用 torch.save，我们可以保存模型和优化器的 state_dict 内容，如下所示
```
torch.save({
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    }, 
    "model_and_optimizer.pth"
)
```
然后，我们可以通过以下方式恢复模型和优化器状态：首先通过 torch.load 加载保存的数据，然后使用 load_state_dict 方法：
```
checkpoint = torch.load("model_and_optimizer.pth", weights_only=True)

model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])

optimizer = torch.optim.AdamW(model.parameters(), lr=0.0005, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train();
```