无法将pbjs 项识别_无法将Tensorflow模型冻结到冻结的(.pb)文件中

importos,argparseimporttensorflowastf# The original freeze_graph function# from tensorflow.python.tools.freeze_graph import freeze_graphdir=os.path.dirname(os.path.realpath(__file__))deffreeze_graph(model_dir,output_node_names):"""Extract the sub graph defined by the output nodes and convert

all its variables into constant

Args:

model_dir: the root folder containing the checkpoint state file

output_node_names: a string, containing all the output node's names,

comma separated

"""ifnottf.gfile.Exists(model_dir):raiseAssertionError("Export directory doesn't exists. Please specify an export ""directory: %s"%model_dir)ifnotoutput_node_names:print("You need to supply the name of a node to --output_node_names.")return-1# We retrieve our checkpoint fullpathcheckpoint=tf.train.get_checkpoint_state(model_dir)input_checkpoint=checkpoint.model_checkpoint_path# We precise the file fullname of our freezed graphabsolute_model_dir="/".join(input_checkpoint.split('/')[:-1])output_graph=absolute_model_dir+"/frozen_model.pb"# We clear devices to allow TensorFlow to control on which device it will load operationsclear_devices=True# We start a session using a temporary fresh Graphwithtf.Session(graph=tf.Graph())assess:# We import the meta graph in the current default Graphsaver=tf.train.import_meta_graph(input_checkpoint+'.meta',clear_devices=clear_devices)# We restore the weightssaver.restore(sess,input_checkpoint)# We use a built-in TF helper to export variables to constantsoutput_graph_def=tf.graph_util.convert_variables_to_constants(sess,# The session is used to retrieve the weightstf.get_default_graph().as_graph_def(),# The graph_def is used to retrieve the nodesoutput_node_names.split(",")# The output node names are used to select the usefull nodes)# Finally we serialize and dump the output graph to the filesystemwithtf.gfile.GFile(output_graph,"wb")asf:f.write(output_graph_def.SerializeToString())print("%d ops in the final graph."%len(output_graph_def.node))returnoutput_graph_defif__name__=='__main__':parser=argparse.ArgumentParser()parser.add_argument("--model_dir",type=str,default="",help="Model folder to export")parser.add_argument("--output_node_names",type=str,default="",help="The name of the output nodes, comma separated.")args=parser.parse_args()freeze_graph(args.model_dir,args.output_node_names)

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值