TensorFlow如何加载部分模型

本文介绍如何使用TensorFlow读取checkpoint文件中的变量名及值,包括如何打印变量信息和加载部分模型。通过实例展示了从checkpoint文件中获取变量及其形状,并演示了如何排除特定层的参数,仅加载所需部分。

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

查看TensorFlow checkpoint文件中的变量名和对应值

from tensorflow.python import pywrap_tensorflow
checkpoint_path = os.path.join(model_dir, "model.ckpt")
reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in var_to_shape_map:
    print("tensor_name: ", key)
    print(reader.get_tensor(key)) # Remove this is you want to print only variable names

TensorFlow加载部分模型

#得到该网络中,所有可以加载的参数
variables = tf.contrib.framework.get_variables_to_restore()
for v in variables:
    print(v.name,v.shape)
#删除vgg_16_text层中的参数
variables_to_restore = [v for v in variables if v.name.split('/')[0]!='vgg_16_text']
#构建这部分参数的saver
saver = tf.train.Saver(variables_to_restore)
saver.restore(sess,checkpoint_path)

参考资料:

tensorflow 恢复部分参数、加载指定参数

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值