关于深度学习:pytorch的pth权重转化为tensorflow的ckpt文件

33次阅读

共计 391 个字符，预计需要花费 1 分钟才能阅读完成。

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)  # 输入权重文件中的变量名
            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)

正文完