ckpt-pb-onnx-trt

ckpt

这应该是ckpt转换成自己网络名称的操作

import argparse
import tensorflow as tf
from core.yolov3 import YOLOV3
from core.config import cfg
parser = argparse.ArgumentParser()
parser.add_argument("--train_from_coco", action='store_true')
flag = parser.parse_args()

org_weights_path = cfg.YOLO.ORIGINAL_WEIGHT
cur_weights_path = cfg.YOLO.DEMO_WEIGHT
preserve_cur_names = ['conv_sbbox', 'conv_mbbox', 'conv_lbbox']
preserve_org_names = ['Conv_6', 'Conv_14', 'Conv_22']


org_weights_mess = []
tf.Graph().as_default()
load = tf.train.import_meta_graph(org_weights_path + '.meta')
with tf.Session() as sess:
    load.restore(sess, org_weights_path)
    for var in tf.global_variables():
        var_name = var.op.name
        var_name_mess = str(var_name).split('/')
        var_shape = var.shape
        if flag.train_from_coco:
            if (var_name_mess[-1] not in ['weights', 'gamma', 'beta', 'moving_mean', 'moving_variance']) or \
                    (var_name_mess[1] == 'yolo-v3' and (var_name_mess[-2] in preserve_org_names)): continue
        org_weights_mess.append([var_name, var_shape])
        print("=> " + str(var_name).ljust(50), var_shape)
print()
tf.reset_default_graph()

cur_weights_mess = []
tf.Graph().as_default()
with tf.name_scope('input'):
    input_data = tf.placeholder(dtype=tf.float32, shape=(1, 416, 416, 3), name='input_data')
    training = tf.placeholder(dtype=tf.bool, name='trainable')
model = YOLOV3(input_data, training)
for var in tf.global_variables():
    var_name = var.op.name
    var_name_mess = str(var_name).split('/')
    var_shape = var.shape
    print(var_name_mess[0])
    if flag.train_from_coco:
        if var_name_mess[0] in preserve_cur_names: continue
    cur_weights_mess.append([var_name, var_shape])
    print("=> " + str(var_name).ljust(50), var_shape)

org_weights_num = len(org_weights_mess)
cur_weights_num = len(cur_weights_mess)
if cur_weights_num != org_weights_num:
    raise RuntimeError

print('=> Number of weights that will rename:\t%d' % cur_weights_num)
cur_to_org_dict = {}
for index in range(org_weights_num):
    org_name, org_shape = org_weights_mess[index]
    cur_name, cur_shape = cur_weights_mess[index]
    if cur_shape != org_shape:
        print(org_weights_mess[index])
        print(cur_weights_mess[index])
        raise RuntimeError
    cur_to_org_dict[cur_name] = org_name
    print("=> " + str(cur_name).ljust(50) + ' : ' + org_name)

with tf.name_scope('load_save'):
    name_to_var_dict = {var.op.name: var for var in tf.global_variables()}
    restore_dict = {cur_to_org_dict[cur_name]: name_to_var_dict[cur_name] for cur_name in cur_to_org_dict}
    load = tf.train.Saver(restore_dict)
    save = tf.train.Saver(tf.global_variables())
    for var in tf.global_variables():
        print("=> " + var.op.name)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print('=> Restoring weights from:\t %s' % org_weights_path)
    load.restore(sess, org_weights_path)
    save.save(sess, cur_weights_path)
tf.reset_default_graph()

pb

将ckpt转化成pb

import tensorflow as tf
from core.yolov3 import YOLOV3

pb_file = "/data/zhangyong/ocr/tensorflow-yolov3-master/last1.pb"
ckpt_file = "/data/zhangyong/ocr/tensorflow-yolov3-master/new/yolov3_coco_demo.ckpt"
output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"]

with tf.name_scope('input'):
    input_data = tf.placeholder(dtype=tf.float32, name='input_data')

model = YOLOV3(input_data, trainable=False)
print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox)

sess  = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
saver = tf.train.Saver()
saver.restore(sess, ckpt_file)

converted_graph_def = tf.graph_util.convert_variables_to_constants(sess,
                            input_graph_def  = sess.graph.as_graph_def(),
                            output_node_names = output_node_names)

with tf.gfile.GFile(pb_file, "wb") as f:
    f.write(converted_graph_def.SerializeToString())

预测

def read_pb_return_tensors(graph, pb_file, return_elements):

    with tf.gfile.FastGFile(pb_file, 'rb') as f:
        frozen_graph_def = tf.GraphDef()
        frozen_graph_def.ParseFromString(f.read())

    with graph.as_default():
        return_elements = tf.import_graph_def(frozen_graph_def,
                                              return_elements=return_elements)
    return return_elements

import cv2
import numpy as np
import core.utils as utils
import tensorflow as tf
from PIL import Image

return_elements = ["input/input_data:0", "pred_sbbox/concat_2:0", "pred_mbbox/concat_2:0", "pred_lbbox/concat_2:0"]
pb_file         = "/data/zhangyong/ocr/tensorflow-yolov3-master/new.pb"
image_path      = "./docs/images/road.jpeg"
num_classes     = 80
input_size      = 416
graph           = tf.Graph()
original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preporcess(np.copy(original_image), [input_size, input_size])
image_data = image_data[np.newaxis, ...]
return_tensors = utils.read_pb_return_tensors(graph, pb_file, return_elements)
with tf.Session(graph=graph) as sess:
    a="import/input/trainable:0"
    pred_sbbox, pred_mbbox, pred_lbbox = sess.run(
        [return_tensors[1], return_tensors[2], return_tensors[3]],
                feed_dict={ return_tensors[0]: image_data,a: False})
pred_bbox = np.concatenate([np.reshape(pred_sbbox, (-1, 5 + num_classes)),
                            np.reshape(pred_mbbox, (-1, 5 + num_classes)),
                            np.reshape(pred_lbbox, (-1, 5 + num_classes))], axis=0)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size, input_size, 0.3)
bboxes = utils.nms(bboxes, 0.45, method='nms')
image = utils.draw_bbox(original_image, bboxes)
image = Image.fromarray(image)
image.show()
input("wait:")

onnx

pb2onnx使用的是tensorflow-onnx命令为

python -m tf2onnx.convert --input last1.pb --inputs input/input_data:0[1,416,416,3] --inputs-as-nchw input/input_data:0 --outputs pred_sbbox/concat_2:0,pred_mbbox/concat_2:0,pred_lbbox/concat_2:0 --output ./0802_1.onnx --verbose --opset 10

onnx推断为

import onnxruntime
import numpy as np
import torch
# Input to the model
from PIL import Image
import cv2
import core.utils as utils
import torchvision.transforms as transforms
img = Image.open("./road.jpg")
# ort_session = onnxruntime.InferenceSession("./last1_sim.onnx")
ort_session = onnxruntime.InferenceSession("./0802_1sim.onnx")
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
# resize = transforms.Resize([416, 416])
# img = resize(img)
# img_ycbcr = img.convert('YCbCr')
# img_y, img_cb, img_cr = img_ycbcr.split()
# to_tensor = transforms.ToTensor()
# img_y = to_tensor(img_y)
# img_y.unsqueeze_(0)
# print(img_y.shape)
# img_y=img_y.permute(0,3,1,2)
# image = Image.open("./road.jpg")
def letterbox_image(image, size):
    iw, ih = image.size
    w, h = size
    scale = min(w/iw, h/ih)
    nw = int(iw*scale)
    nh = int(ih*scale)

    image = image.resize((nw,nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (128,128,128))
    new_image.paste(image, ((w-nw)//2, (h-nh)//2))
    return new_image
original_image = cv2.imread("./road.jpg")
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preporcess(np.copy(original_image), [416, 416])
image_data = image_data[np.newaxis, ...]
image_data=image_data.astype(np.float32)
image_data=np.transpose(image_data,(0,3,1,2))
image_data1=image_data[0][1]
image_data2=image_data[0][2]
# with open("./now.txt","a+") as f:
#
#     for i in range(3):
#         for j in range(416):
#             for k in range(416):
#                 f.write(str(image_data[0][i][j][k])+"\n")

# image_data=np.transpose(image_data,(0,3,1,2)).astype(np.float32)
# image_shape = np.array(np.shape(image)[0:2])
#
# crop_img = np.array(letterbox_image(image, (416, 416)))
# photo = np.array(crop_img, dtype=np.float32)
# photo /= 255.0
# photo = photo.astype(np.float32)
# photo = photo[np.newaxis, ...]
# photo = np.transpose(photo, (0, 3, 1,2))
ort_inputs = {ort_session.get_inputs()[0].name: image_data}
ort_outs = ort_session.run(None, ort_inputs)
img_out_y = ort_outs[0].reshape(-1,85)
img2=ort_outs[1].reshape(-1,85)
img3=ort_outs[2].reshape(-1,85)
pred_bbox=np.concatenate([img_out_y,img2,img3],axis=0)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size, 416, 0.3)
bboxes = utils.nms(bboxes, 0.45, method='nms')
image = utils.draw_bbox(original_image, bboxes)
image=np.array(image)
image = Image.fromarray(image)
image.save("./1.jpg")
image.show()
input("wait:")

trt

首先将onnx sim命令为

python3 -m onnxsim 0802_1.onnx 0802_1sim.onnx

转化trt使用的是onnx-trt命令为

 onnx2trt 0802_1sim.onnx -b 1 -o 0802_1.trt

bug

转化完成后,发现最终解码的值x,w,h,score,class都正确,只有y不正确,经过debug发现y似乎没有被解码,自己在c++上解码后发现结果大部分正确

int row1=i/3;
int row=row1/yolo.height;
float a=my_Logist(intput[i*85+1]);
det.bbox[0]=intput[i*85+0];
det.bbox[1]=(row+a)/yolo.height*416;

思考

结果不是自己想要的不要慌,不能急,还是要一步步的调试,首先看看不同模型文件能否正确预测,即pb,onnx,trt能否正确预测,定位那个模型出现了问题,然后找到该模型文件的预测代码,不断的debug,从输入到输出对比,最终确定出错的原因.