mirror of https://github.com/hpcaitech/ColossalAI
62 lines
1.6 KiB
Markdown
62 lines
1.6 KiB
Markdown
# 模型检查点
|
||
|
||
作者 : Guangyang Lu
|
||
|
||
**预备知识:**
|
||
- [Launch Colossal-AI](./launch_colossalai.md)
|
||
- [Initialize Colossal-AI](./initialize_features.md)
|
||
|
||
**示例代码:**
|
||
- [ColossalAI-Examples Model Checkpoint](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/utils/checkpoint)
|
||
|
||
**函数是经验函数.**
|
||
|
||
## 简介
|
||
|
||
本教程将介绍如何保存和加载模型检查点。
|
||
|
||
为了充分利用Colossal-AI的强大并行策略，我们需要修改模型和张量，可以直接使用 `torch.save` 或者 `torch.load` 保存或加载模型检查点。在Colossal-AI中，我们提供了应用程序接口实现上述同样的效果。
|
||
|
||
但是，在加载时，你不需要使用与存储相同的保存策略。
|
||
|
||
## 使用方法
|
||
|
||
### 保存
|
||
|
||
有两种方法可以使用Colossal-AI训练模型，即使用engine或使用trainer。
|
||
**注意我们只保存 `state_dict`.** 因此，在加载检查点时，需要首先定义模型。
|
||
|
||
#### 同 engine 保存
|
||
|
||
```python
|
||
from colossalai.utils import save_checkpoint
|
||
model = ...
|
||
engine, _, _, _ = colossalai.initialize(model=model, ...)
|
||
for epoch in range(num_epochs):
|
||
... # do some training
|
||
save_checkpoint('xxx.pt', epoch, model)
|
||
```
|
||
|
||
#### 用 trainer 保存
|
||
```python
|
||
from colossalai.trainer import Trainer, hooks
|
||
model = ...
|
||
engine, _, _, _ = colossalai.initialize(model=model, ...)
|
||
trainer = Trainer(engine, ...)
|
||
hook_list = [
|
||
hooks.SaveCheckpointHook(1, 'xxx.pt', model)
|
||
...]
|
||
|
||
trainer.fit(...
|
||
hook=hook_list)
|
||
```
|
||
|
||
### 加载
|
||
|
||
```python
|
||
from colossalai.utils import load_checkpoint
|
||
model = ...
|
||
load_checkpoint('xxx.pt', model)
|
||
... # train or test
|
||
```
|