【ALBERT】TensorFlow 模型转 PyTorch 模型

任务简介：

由于目前 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 模型