【ALBERT】TensorFlow 模型转 PyTorch 模型

原创

已于 2022-05-23 18:09:46 修改 · 2.3k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #tensorflow #深度学习

于 2021-10-29 15:54:00 首次发布

本文介绍了如何将TensorFlow的ALBERT模型转换为PyTorch模型，通过`load_tf_weights_in_albert`函数加载权重，并保存为PyTorch的bin模型文件，便于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):