Xlearning上TensorFlow保存checkpoint

场景描述：

1. 构建模型时，会加载本地预训练的embedding向量到numpy array，数组大小为100w*128
2. 再把numpy array赋值给模型的一个变量，后面参与训练

#代码片段
#self.init_emb is a file path
emb_dict= np.loadtxt(self.init_emb, dtype=np.float32)
self.item_mat = tf.get_variable(name="item_mat", shape=[938280, 128], 
                                dtype=tf.float32,
                                initializer=tf.constant(emb_dict))

3. 定时保存checkpoint到HDFS路径

出现问题：

保存 checkpoint时报OOM；通过Xlearning上的log可以查看到；

跟进排查：

1. 调用tf.saver.save(session,path)保存checkpoint时，默认会生成4个文件：

checkpoint   //存储checkpoint的状态，路径信息等
xxx.data      //graph中变量的value
xxx.index    //变量名和在data中的位置索引表
xxx.meta    //graph的结构，包括 GraphDef, SaverDef等

2. 发现.meta文件很大，接近400M+；
3. 把checkpoint保存到本地，正常。另外发现和HDFS上的文件大小不一样；确认是写.meta文件到HDFS时写挂了；

初步结论

• TensorFlow通过numpy等外部数据初始化的变量时，这个外部数据会作为一个常量(Const)存储在Graph中。导致写checkpoint时meta文件很大，触发写HDFS时报OOM的问题；

解决方式：

1. 减少metagraph的大小，把变量通过feed_dict的方式进行初始化，不过这样写变量初始化的逻辑要比较tricky。其他变量需要区分出来单独进行初始化，使用tf.global_variables_initializer()需要小心一点。

pl = tf.placeholder(np.float32, shape=[1000, 2000])
v3 = tf.get_variable("v3", initializer=pl)
with tf.Session() as sess:
  matrix = np.random.rand(1000, 2000)
  sess.run(v3.initializer, feed_dict={pl: matrix},run_metadata=meta)

2. 保存checkpoint时不写meta文件，可以暂时绕过。不过最后存储模型时还是需要保存Graph结构的。

  #把write_meta_graph 设置成False
  save_path = saver.save(sess, "/tmp/model/ckpt", write_meta_graph=False, global_step=xxx)

3. 通过框架层面来解决checkpoint上传到HDFS的问题（tiny+目前支持把checkpoint上传到指定HDFS路径）
4. 查证写HDFS的原因，时间原因暂时没去深入跟。

后续跟进

1.TensorFlow的Graph 及模型保存
2.变量的初始化(通过tf op进行初始化和np array初始化的流程细节)
3.写HDFS的这个bug

参考资料

1.tensorflow：A Tool Developer Guide To TFModelFIles
2.图的核心数据结构
3.files in save model

附件图片

1.通过赋值np array的const

2.通过tf.constant构造的const