任务简介:
由于目前 PyTorch 使用者非常多,BERT 源码又是 Google 用 TensorFlow 写的,Google 开源的 BERT 预训练模型都是 TensorFlow 生成的 ckpt 模型文件, PyTorch 使用者需要将 ckpt 文件转换为 PyTorch 的 bin 模型文件。
本文以 albert 模型为例,将 albert 的 TensorFlow 模型转 PyTorch 模型。
albert 模型转换代码:
"""
@Author : ChenXin
@Date : 2021/10/23 1:26
@Brief : TensorFlow 的模型转为 PyTorch 模型
"""
import argparse
import logging
import torch
from transformers import AlbertConfig, AlbertForPreTraining, load_tf_weights_in_albert
logging.basicConfig(level=logging.INFO)
def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, albert_config_file, pytorch_dump_path):
# 初始化 PyTorch 模型
config = AlbertConfig.from_json_file(albert_config_file)
print("Building PyTorch model from configuration: {}".format(str(config)))
model = AlbertForPreTraining(config)
# 从 tf 的 checkpoint 文件中加载模型权重
load_tf_weights_in_albert(model, config, tf_checkpoint_path)
# 保存 PyTorch 模型
print("Save PyTorch model to {}".format(pytorch_dump_path))
torch.save(model.state_dict(), pytorch_dump_path)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# 参数设置
parser.add_argument("--tf_checkpoint_path", default='./resources/albert_base_zh/model.ckpt-best', type=str,
required=False, help="Path to the TensorFlow checkpoint path.")
parser.add_argument("--albert_config_file", default='./resources/albert_base_zh/albert_config.json', type=str,
required=False,
help="The config json file corresponding to the pre-trained ALBERT model. \n""This specifies the model architecture.", )
parser.add_argument("--pytorch_dump_path", default='./resources/albert_base_zh/pytorch_model.bin', type=str, required=False,
help="Path to the output PyTorch model.")
args = parser.parse_args()
convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.albert_config_file, args.pytorch_dump_path)
输出:
生成了 PyTorch 模型