yolov3模型转换为OM模型

1.下载安装依赖:https://github.com/wizyoung/YOLOv3_TensorFlow
2.将yolov3.weights放到~/test/YOLOv3_TensorFlow-master/data/darknet_weights文件夹中,voc_names.txt里写入训练的类别
3.修改test/YOLOv3_TensorFlow-master中convert_weight.py里的类别个数,运行python3 convert_weight.py
4.~/test/YOLOv3_TensorFlow-master/data/darknet_weights中生成的3个ckpt后运行python3 toPb.py生成 yolov3.pb模型,
5.拷贝aipp_yolov3_picture.cfg文件到.pb文件同一文件夹下,运行命令 omg --model ./yolov3.pb --framework 3 --output ./yolov3_facemask --insert_op_conf ./aipp_yolov3_picture.cfg --input_shape "Placeholder:1,416,416,3",生成OM文件

 

toPb.py:

import tensorflow as tf
from tensorflow.python.framework import graph_util
from tensorflow.python.platform import gfile

def freeze_graph(input_path,output_path):
output_node_names="yolov3/yolov3_head/feature_map_1,yolov3/yolov3_head/feature_map_2,yolov3/yolov3_head/feature_map_3"
saver=tf.train.import_meta_graph(input_path+".meta",clear_devices=True)
graph=tf.get_default_graph()
input_graph_def=graph.as_graph_def()

with tf.Session() as sess:
saver.restore(sess,input_path)
output_graph_def=graph_util.convert_variables_to_constants(
sess=sess,
input_graph_def=input_graph_def,
output_node_names=output_node_names.split(",")
)

with tf.gfile.GFile(output_path,'wb') as fgraph:
fgraph.write(output_graph_def.SerializeToString())

if __name__ == '__main__':
input_path ="./yolov3.ckpt"
output_path="./yolov3/yolov3.pb"
freeze_graph(input_path,output_path)

posted @ 2020-08-26 14:35  闪光123  阅读(1432)  评论(0编辑  收藏  举报