optimize gpu inference for tflite model

Author: hunglc007

convert_tflite.py

@@ -8,8 +8,8 @@
 import os
 from core.config import cfg
 
-flags.DEFINE_string('weights', './checkpoints/yolov3-416', 'path to weights file')
-flags.DEFINE_string('output', './checkpoints/yolov3-416-int8.tflite', 'path to output')
+flags.DEFINE_string('weights', './checkpoints/yolov4-416', 'path to weights file')
+flags.DEFINE_string('output', './checkpoints/yolov4-416-fp32.tflite', 'path to output')
 flags.DEFINE_integer('input_size', 416, 'path to output')
 flags.DEFINE_string('quantize_mode', 'float32', 'quantize mode (int8, float16, float32)')
 flags.DEFINE_string('dataset', "/Volumes/Elements/data/coco_dataset/coco/5k.txt", 'path to dataset')

core/yolov4.py

@@ -192,15 +192,16 @@ def decode_train(conv_output, output_size, NUM_CLASS, STRIDES, ANCHORS, i=0, XYS
     return tf.concat([pred_xywh, pred_conf, pred_prob], axis=-1)
 
 def decode_tf(conv_output, output_size, NUM_CLASS, STRIDES, ANCHORS, i=0, XYSCALE=[1, 1, 1]):
+    batch_size = tf.shape(conv_output)[0]
     conv_output = tf.reshape(conv_output,
-                             (tf.shape(conv_output)[0], output_size, output_size, 3, 5 + NUM_CLASS))
+                             (batch_size, output_size, output_size, 3, 5 + NUM_CLASS))
 
     conv_raw_dxdy, conv_raw_dwdh, conv_raw_conf, conv_raw_prob = tf.split(conv_output, (2, 2, 1, NUM_CLASS),
                                                                           axis=-1)
 
     xy_grid = tf.meshgrid(tf.range(output_size), tf.range(output_size))
     xy_grid = tf.expand_dims(tf.stack(xy_grid, axis=-1), axis=2)  # [gx, gy, 1, 2]
-    xy_grid = tf.tile(tf.expand_dims(xy_grid, axis=0), [tf.shape(conv_output)[0], 1, 1, 3, 1])
+    xy_grid = tf.tile(tf.expand_dims(xy_grid, axis=0), [batch_size, 1, 1, 3, 1])
 
     xy_grid = tf.cast(xy_grid, tf.float32)
 
@@ -213,40 +214,55 @@ def decode_tf(conv_output, output_size, NUM_CLASS, STRIDES, ANCHORS, i=0, XYSCAL
     pred_prob = tf.sigmoid(conv_raw_prob)
 
     pred_prob = pred_conf * pred_prob
+    pred_prob = tf.reshape(pred_prob, (batch_size, -1, NUM_CLASS))
+    pred_xywh = tf.reshape(pred_xywh, (batch_size, -1, 4))
+
     return pred_xywh, pred_prob
     # return tf.concat([pred_xywh, pred_conf, pred_prob], axis=-1)
 
 def decode_tflite(conv_output, output_size, NUM_CLASS, STRIDES, ANCHORS, i=0, XYSCALE=[1,1,1]):
-    conv_output = tf.reshape(conv_output, (1, output_size, output_size, 3, 5 + NUM_CLASS))
-
-    conv_raw_dxdy, conv_raw_dwdh, conv_raw_conf, conv_raw_prob = tf.split(conv_output, (2, 2, 1, NUM_CLASS), axis=-1)
+    conv_raw_dxdy_0, conv_raw_dwdh_0, conv_raw_score_0,\
+    conv_raw_dxdy_1, conv_raw_dwdh_1, conv_raw_score_1,\
+    conv_raw_dxdy_2, conv_raw_dwdh_2, conv_raw_score_2 = tf.split(conv_output, (2, 2, 1+NUM_CLASS, 2, 2, 1+NUM_CLASS,
+                                                                                2, 2, 1+NUM_CLASS), axis=-1)
+
+    conv_raw_score = [conv_raw_score_0, conv_raw_score_1, conv_raw_score_2]
+    for idx, score in enumerate(conv_raw_score):
+        score = tf.sigmoid(score)
+        score = score[:, :, :, 0:1] * score[:, :, :, 1:]
+        conv_raw_score[idx] = tf.reshape(score, (1, -1, NUM_CLASS))
+    pred_prob = tf.concat(conv_raw_score, axis=1)
+
+    conv_raw_dwdh = [conv_raw_dwdh_0, conv_raw_dwdh_1, conv_raw_dwdh_2]
+    for idx, dwdh in enumerate(conv_raw_dwdh):
+        dwdh = tf.exp(dwdh) * ANCHORS[i][idx]
+        conv_raw_dwdh[idx] = tf.reshape(dwdh, (1, -1, 2))
+    pred_wh = tf.concat(conv_raw_dwdh, axis=1)
 
     xy_grid = tf.meshgrid(tf.range(output_size), tf.range(output_size))
-    xy_grid = tf.expand_dims(tf.stack(xy_grid, axis=-1), axis=2)  # [gx, gy, 1, 2]
-    xy_grid = tf.tile(tf.expand_dims(xy_grid, axis=0), [1, 1, 1, 3, 1])
-
+    xy_grid = tf.stack(xy_grid, axis=-1)  # [gx, gy, 2]
+    xy_grid = tf.expand_dims(xy_grid, axis=0)
     xy_grid = tf.cast(xy_grid, tf.float32)
 
-    pred_xy = ((tf.sigmoid(conv_raw_dxdy) * XYSCALE[i]) - 0.5 * (XYSCALE[i] - 1) + xy_grid) * \
+    conv_raw_dxdy = [conv_raw_dxdy_0, conv_raw_dxdy_1, conv_raw_dxdy_2]
+    for idx, dxdy in enumerate(conv_raw_dxdy):
+        dxdy = ((tf.sigmoid(dxdy) * XYSCALE[i]) - 0.5 * (XYSCALE[i] - 1) + xy_grid) * \
               STRIDES[i]
-    pred_wh = (tf.exp(conv_raw_dwdh) * ANCHORS[i])
+        conv_raw_dxdy[idx] = tf.reshape(dxdy, (1, -1, 2))
+    pred_xy = tf.concat(conv_raw_dxdy, axis=1)
     pred_xywh = tf.concat([pred_xy, pred_wh], axis=-1)
-
-    pred_conf = tf.sigmoid(conv_raw_conf)
-    pred_prob = tf.sigmoid(conv_raw_prob)
-
-    pred_prob = pred_conf * pred_prob
     return pred_xywh, pred_prob
     # return tf.concat([pred_xywh, pred_conf, pred_prob], axis=-1)
 
 def decode_trt(conv_output, output_size, NUM_CLASS, STRIDES, ANCHORS, i=0, XYSCALE=[1,1,1]):
-    conv_output = tf.reshape(conv_output, (tf.shape(conv_output)[0], output_size, output_size, 3, 5 + NUM_CLASS))
+    batch_size = tf.shape(conv_output)[0]
+    conv_output = tf.reshape(conv_output, (batch_size, output_size, output_size, 3, 5 + NUM_CLASS))
 
     conv_raw_dxdy, conv_raw_dwdh, conv_raw_conf, conv_raw_prob = tf.split(conv_output, (2, 2, 1, NUM_CLASS), axis=-1)
 
     xy_grid = tf.meshgrid(tf.range(output_size), tf.range(output_size))
     xy_grid = tf.expand_dims(tf.stack(xy_grid, axis=-1), axis=2)  # [gx, gy, 1, 2]
-    xy_grid = tf.tile(tf.expand_dims(xy_grid, axis=0), [tf.shape(conv_output)[0], 1, 1, 3, 1])
+    xy_grid = tf.tile(tf.expand_dims(xy_grid, axis=0), [batch_size, 1, 1, 3, 1])
 
     # x = tf.tile(tf.expand_dims(tf.range(output_size, dtype=tf.float32), axis=0), [output_size, 1])
     # y = tf.tile(tf.expand_dims(tf.range(output_size, dtype=tf.float32), axis=1), [1, output_size])
@@ -258,14 +274,17 @@ def decode_trt(conv_output, output_size, NUM_CLASS, STRIDES, ANCHORS, i=0, XYSCA
     # pred_xy = ((tf.sigmoid(conv_raw_dxdy) * XYSCALE[i]) - 0.5 * (XYSCALE[i] - 1) + xy_grid) * \
     #           STRIDES[i]
     pred_xy = (tf.reshape(tf.sigmoid(conv_raw_dxdy), (-1, 2)) * XYSCALE[i] - 0.5 * (XYSCALE[i] - 1) + tf.reshape(xy_grid, (-1, 2))) * STRIDES[i]
-    pred_xy = tf.reshape(pred_xy, (tf.shape(conv_output)[0], output_size, output_size, 3, 2))
+    pred_xy = tf.reshape(pred_xy, (batch_size, output_size, output_size, 3, 2))
     pred_wh = (tf.exp(conv_raw_dwdh) * ANCHORS[i])
     pred_xywh = tf.concat([pred_xy, pred_wh], axis=-1)
 
     pred_conf = tf.sigmoid(conv_raw_conf)
     pred_prob = tf.sigmoid(conv_raw_prob)
 
     pred_prob = pred_conf * pred_prob
+
+    pred_prob = tf.reshape(pred_prob, (batch_size, -1, NUM_CLASS))
+    pred_xywh = tf.reshape(pred_xywh, (batch_size, -1, 4))
     return pred_xywh, pred_prob
     # return tf.concat([pred_xywh, pred_conf, pred_prob], axis=-1)
 

detect.py

@@ -55,10 +55,10 @@ def main(_argv):
         interpreter.set_tensor(input_details[0]['index'], images_data)
         interpreter.invoke()
         pred = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))]
-        if FLAGS.model == 'yolov4' and FLAGS.tiny == True:
-            boxes, pred_conf = filter_boxes(pred[1], pred[0], score_threshold=0.25)
+        if FLAGS.model == 'yolov3' and FLAGS.tiny == True:
+            boxes, pred_conf = filter_boxes(pred[1], pred[0], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))
         else:
-            boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25)
+            boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25, input_shape=tf.constant([input_size, input_size]))
     else:
         saved_model_loaded = tf.saved_model.load(FLAGS.weights, tags=[tag_constants.SERVING])
         infer = saved_model_loaded.signatures['serving_default']

detectvideo.py

@@ -63,10 +63,12 @@ def main(_argv):
             interpreter.set_tensor(input_details[0]['index'], image_data)
             interpreter.invoke()
             pred = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))]
-            if FLAGS.model == 'yolov4' and FLAGS.tiny == True:
-                boxes, pred_conf = filter_boxes(pred[1], pred[0], score_threshold=0.25)
+            if FLAGS.model == 'yolov3' and FLAGS.tiny == True:
+                boxes, pred_conf = filter_boxes(pred[1], pred[0], score_threshold=0.25,
+                                                input_shape=tf.constant([input_size, input_size]))
             else:
-                boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25)
+                boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25,
+                                                input_shape=tf.constant([input_size, input_size]))
         else:
             batch_data = tf.constant(image_data)
             pred_bbox = infer(batch_data)

save_model.py

@@ -38,14 +38,12 @@ def save_tf():
         output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE, FLAGS.framework)
       bbox_tensors.append(output_tensors[0])
       prob_tensors.append(output_tensors[1])
-  pred_bbox = [tf.reshape(x, (tf.shape(x)[0], -1, tf.shape(x)[-1])) for x in bbox_tensors]
-  pred_bbox = tf.concat(pred_bbox, axis=1)
-  pred_prob = [tf.reshape(x, (tf.shape(x)[0], -1, tf.shape(x)[-1])) for x in prob_tensors]
-  pred_prob = tf.concat(pred_prob, axis=1)
+  pred_bbox = tf.concat(bbox_tensors, axis=1)
+  pred_prob = tf.concat(prob_tensors, axis=1)
   if FLAGS.framework == 'tflite':
     pred = (pred_bbox, pred_prob)
   else:
-    boxes, pred_conf = filter_boxes(pred_bbox, pred_prob, score_threshold=FLAGS.score_thres)
+    boxes, pred_conf = filter_boxes(pred_bbox, pred_prob, score_threshold=FLAGS.score_thres, input_shape=tf.constant([FLAGS.input_size, FLAGS.input_size]))
     pred = tf.concat([boxes, pred_conf], axis=-1)
   model = tf.keras.Model(input_layer, pred)
   utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny)