def save_model(model, filePath, name): # 路径需要预先建立好 with open(filePath, name) as f: dill.dump(model, f)
时间: 2023-12-14 10:38:20 浏览: 85
这是一个保存模型的函数,使用了dill库来进行序列化并将模型保存到指定的文件路径下。
参数说明:
- `model`: 要保存的模型对象
- `filePath`: 保存模型的文件路径
- `name`: 打开文件的模式,可选值为`'wb'`(写入二进制)和`'ab'`(追加二进制)
使用示例:
```python
import torch
from transformers import AutoModel
# 加载一个预训练的GPT-2模型
model = AutoModel.from_pretrained("gpt2")
# 保存模型
save_model(model, "./model/gpt2_model.pkl", "wb")
```
以上代码将预训练的GPT-2模型保存到了`./model/gpt2_model.pkl`文件中。
阅读全文