场景描述:
- 构建模型时,会加载本地预训练的embedding向量到numpy array,数组大小为100w*128
- 再把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))
- 定时保存checkpoint到HDFS路径
出现问题:
保存 checkpoint时报OOM;通过Xlearning上的log可以查看到;
跟进排查:
- 调用tf.saver.save(session,path)保存checkpoint时,默认会生成4个文件:
checkpoint //存储checkpoint的状态,路径信息等
xxx.data //graph中变量的value
xxx.index //变量名和在data中的位置索引表
xxx.meta //graph的结构,包括 GraphDef, SaverDef等
- 发现.meta文件很大,接近400M+;
- 把checkpoint保存到本地,正常。另外发现和HDFS上的文件大小不一样;确认是写.meta文件到HDFS时写挂了;
初步结论
- TensorFlow通过numpy等外部数据初始化的变量时,这个外部数据会作为一个常量(Const)存储在Graph中。导致写checkpoint时meta文件很大,触发写HDFS时报OOM的问题;
解决方式:
- 减少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)
- 保存checkpoint时不写meta文件,可以暂时绕过。不过最后存储模型时还是需要保存Graph结构的。
#把write_meta_graph 设置成False
save_path = saver.save(sess, "/tmp/model/ckpt", write_meta_graph=False, global_step=xxx)
- 通过框架层面来解决checkpoint上传到HDFS的问题(tiny+目前支持把checkpoint上传到指定HDFS路径)
- 查证写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