Convert the pth weight of pytorch to the ckpt file of tensorflow

Time:2022-8-5
import tensorflow as tf
import torch

def convert(bin_path, ckptpath):
    with tf.Session() as sess:
        for var_name, value in torch.load(bin_path, map_location='cpu').items():
            print(var_name) # output the variable name in the weights file
            tf.Variable(initial_value=value, name=var_name)
        saver = tf.train.Saver()
        sess.run(tf.global_variables_initializer())
        saver.save(sess, ckpt_path)

bin_path = 'model.pth'
ckpt_path = 'new_model.ckpt'
convert(bin_path, ckpt_path)