tensorflow模型的保存与恢复,以及ckpt到pb的转化
一、模型的保存
使用tensorflow训练模型的过程中,需要适时对模型进行保存,以及对保存的模型进行restore,以便后续对模型进行处理。如:测试、部署、拿别的模型进行fine-tune等。
保存模型是整个内容的第一步,操作十分简单,只需要创建一个saver,并在一个Session里完成保存。
saver = tf.train.Saver()
with tf.Session() as sess:
saver.save(sess, model_name)
以上代码在0.11以下版本的tensorflow里会保存与下面类似的3个文件
checkpoint model.ckpt-1000.meta model.ckpt-1000.ckpt
在0.11及以上版本的tensorflow里则会保存与下类似的4个文件
checkpoint model.ckpt-1000.index model.ckpt-1000.data-00000-of-00001 model.ckpt-1000.meta
其中checkpoint列出保存的所有模型以及最近的模型;meta文件是模型定义的内容;ckpt(或data和index)文件是保存的模型数据。
除了上面最简单的保存方式,也可以指定保存的步数,多长时间保存一次,磁盘上最多保存几个模型(将前面的删除以保持固定个数),需要做的是在创建saver时指定参数
saver = tf.train.Saver(savable_variables, max_to_keep=n, keep_checkpoint_every_n_hours=m)
其中,savable_variables指定待保存的变量,比如指定为tf.global_variables()保存所有global变量;指定为[v1, v2]保存v1和v2两个变量,如果省略,则保存所有。
max_to_keep指定磁盘上最多保存有几个模型。
keep_checkpoint_every_n_hours指定多少小时保存一次。
保存模型时指定参数
saver.save(sess, 'model_name', global_step=step, write_meta_graph=False)
其中,可以指定模型文件名,步数,write_meta_graph则用来指定是否保存meta文件记录graph,等等。
二、模型的恢复及查看模型参数
with tf.Session() as sess: # 加载模型定义的graph saver = tf.train.import_meta_graph('model.ckpt-1000.meta') # 方式一:加载指定文件夹下最近保存的一个模型的数据 saver.restore(sess, tf.train.latest_checkpoint('./')) # 方式二:指定具体某个数据,需要注意的是,指定的文件不要包含后缀 # saver.restore(sess, os.path.join(path, 'model.ckpt-1000')) # 查看模型中的trainable variables tvs = [v for v in tf.trainable_variables()] for v in tvs: print(v.name) print(sess.run(v)) # 查看模型中的所有tensor或者operations gv = [v for v in tf.global_variables()] for v in gv: print(v.name) # 获得几乎所有的operations相关的tensor ops = [o for o in sess.graph.get_operations()] for o in ops: print(o.name)
说明:
1、global_variables()比trainable_variables()多了一些非trainable的变量,比如定义时指定为trainable=False的变量,或Optimizer相关的变量。
2、sess.graph.get_operations()可以换为tf.get_default_graph().get_operations(),二者区别无非是graph明确的时候可以直接使用前者,否则需要使用后者。
三、将ckpt转化为pb
freeze_graph就是将模型固化,具体说就是将训练数据和模型固化成pb文件。
参数: (必选: 表示必须有值;可选: 表示可以为空):
1、input_graph:(必选)模型文件,可以是二进制的pb文件,或文本的meta文件,用input_binary来指定区分(见下面说明)
2、input_saver:(可选)Saver解析器。保存模型和权限时,Saver也可以自身序列化保存,以便在加载时应用合适的版本。主要用于版本不兼容时使用。可以为空,为空时用当前版本的Saver。
3、input_binary:(可选)配合input_graph用,为true时,input_graph为二进制,为false时,input_graph为文件。默认False
4、input_checkpoint:(必选)检查点数据文件。训练时,给Saver用于保存权重、偏置等变量值。这时用于模型恢复变量值。
5、output_node_names:(必选)输出节点的名字,有多个时用逗号分开。用于指定输出节点,将没有在输出线上的其它节点剔除。
6、restore_op_name:(可选)从模型恢复节点的名字。升级版中已弃用。默认:save/restore_all
7、filename_tensor_name:(可选)已弃用。默认:save/Const:0
8、output_graph:(必选)用来保存整合后的模型输出文件。
9、clear_devices:(可选),默认True。指定是否清除训练时节点指定的运算设备(如cpu、gpu、tpu。cpu是默认)
10、initializer_nodes:(可选)默认空。权限加载后,可通过此参数来指定需要初始化的节点,用逗号分隔多个节点名字。
11、variable_names_blacklist:(可先)默认空。变量黑名单,用于指定不用恢复值的变量,用逗号分隔多个变量名字。
if __name__ == '__main__': args = parse_args() # model path demonet = args.demo_net dataset = args.dataset tfmodel = os.path.join('output', demonet, DATASETS[dataset][0], 'default', NETS[demonet][0]) if not os.path.isfile(tfmodel + '.meta'): print(tfmodel) raise IOError(('{:s} not found.\nDid you download the proper networks from ' 'our server and place them properly?').format(tfmodel + '.meta')) # set config tfconfig = tf.ConfigProto(allow_soft_placement=True) tfconfig.gpu_options.allow_growth = True # init session sess = tf.Session(config=tfconfig) # load network if demonet == 'vgg16': net = vgg16(batch_size=1) else: raise NotImplementedError net.create_architecture(sess, "TEST", 4, tag='default', anchor_scales=[8, 16, 32]) saver = tf.train.Saver() saver.restore(sess, tfmodel) # 保存图 tf.train.write_graph(sess.graph_def, 'pb/pb_model', 'model.pb') # 把图和参数结构一起 freeze_graph.freeze_graph('pb/pb_model/model.pb', '', False, tfmodel, 'vgg_16/cls_score/BiasAdd,vgg_16/cls_prob,vgg_16/bbox_pred/BiasAdd,vgg_16/rois/PyFunc', 'save/restore_all', 'save/Const:0', 'pb/pb_model/frozen_model.pb', False, "")