526互联

：）关于transformers模型的保存与加载-|

发布时间 2023-03-22 22:03:17作者: lexn

关于transformers模型的保存与加载

两种情况， 自定义模型训练后保存， transformers预训练模型保存。

参考代码

# -*- coding: utf-8 -*-
import torch
from transformers import GPT2LMHeadModel
from transformers import AutoTokenizer
from transformers import TextGenerationPipeline

model = GPT2LMHeadModel.from_pretrained("SkyWork/SkyTextTiny")
# torch.save(model.state_dict(), 'model.bin')
model.save_pretrained("GPTMODEL")
tokenizer = AutoTokenizer.from_pretrained("SkyWork/SkyTextTiny", trust_remote_code=True)
tokenizer.save_pretrained('tokenizer_directory')
text_generator = TextGenerationPipeline(model, tokenizer, device=0)
input_str = "今天是个好天气"
max_new_tokens = 20
print(text_generator(input_str, max_new_tokens=max_new_tokens, do_sample=True))

　　

torch架构把模型分为两部分，身体和大脑， 身体即模型的结构，大脑即权重。

模型结构可能存在config.json中， 权重存在pth文件或bin文件。

1 模型训练后保存及加载

1.1保存

 1)保存结构

　　------- 待解决存在问题 Object of type GPT2Config is not JSON serializable

　　

 1 # -*- coding: utf-8 -*-
 2 import torch
 3 from transformers import GPT2LMHeadModel
 4 from transformers import AutoTokenizer
 5 from transformers import TextGenerationPipeline
 6 
 7 def save_struc(model):
 8     import json
 9     with open('config.json', 'w') as f:
10         json.dump(model.config, f)
11     torch.save(model.state_dict(), 'model.bin')
12 
13 model = GPT2LMHeadModel.from_pretrained('GPTMODEL')
14 
15 # 保存模型结构
16 save_struc(model)
View Code

2）保存权重

model = GPT2LMHeadModel.from_pretrained("SkyWork/SkyTextTiny")
torch.save(model.state_dict(), 'model.bin')

　　其中

model.state_dict()

把模型得参数权重导出到字典。

1.2 加载

model.load_state_dict(torch.load('model.bin'))

　　

先load 成为字典，再load_state_dict 加载入模型，注意这个model已经是 实例化后的即已经有了结构得model。所以只要读入权重即可。

 

2 预训练模型保存及加载

2.1

保存直接config.json和权重bin都保存再一个目录

model = GPT2LMHeadModel.from_pretrained("SkyWork/SkyTextTiny")
model.save_pretrained("GPTMODEL")

　　

目录为

 

 

2.2 加载

model = GPT2LMHeadModel.from_pretrained('GPTMODEL')