tf预训练模型转换为torch预训练模型

最新推荐文章于 2025-01-17 18:00:00 发布

wanzi_antang

最新推荐文章于 2025-01-17 18:00:00 发布

阅读量1.6k

点赞数 2

分类专栏： python 自然语言处理 nlp随笔记

本文链接：https://blog.csdn.net/wanzi_antang/article/details/115011692

版权

本文详述了将TensorFlow的Albert预训练模型成功转换为PyTorch模型的过程，期间遇到并解决了诸多问题，对于需要进行模型转换的开发者具有参考价值。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在将albert的tensorflow预训练模型转换为 torch类型预训练模型，踩了很多坑。终于解决，希望对大家有用

前期准备
创建一个环境带有torch和tf的环境，步骤如下：
首先创建环境
python conda create -n torchtf_env python=3.7
然后，安装torch（根据自己电脑的cuda安装）
python conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge
之后，继续安装tensorflow-gpu版本
python conda install tensorflow-gpu==1.15
最后安装transformers
pip install transformers

2 .从github上下载tensorflow预训练的albert版本

#! usr/bin/env python3
# -*- coding:utf-8 -*-
"""
Created on 19/03/2021 20:22 
@Author: lixj
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import torch
from transformers import