Object Detection Refactored

Author: DEKHTIARJonathan

tftrt/examples/benchmark_utils.py

@@ -183,7 +183,7 @@ def aggregate_data(self, y_pred, y):
             self._total_samples_processed += step_batch_size
             idx_stop = self._total_samples_processed
 
-            with timed_section("Copy Time",
+            with timed_section("Numpy Copy Time",
                                activate=self._args.debug_performance,
                                start_end_mode=False):
                 for key, val in self._predicted.items():

tftrt/examples/image_classification/image_classification.py

@@ -15,14 +15,9 @@
 # limitations under the License.
 # =============================================================================
 
-import multiprocessing
 import os
 import sys
 
-import multiprocessing
-
-from functools import partial
-
 import numpy as np
 
 import tensorflow as tf
@@ -169,7 +164,7 @@ def preprocess_sample_fn(record):
 
         dataset = dataset.interleave(
             tf.data.TFRecordDataset,
-            cycle_length=min(8, multiprocessing.cpu_count()),
+            cycle_length=tf.data.experimental.AUTOTUNE,
             block_length=max(self._args.batch_size, 32)
         )
 
@@ -181,7 +176,7 @@ def preprocess_sample_fn(record):
 
         dataset = dataset.map(
             map_func=preprocess_fn,
-            num_parallel_calls=min(8, multiprocessing.cpu_count())
+            num_parallel_calls=tf.data.experimental.AUTOTUNE,
         )
 
         dataset = dataset.batch(self._args.batch_size, drop_remainder=False)

tftrt/examples/object_detection/object_detection.py

@@ -16,16 +16,12 @@
 # =============================================================================
 
 import os
-import sys
-
-import logging
-import time
 import shutil
+import sys
 
-from functools import partial
+import numpy as np
 import ujson as json
 
-import numpy as np
 import tensorflow as tf
 
 from pycocotools.coco import COCO
@@ -46,7 +42,7 @@
 
 class CommandLineAPI(BaseCommandLineAPI):
 
-    SAMPLES_IN_VALIDATION_SET = 5000
+    SAMPLES_IN_VALIDATION_SET = 50000
 
     def __init__(self):
         super(CommandLineAPI, self).__init__()
@@ -66,40 +62,110 @@ def __init__(self):
         )
 
 
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
+# %%%%%%%%%%%%%%%%% IMPLEMENT MODEL-SPECIFIC FUNCTIONS HERE %%%%%%%%%%%%%%%%%% #
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
+
+
 class BenchmarkRunner(BaseBenchmarkRunner):
 
-    ACCURACY_METRIC_NAME = "mAP"
-
-    def before_benchmark(self, **kwargs):
-        self._output_name_map = (
-            # <tf.Tensor 'detection_boxes:0' shape=(8, None, None) dtype=float32>
-            (0, 'boxes'),
-            # <tf.Tensor 'detection_classes:0' shape=(8, None) dtype=float32>
-            (1, 'classes'),
-            # <tf.Tensor 'num_detections:0' shape=(8,) dtype=float32>
-            (2, 'num_detections'),
-            # <tf.Tensor 'detection_scores:0' shape=(8, None) dtype=float32>
-            (3, 'scores'),
+    def get_dataset_batches(self):
+        """Returns a list of batches of input samples.
+
+        Each batch should be in the form [x, y], where
+        x is a numpy array of the input samples for the batch, and
+        y is a numpy array of the expected model outputs for the batch
+
+        Returns:
+        - dataset: a TF Dataset object
+        - bypass_data_to_eval: any object type that will be passed unmodified to
+                            `evaluate_result()`. If not necessary: `None`
+
+        Note: script arguments can be accessed using `self._args.attr`
+        """
+
+        coco_api = COCO(annotation_file=self._args.annotation_path)
+        image_ids = coco_api.getImgIds()
+
+        image_paths = []
+        for image_id in image_ids:
+            coco_img = coco_api.imgs[image_id]
+            image_paths.append(
+                os.path.join(self._args.data_dir, coco_img['file_name'])
+            )
+
+        dataset = tf.data.Dataset.from_tensor_slices(image_paths)
+
+        def load_image_op(path):
+            image = tf.io.read_file(path)
+            image = tf.image.decode_jpeg(image, channels=3)
+
+            return tf.data.Dataset.from_tensor_slices([image])
+
+        dataset = dataset.interleave(
+            load_image_op,
+            cycle_length=tf.data.experimental.AUTOTUNE,
+            block_length=8,
+            num_parallel_calls=tf.data.experimental.AUTOTUNE
         )
 
-    def compute_accuracy_metric(self, predictions, expected, **kwargs):
-        return self._eval_model(
-            predictions=predictions,
-            image_ids=kwargs["image_ids"],
-            annotation_path=kwargs["annotation_path"]
+        def preprocess_fn(image):
+            if self._args.input_size is not None:
+                image = tf.image.resize(
+                    image,
+                    size=(self._args.input_size, self._args.input_size)
+                )
+                image = tf.cast(image, tf.uint8)
+            return image
+
+        dataset = dataset.map(
+            map_func=preprocess_fn,
+            num_parallel_calls=tf.data.experimental.AUTOTUNE,
         )
 
-    def _eval_model(self, predictions, image_ids, annotation_path):
+        dataset = dataset.batch(self._args.batch_size, drop_remainder=False)
+
+        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+
+        return dataset, None
+
+    def preprocess_model_inputs(self, data_batch):
+        """This function prepare the `data_batch` generated from the dataset.
+        Returns:
+            x: input of the model
+            y: data to be used for model evaluation
+
+        Note: script arguments can be accessed using `self._args.attr`
+        """
+
+        return data_batch, np.array([])
+
+    def postprocess_model_outputs(self, predictions, expected):
+        """Post process if needed the predictions and expected tensors. At the
+        minimum, this function transforms all TF Tensors into a numpy arrays.
+        Most models will not need to modify this function.
+
+        Note: script arguments can be accessed using `self._args.attr`
+        """
 
-        # for key in predictions:
-        #     predictions[key] = np.vstack(predictions[key])
-        #     if key == 'num_detections':
-        #         predictions[key] = predictions[key].ravel()
+        predictions = {k: t.numpy() for k, t in predictions.items()}
+
+        return predictions, expected
+
+    def evaluate_model(self, predictions, expected, bypass_data_to_eval):
+        """Evaluate result predictions for entire dataset.
+
+        This computes overall accuracy, mAP,  etc.  Returns the
+        metric value and a metric_units string naming the metric.
+
+        Note: script arguments can be accessed using `args.attr`
+        """
+        coco_api = COCO(annotation_file=self._args.annotation_path)
+        image_ids = coco_api.getImgIds()
 
-        coco = COCO(annotation_file=annotation_path)
         coco_detections = []
         for i, image_id in enumerate(image_ids):
-            coco_img = coco.imgs[image_id]
+            coco_img = coco_api.imgs[image_id]
             image_width = coco_img['width']
             image_height = coco_img['height']
 
@@ -127,163 +193,27 @@ def _eval_model(self, predictions, image_ids, annotation_path):
         coco_detections_path = os.path.join(tmp_dir, 'coco_detections.json')
         with open(coco_detections_path, 'w') as f:
             json.dump(coco_detections, f)
-        cocoDt = coco.loadRes(coco_detections_path)
+
+        cocoDt = coco_api.loadRes(coco_detections_path)
 
         shutil.rmtree(tmp_dir)
 
         # compute coco metrics
-        eval = COCOeval(coco, cocoDt, 'bbox')
+        eval = COCOeval(coco_api, cocoDt, 'bbox')
         eval.params.imgIds = image_ids
 
         eval.evaluate()
         eval.accumulate()
         eval.summarize()
 
-        return eval.stats[0]
-
-    def process_model_output(self, outputs, **kwargs):
-        # outputs = graph_func(batch_images)
-        if isinstance(outputs, dict):
-            outputs = {k: t.numpy() for k, t in outputs.items()}
-        else:
-            outputs = {
-                name: outputs[idx].numpy()
-                for idx, name in self._output_name_map
-            }
-
-        return outputs
-
-
-def get_dataset(
-    batch_size, images_dir, image_ids, input_size, use_synthetic_data
-):
-
-    image_paths = []
-
-    for image_id in image_ids:
-        coco_img = coco.imgs[image_id]
-        image_paths.append(os.path.join(images_dir, coco_img['file_name']))
-
-    dataset = tf.data.Dataset.from_tensor_slices(image_paths)
-
-    def load_image_op(path):
-        image = tf.io.read_file(path)
-        image = tf.image.decode_jpeg(image, channels=3)
-
-        return tf.data.Dataset.from_tensor_slices([image])
-
-    dataset = dataset.interleave(
-        lambda path: load_image_op(path),
-        cycle_length=tf.data.experimental.AUTOTUNE,
-        block_length=8,
-        num_parallel_calls=tf.data.experimental.AUTOTUNE
-    )
-
-    def preprocess_fn(image):
-        if input_size is not None:
-            image = tf.image.resize(image, size=(input_size, input_size))
-            image = tf.cast(image, tf.uint8)
-        return image
-
-    dataset = dataset.apply(
-        tf.data.experimental.map_and_batch(
-            map_func=preprocess_fn,
-            batch_size=batch_size,
-            num_parallel_calls=tf.data.experimental.AUTOTUNE,
-            drop_remainder=True
-        )
-    )
-
-    if use_synthetic_data:
-        dataset = dataset.take(count=1)  # loop over 1 batch
-        dataset = dataset.cache()
-        dataset = dataset.repeat()
-
-    dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
-
-    return dataset
+        return eval.stats[0] * 100, "mAP %"
 
 
 if __name__ == '__main__':
 
     cmdline_api = CommandLineAPI()
     args = cmdline_api.parse_args()
 
-    coco = COCO(annotation_file=args.annotation_path)
-    image_ids = coco.getImgIds()
-
-    def _input_fn(input_data_dir, build_steps, model_phase):
-
-        dataset = get_dataset(
-            batch_size=args.batch_size,
-            images_dir=input_data_dir,
-            image_ids=image_ids,
-            input_size=args.input_size,
-            # even when using synthetic data, we need to
-            # build and/or calibrate using real training data
-            # to be in a realistic scenario
-            use_synthetic_data=False,
-        )
+    runner = BenchmarkRunner(args)
 
-        for i, batch_images in enumerate(dataset):
-            if i >= build_steps:
-                break
-
-            print("* [%s] - step %04d/%04d" % (model_phase, i + 1, build_steps))
-            yield batch_images,
-
-    calibration_input_fn = partial(
-        _input_fn,
-        input_data_dir=args.calib_data_dir,
-        build_steps=args.num_calib_batches // args.batch_size,
-        model_phase="Calibration"
-    )
-
-    optimize_offline_input_fn = partial(
-        _input_fn,
-        input_data_dir=args.data_dir,
-        build_steps=1,
-        model_phase="Building"
-    )
-
-    runner = BenchmarkRunner(
-        input_saved_model_dir=args.input_saved_model_dir,
-        output_saved_model_dir=args.output_saved_model_dir,
-        allow_build_at_runtime=args.allow_build_at_runtime,
-        calibration_input_fn=calibration_input_fn,
-        debug=args.debug,
-        gpu_mem_cap=args.gpu_mem_cap,
-        input_signature_key=args.input_signature_key,
-        max_workspace_size_bytes=args.max_workspace_size,
-        minimum_segment_size=args.minimum_segment_size,
-        num_calib_batches=args.num_calib_batches,
-        optimize_offline=args.optimize_offline,
-        optimize_offline_input_fn=optimize_offline_input_fn,
-        output_tensor_names=args.output_tensor_names,
-        precision_mode=args.precision,
-        use_dynamic_shape=args.use_dynamic_shape,
-        use_tftrt=args.use_tftrt
-    )
-
-    get_benchmark_input_fn = partial(
-        get_dataset,
-        images_dir=args.data_dir,
-        image_ids=image_ids,
-        input_size=args.input_size
-    )
-
-    runner.execute_benchmark(
-        batch_size=args.batch_size,
-        display_every=args.display_every,
-        get_benchmark_input_fn=get_benchmark_input_fn,
-        num_iterations=args.num_iterations,
-        num_warmup_iterations=args.num_warmup_iterations,
-        skip_accuracy_testing=(
-            args.use_synthetic_data or args.skip_accuracy_testing
-        ),
-        use_synthetic_data=args.use_synthetic_data,
-        use_xla=args.use_xla,
-        ########### Additional Settings ############
-        image_ids=image_ids,
-        annotation_path=args.annotation_path
-    )
+    runner.execute_benchmark()