Tensorflow ckpt模型加载时的命名映射

举报
dailydreamer 发表于 2019/10/10 19:17:48 2019/10/10
【摘要】 tensorflow代码在加载ckpt模型时一般会使用如下代码:saver = tf.train.Saver()with tf.Session() as sess: saver.restore(sess, "model.ckpt")那么问题来了,如果想把开源的ckpt模型加载到自己代码里不同的变量命名空间中, 应该如何操作呢?查看tensorflow文档可以发现,tf.train.Save...

tensorflow代码在加载ckpt模型时一般会使用如下代码:

saver = tf.train.Saver()
with tf.Session() as sess:
		saver.restore(sess, "model.ckpt")

那么问题来了,如果想把开源的ckpt模型加载到自己代码里不同的变量命名空间中, 应该如何操作呢?


查看tensorflow文档可以发现,tf.train.Saver() 接受一个 var_list 参数,可以是一个 dict[str, SaveableObject],即一个名字到tf变量的映射,利用它即可实现模型加载时的命名映射。


举一个栗子。在kinetics-i3d代码中,有如下一段

rgb_variable_map = {}
for variable in tf.global_variables():
if variable.name.split('/')[0] == 'RGB':
  if eval_type == 'rgb600':
    rgb_variable_map[variable.name.replace(':0', '')[len('RGB/inception_i3d/'):]] = variable
  else:
    rgb_variable_map[variable.name.replace(':0', '')] = variable

rgb_saver = tf.train.Saver(var_list=rgb_variable_map, reshape=True)

功能是根据不同 eval_type 加载模型的时候,如果 eval_type == 'rgb600' ,就需要把待加载模型中的每个权重的名字都映射到计算图中加了 'RGB/inception_i3d/' 前缀的名字的变量。这样就实现了命名映射。


【版权声明】本文为华为云社区用户原创内容,转载时必须标注文章的来源(华为云社区)、文章链接、文章作者等基本信息, 否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。