[Tensorflow] Remove nodes & add nodes with pre-trained model

Chenchou LO
Sep 7, 2018 · 5 min read

最近有個需求是希望在原本pre-trained好的model中,把原本某個embedding look-up換掉改成不同的shape。因為這個embedding lookup的shape是綁定特定的dataset,想要對另一個dataset fine-tune所以shape要match新的dataset,這樣就等於把原本pre-trained model中的nodes拿掉,並且補上不同新nodes。


首先基本的restore from checkpoint的方式就是官方的guide…

這樣子創saver = tf.train.Saver()沒有指定var_list的話就是預設全部vars weights都restore,然後載入的ckpt是’/temp/model.ckpt’

可是由於我現在model裡面有新的nodes,直接這樣做會有以下的error

NotFoundError (see above for traceback): Key y_vcc/y_emb/Adam_1 not found in checkpoint
[[Node: save/RestoreV2_122 = RestoreV2[dtypes=[DT_FLOAT], _device=”/job:localhost/replica:0/task:0/device:CPU:0"](_arg_save/Const_0_0, save/RestoreV2_122/tensor_names, save/RestoreV2_122/shape_and_slices)]]

其中y_vcc是我新創的nodes,因為不在model.ckpt中。

所以在創建tf.train.Saver()的時候要帶你要restore的var_list才可以避免這個錯誤,

這code是從這討論串來的,先用optimistic_restore_vars拿checkpoint中vars有在現在model中出現的var_list,然後再拿這個restore_vars來init tf.train.Saver()。然後因為新增的nodes還是要initial,所以這邊也加上了tf.global_variables_initializer()給tf.train.Supervisor,不然會報Variables not initialized的error。

以上是不要全部restore並且要加新的nodes的方式。


BUT…

照上面的作法,有specify restore var_list並且設global_variables_initializer給新nodes後,再用設定好的tf.train.Supervisor來開session。

還是報Variables not initialized的error。

好吧,看起來用tf.train.Supervisor並給上述的還是沒辦法work,

目前就我的認知一般常見開session的方式有以下三種:

1. tf.Session() : 最naive。

2. tf.train.Supervisor() : 最fancy,可以設定loop來定時吐log、存model.ckpt、設定initializer、幫你從最後的model.ckpt繼續train以及設定saver。

3. tf.train.MonitoredTrainingSession() : 介於上述兩個方式之間,可以給定hook class來設定Session跑前跑後要做什麼事情。

於是我就換成最naive的方式,在tf.Session下先session.run(initializer)後再用設定好restore_var的Saver來restore weights,然後一進iteration就卡住了,完全不知道為什麼,什麼error都沒有就是不動,而且我把vars的weights全部dump出來看也確定該restore的有好好的restore然後該init的也有好好的init。

因為我的input是透過tf.train.string_input_producer及tf.train.shuffle_batch來餵的,合理懷疑是因為這樣的關係沒辦法用tf.Session來開session並且正常的queue data進來。

總之最後改成用tf.train.MonitoredTrainingSession,然後其他code都沒動的情況下,就可以正常的開始training惹,以上是對於這整個流程的workaround…

Welcome to a place where words matter. On Medium, smart voices and original ideas take center stage - with no ads in sight. Watch
Follow all the topics you care about, and we’ll deliver the best stories for you to your homepage and inbox. Explore
Get unlimited access to the best stories on Medium — and support writers while you’re at it. Just $5/month. Upgrade