加載模型圖_Tensorflow ckpt模型加載時的命名映射

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/'

字首的名字的變量。這樣就實作了命名映射。

更多精彩内容，請滑至頂部點選右上角關注小宅哦~

作者：dailydreamer