pytorch模型转为TensorFlow模型的代码

将模型pytorch_model.bin转成bert_model.ckpt的方法:

import numpy as np
import tensorflow as tf
import torch
v_list=[]

for k,v in torch.load('pytorch_model.bin',map_location='cpu').items():
    v_list.append(np.array(v))

def change(ckpt_path, new_ckpt_path):
    index = 0
    with tf.Session() as sess:
        for var_name, _ in tf.contrib.framework.list_variables(ckpt_path):
            print(var_name)
            var = tf.contrib.framework.load_variable(ckpt_path, var_name)
            var = tf.Variable(v_list[index])
            index+=1

        saver = tf.train.Saver()
        sess.run(tf.global_variables_initializer())
        saver.save(sess, new_ckpt_path)

ckpt_path = './chinese_L-12_H-768_A-12/bert_model.ckpt'
new_ckpt_path = './bert_model.ckpt'
change(ckpt_path, new_ckpt_path)