Adding support for NGC resnet50 v1.5

Author: DEKHTIARJonathan

tftrt/examples/image_classification/image_classification.py

@@ -55,11 +55,13 @@ def __init__(self):
                                        'the model')
 
         self._parser.add_argument('--preprocess_method', type=str,
-                                  choices=['vgg', 'inception'], default='vgg',
+                                  choices=['vgg', 'inception',
+                                           'resnet50_v1_5_tf1_ngc_preprocess'
+                                  ],
+                                  default='vgg',
                                   help='The image preprocessing method used in '
                                        'dataloading.')
 
-
 class BenchmarkRunner(BaseBenchmarkRunner):
 
     ACCURACY_METRIC_NAME = "accuracy"
@@ -107,6 +109,8 @@ def get_preprocess_fn(preprocess_method, input_size):
             preprocess_fn = preprocessing.vgg_preprocess
         elif preprocess_method == 'inception':
             preprocess_fn = preprocessing.inception_preprocess
+        elif preprocess_method == 'resnet50_v1_5_tf1_ngc_preprocess':
+            preprocess_fn = preprocessing.resnet50_v1_5_tf1_ngc_preprocess
         else:
             raise ValueError(
                 'Invalid preprocessing method {}'.format(preprocess_method)

tftrt/examples/image_classification/preprocessing.py

@@ -188,4 +188,43 @@ def vgg_preprocess(image, output_height, output_width):
 	image = _central_crop([image], output_height, output_width)[0]
 	image.set_shape([output_height, output_width, 3])
 	image = tf.cast(image, tf.float32)
-	return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN]) 
+	return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
+
+
+def resnet50_v1_5_tf1_ngc_preprocess(image,
+                  height,
+                  width,
+                  central_fraction=0.875,
+                  scope=None,
+                  central_crop=False):
+  """Prepare one image for evaluation.
+  If height and width are specified it would output an image with that size by
+  applying resize_bilinear.
+  If central_fraction is specified it would crop the central fraction of the
+  input image.
+  Args:
+  image: 3-D Tensor of image. If dtype is tf.float32 then the range should be
+    [0, 1], otherwise it would converted to tf.float32 assuming that the range
+    is [0, MAX], where MAX is largest positive representable number for
+    int(8/16/32) data type (see `tf.image.convert_image_dtype` for details).
+  height: integer
+  width: integer
+  central_fraction: Optional Float, fraction of the image to crop.
+  scope: Optional scope for name_scope.
+  central_crop: Enable central cropping of images during preprocessing for
+    evaluation.
+  Returns:
+  3-D float Tensor of prepared image.
+  """
+  if image.dtype != tf.float32:
+      image = tf.image.convert_image_dtype(image, dtype=tf.float32)
+      # Crop the central region of the image with an area containing 87.5% of
+      # the original image.
+  image = tf.image.central_crop(image, central_fraction=central_fraction)
+  if height and width:
+      # Resize the image to the specified height and width.
+      image = tf.expand_dims(image, 0)
+      image = tf.image.resize(image, [height, width])
+      image = tf.squeeze(image, [0])
+  image = image * 255
+  return image

tftrt/examples/image_classification/scripts/base_script.sh

@@ -75,6 +75,13 @@ case ${MODEL_NAME} in
     NUM_CLASSES=1000
     ;;
 
+  "resnet50-v1.5_tf1_ngc" )
+    NUM_CLASSES=1000
+    OUTPUT_TENSOR_IDX_FLAG="--output_tensor_indices=0"
+    OUTPUT_TENSOR_NAME_FLAG="--output_tensor_names=classes"
+    PREPROCESS_METHOD="resnet50_v1_5_tf1_ngc_preprocess"
+    ;;
+
   "resnet50v2_backbone" | "resnet50v2_sparse_backbone" )
     INPUT_SIZE=256
     ;;

tftrt/examples/image_classification/scripts/resnet50-v1.5_tf1_ngc.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
+
+bash ${SCRIPT_DIR}/base_script.sh --model_name="resnet50-v1.5_tf1_ngc" ${@}