Transformers 保存並加載模型 | 八

做者|huggingface
編譯|VK
來源|Github

本節說明如何保存和從新加載微調模型(BERT，GPT，GPT-2和Transformer-XL)。你須要保存三種文件類型才能從新加載通過微調的模型：

• 模型自己應該是PyTorch序列化保存的模型(https://pytorch.org/docs/stab...
• 模型的配置文件是保存爲JSON文件
• 詞彙表(以及基於GPT和GPT-2合併的BPE的模型)。

這些文件的默認文件名以下：

• 模型權重文件：pytorch_model.bin
• 配置文件：config.json
• 詞彙文件：vocab.txt表明BERT和Transformer-XL，vocab.json表明GPT/GPT-2(BPE詞彙)，
• 表明GPT/GPT-2(BPE詞彙)額外的合併文件：merges.txt。

若是使用這些默認文件名保存模型，則能夠使用from_pretrained()方法從新加載模型和tokenizer。

這是保存模型，配置和配置文件的推薦方法。詞彙到output_dir目錄，而後從新加載模型和tokenizer：

from transformers import WEIGHTS_NAME, CONFIG_NAME

output_dir = "./models/"

# 步驟1：保存一個通過微調的模型、配置和詞彙表

#若是咱們有一個分佈式模型，只保存封裝的模型
#它包裝在PyTorch DistributedDataParallel或DataParallel中
model_to_save = model.module if hasattr(model, 'module') else model
#若是使用預約義的名稱保存，則能夠使用`from_pretrained`加載
output_model_file = os.path.join(output_dir, WEIGHTS_NAME)
output_config_file = os.path.join(output_dir, CONFIG_NAME)

torch.save(model_to_save.state_dict(), output_model_file)
model_to_save.config.to_json_file(output_config_file)
tokenizer.save_vocabulary(output_dir)

# 步驟2: 從新加載保存的模型

#Bert模型示例
model = BertForQuestionAnswering.from_pretrained(output_dir)
tokenizer = BertTokenizer.from_pretrained(output_dir, do_lower_case=args.do_lower_case)  # Add specific options if needed
#GPT模型示例
model = OpenAIGPTDoubleHeadsModel.from_pretrained(output_dir)
tokenizer = OpenAIGPTTokenizer.from_pretrained(output_dir)

若是要爲每種類型的文件使用特定路徑，則能夠使用另外一種方法保存和從新加載模型：

output_model_file = "./models/my_own_model_file.bin"
output_config_file = "./models/my_own_config_file.bin"
output_vocab_file = "./models/my_own_vocab_file.bin"

# 步驟1：保存一個通過微調的模型、配置和詞彙表

#若是咱們有一個分佈式模型，只保存封裝的模型
#它包裝在PyTorch DistributedDataParallel或DataParallel中
model_to_save = model.module if hasattr(model, 'module') else model

torch.save(model_to_save.state_dict(), output_model_file)
model_to_save.config.to_json_file(output_config_file)
tokenizer.save_vocabulary(output_vocab_file)

# 步驟2: 從新加載保存的模型

# 咱們沒有使用預約義權重名稱、配置名稱進行保存，沒法使用`from_pretrained`進行加載。
# 下面是在這種狀況下的操做方法：

#Bert模型示例
config = BertConfig.from_json_file(output_config_file)
model = BertForQuestionAnswering(config)
state_dict = torch.load(output_model_file)
model.load_state_dict(state_dict)
tokenizer = BertTokenizer(output_vocab_file, do_lower_case=args.do_lower_case)

#GPT模型示例
config = OpenAIGPTConfig.from_json_file(output_config_file)
model = OpenAIGPTDoubleHeadsModel(config)
state_dict = torch.load(output_model_file)
model.load_state_dict(state_dict)
tokenizer = OpenAIGPTTokenizer(output_vocab_file)

原文連接：https://huggingface.co/transf...

歡迎關注磐創AI博客站：
http://panchuang.net/

OpenCV中文官方文檔：
http://woshicver.com/

歡迎關注磐創博客資源彙總站：
http://docs.panchuang.net/