大模型原理：保存、加载模型和优化器的权重-马育民老师

# 介绍

虽然使用的大语言模型和数据集都相对较小，但训练时间仍然较长，如果数据量大，需要使用价格高昂的GPU，所以训练大语言模型代价非常高昂。
因此，保存大语言模型的参数非常重要，这样就不必每次使用它时都重新运行训练。

本文实现对模型进行保存，通常能为以后的继续训练带来方便（第(6)步）

[![](https://www.malaoshi.top/upload/0/0/1GW2gazfDTAG.png)](https://www.malaoshi.top/upload/0/0/1GW2gazfDTAG.png)

# 保存模型中的参数

推荐使用 `torch.save()` 函数保存模型的 `state_dict`，即将每个层映射到其参数的字典

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

`model.pth` 是保存 `state_dict` 的文件名。`.pth` 扩展名是PyTorch文件的规范，尽管从技术上讲可以使用任何文件扩展名。

# 加载模型中的参数

在通过 `state_dict` 保存模型权重之后，可以将模型权重加载到一个新的 `GPTModel` 模型实例中：

```
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(torch.load("model.pth", map_location=device))
model.eval()
```

`dropout` 通过在训练过程中 **随机 “丢弃”**一层的神经元，有助于防止模型对训练数据过拟合。
但是，在推断过程中，要使用网络学习到的所有信息。因此，使用 `model.eval()` 将模型 **切换到推断模式**，这样就会 **禁用模型的 `dropout` 层**。

如果继续训练模型，可以使用前面定义的 `train_model_simple()` 函数，建议同时保存优化器状态。

# 保存优化器中的参数

像AdamW这样的自适应优化器可以为每个模型权重存储额外的参数。AdamW可以使用历史数据动态地调整每个模型参数的学习率。如果没有它，那么优化器就会重置，模型可能学习效果不佳，甚至无法正确收敛，这意味着模型将失去生成连贯文本的能力。可以使用 `torch.save()` 保存模型和优化器的 `state_dict` 内容：

```
torch.save({
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    },
    "model_and_optimizer.pth"
)
```

# 加载优化器中的参数

```
checkpoint = torch.load("model_and_optimizer.pth", map_location=device)
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train();
```

原文出处：http://www.malaoshi.top/show_1GW2gfmlBW0J.html