update convert tftrt

Author: hunglc007

convert_trt.py

@@ -1,6 +1,9 @@
-import tensorflow as tf
 from absl import app, flags, logging
 from absl.flags import FLAGS
+import tensorflow as tf
+physical_devices = tf.config.experimental.list_physical_devices('GPU')
+if len(physical_devices) > 0:
+    tf.config.experimental.set_memory_growth(physical_devices[0], True)
 import numpy as np
 import cv2
 from tensorflow.python.compiler.tensorrt import trt_convert as trt
@@ -14,47 +17,53 @@
 flags.DEFINE_string('output', './checkpoints/yolov4-trt-fp16-416', 'path to output')
 flags.DEFINE_integer('input_size', 416, 'path to output')
 flags.DEFINE_string('quantize_mode', 'float16', 'quantize mode (int8, float16)')
-flags.DEFINE_string('dataset', "./coco_dataset/coco/5k.txt", 'path to dataset')
-flags.DEFINE_integer('loop', 10, 'loop')
+flags.DEFINE_string('dataset', "/media/user/Source/Data/coco_dataset/coco/5k.txt", 'path to dataset')
+flags.DEFINE_integer('loop', 8, 'loop')
 
 def representative_data_gen():
   fimage = open(FLAGS.dataset).read().split()
+  batched_input = np.zeros((FLAGS.loop, FLAGS.input_size, FLAGS.input_size, 3), dtype=np.float32)
   for input_value in range(FLAGS.loop):
     if os.path.exists(fimage[input_value]):
       original_image=cv2.imread(fimage[input_value])
       original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
-      image_data = utils.image_preprocess(np.copy(original_image), [FLAGS.input_size, FLAGS.input_size])
+      image_data = utils.image_preporcess(np.copy(original_image), [FLAGS.input_size, FLAGS.input_size])
       img_in = image_data[np.newaxis, ...].astype(np.float32)
-      batched_input = tf.constant(img_in)
+      batched_input[input_value, :] = img_in
+      # batched_input = tf.constant(img_in)
       print(input_value)
-      yield (batched_input, )
+      # yield (batched_input, )
+      # yield tf.random.normal((1, 416, 416, 3)),
     else:
       continue
+  batched_input = tf.constant(batched_input)
+  yield (batched_input,)
 
 def save_trt():
+
   if FLAGS.quantize_mode == 'int8':
     conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
       precision_mode=trt.TrtPrecisionMode.INT8,
-      max_workspace_size_bytes=8000000000,
+      max_workspace_size_bytes=4000000000,
       use_calibration=True,
-      max_batch_size=32)
+      max_batch_size=8)
     converter = trt.TrtGraphConverterV2(
       input_saved_model_dir=FLAGS.weights,
       conversion_params=conversion_params)
     converter.convert(calibration_input_fn=representative_data_gen)
   elif FLAGS.quantize_mode == 'float16':
     conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
       precision_mode=trt.TrtPrecisionMode.FP16,
-      max_workspace_size_bytes=8000000000,
-      max_batch_size=32)
+      max_workspace_size_bytes=4000000000,
+      max_batch_size=8)
     converter = trt.TrtGraphConverterV2(
       input_saved_model_dir=FLAGS.weights, conversion_params=conversion_params)
     converter.convert()
   else :
     conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
       precision_mode=trt.TrtPrecisionMode.FP32,
-      max_workspace_size_bytes=8000000000,
-      max_batch_size=32)
+      max_workspace_size_bytes=4000000000,
+      max_batch_size=8)
     converter = trt.TrtGraphConverterV2(
       input_saved_model_dir=FLAGS.weights, conversion_params=conversion_params)
     converter.convert()