読者です 読者をやめる 読者になる 読者になる

KZKY memo

自分用メモ.

TensorFlow: Save and Restore Graph/Variable

tensorflow

TensorflowのSaverでsave/restoreされるのは,MetaGarph/Variableなどの情報.graphだけを永続化したい場合は,tf.train.write_graph/tf.import_graph_defを使う.

同じグラフを常に使う場合,単なるチェックポイントからの復帰は,これのみで可能.

Save and Restore Graph/Variable

元のグラフ,変数を保存後,別のグラフにimportして,RestoreしたVariableにアクセスするサンプル.

  1. Variableを作って
  2. Default Graphに Opを追加
  3. VariableをCollectionに追加しておいて,Restore後にアクセスしやすくしておく
  4. tf.train.Saver.saveでsave
  5. tf.Graphで別のグラフを作って,graph.as_default()でコンテキストをつくり
  6. saver.restore()
  7. ユーザ定義のcollectionへアクセスして,変数へアクセス
#!/usr/bin/env python

import tensorflow as tf
import datetime

def main():
    # Varialbes
    initial_v1 = tf.truncated_normal((10, 5), mean=0, stddev=0.1)
    #v1 = tf.Variable(initial_v1, name="v1")
    v1 = tf.get_variable("v1", initializer=initial_v1)

    initial_v2 = tf.truncated_normal((5, 3), mean=0, stddev=0.1)
    #v2 = tf.Variable(initial_v2, name="v2")
    v2 = tf.get_variable("v2", initializer=initial_v2)
    tf.add_to_collection("v2", v2)
    
    # Ops
    op_name = "prod"
    prod = tf.matmul(v1, v2, name=op_name)
    init_op = tf.initialize_all_variables()

    # Saver
    saver = tf.train.Saver([v1, v2])
    #saver = tf.train.Saver()

    with tf.Session() as sess:
        # Init variables
        sess.run(init_op)

        # Compute
        result0 = sess.run(prod)
        print result0
        print ""

        # Save variables
        now = datetime.datetime.now()
        save_path = "/tmp/{}-{}.ckpt".format(__file__, now)
        saver.save(sess, save_path, write_meta_graph=True)
        print "Model saved in {}".format(save_path)
        print ""
        
    # Restore graph to another graph (not default graph) and variables
    graph = tf.Graph()
    with graph.as_default():
        saver = tf.train.import_meta_graph("{}.meta".format(save_path))

        with tf.Session() as sess:
            saver.restore(sess, save_path)
            print "Variables restored"
            print ""

            print [var.name for var in tf.get_collection(tf.GraphKeys.VARIABLES)]
            print [var.eval() for var in tf.get_collection(tf.GraphKeys.VARIABLES)]

            print tf.get_collection("v2")[0].eval()
            #print tf.get_variable("v2:0") can NOT get
            print ""
            
if __name__ == '__main__':
    main()

正直,save/restore APIが使いにくいので,そのうち変更されるかもしれない... これのようにハックしている人もいる.