Move color distortion ops in SSD model from CPU to GPU.

haoyuz · tensorflower-gardener · commit 119534b5ce7c · 2019-02-05T15:57:34.000-08:00
PiperOrigin-RevId: 229664805
diff --git a/scripts/tf_cnn_benchmarks/models/model.py b/scripts/tf_cnn_benchmarks/models/model.py
@@ -257,6 +257,10 @@ def get_synthetic_inputs(self, input_name, nclass):
         name=self.model_name + '_synthetic_labels')
     return (inputs, labels)
 
+  def gpu_preprocess_nhwc(self, images, phase_train=True):
+    del phase_train
+    return images
+
   def build_network(self,
                     inputs,
                     phase_train=True,
@@ -273,6 +277,7 @@ def build_network(self,
         information.
     """
     images = inputs[0]
+    images = self.gpu_preprocess_nhwc(images, phase_train)
     if self.data_format == 'NCHW':
       images = tf.transpose(images, [0, 3, 1, 2])
     var_type = tf.float32
diff --git a/scripts/tf_cnn_benchmarks/models/ssd_model.py b/scripts/tf_cnn_benchmarks/models/ssd_model.py
@@ -116,6 +116,26 @@ def __init__(self, label_num=ssd_constants.NUM_CLASSES, batch_size=32,
   def skip_final_affine_layer(self):
     return True
 
+  def gpu_preprocess_nhwc(self, images, phase_train=True):
+    try:
+      import ssd_dataloader  # pylint: disable=g-import-not-at-top
+    except ImportError:
+      raise ImportError('To use the COCO dataset, you must clone the '
+                        'repo https://github.com/tensorflow/models and add '
+                        'tensorflow/models and tensorflow/models/research to '
+                        'the PYTHONPATH, and compile the protobufs by '
+                        'following https://github.com/tensorflow/models/blob/'
+                        'master/research/object_detection/g3doc/installation.md'
+                        '#protobuf-compilation ; To evaluate using COCO'
+                        'metric, download and install Python COCO API from'
+                        'https://github.com/cocodataset/cocoapi')
+
+    if phase_train:
+      images = ssd_dataloader.color_jitter(
+          images, brightness=0.125, contrast=0.5, saturation=0.5, hue=0.05)
+      images = ssd_dataloader.normalize_image(images)
+    return images
+
   def add_backbone_model(self, cnn):
     # --------------------------------------------------------------------------
     # Resnet-34 backbone model -- modified for SSD
@@ -152,8 +172,8 @@ def add_backbone_model(self, cnn):
       resnet_model.residual_block(cnn, 256, stride, version, i == 0)
 
     # ResNet-34 block group 4: removed final block group
-    # The following 3 lines are intentially commented out to differentiate from
-    # the original ResNet-34 model
+    # The following 3 lines are intentionally commented out to differentiate
+    # from the original ResNet-34 model
     # for i in range(resnet34_layers[3]):
     #   stride = 2 if i == 0 else 1
     #   resnet_model.residual_block(cnn, 512, stride, version, i == 0)
diff --git a/scripts/tf_cnn_benchmarks/models/trivial_model.py b/scripts/tf_cnn_benchmarks/models/trivial_model.py
@@ -53,15 +53,18 @@ class TrivialSSD300Model(model.CNNModel):
 
   def __init__(self, params=None):
     super(TrivialSSD300Model, self).__init__(
-        'trivial', 300, 32, 0.005, params=params)
+        'trivial', 300, params.batch_size, 0.005, params=params)
 
   def add_inference(self, cnn):
     cnn.reshape([-1, 300 * 300 * 3])
     cnn.affine(1)
     cnn.affine(4096)
 
   def get_input_shapes(self, subset):
-    return [[32, 300, 300, 3], [32, 8732, 4], [32, 8732, 1], [32]]
+    return [[self.batch_size, 300, 300, 3],
+            [self.batch_size, 8732, 4],
+            [self.batch_size, 8732, 1],
+            [self.batch_size]]
 
   def loss_function(self, inputs, build_network_result):
     images, _, _, labels = inputs
diff --git a/scripts/tf_cnn_benchmarks/preprocessing.py b/scripts/tf_cnn_benchmarks/preprocessing.py
@@ -971,9 +971,6 @@ def preprocess(self, data):
       # See https://github.com/tensorflow/models/blob/master/research/object_detection/core/preprocessor.py  # pylint: disable=line-too-long
       mlperf.logger.log(key=mlperf.tags.RANDOM_FLIP_PROBABILITY, value=0.5)
 
-      image = ssd_dataloader.color_jitter(
-          image, brightness=0.125, contrast=0.5, saturation=0.5, hue=0.05)
-      image = ssd_dataloader.normalize_image(image)
       image = tf.cast(image, self.dtype)
 
       encoded_returns = ssd_encoder.encode_labels(boxes, classes)
diff --git a/scripts/tf_cnn_benchmarks/ssd_dataloader.py b/scripts/tf_cnn_benchmarks/ssd_dataloader.py
@@ -356,22 +356,24 @@ def color_jitter(image, brightness=0, contrast=0, saturation=0, hue=0):
     return image
 
 
-def normalize_image(image):
-  """Normalize the image to zero mean and unit variance.
+def normalize_image(images):
+  """Normalize image to zero mean and unit variance.
 
   Args:
-    image: 3D tensor of type float32, value in [0, 1]
+    images: a tensor representing images, at least 3-D.
   Returns:
-    image normalized by mean and stdev.
+    images normalized by mean and stdev.
   """
-  image = tf.subtract(image, ssd_constants.NORMALIZATION_MEAN)
-  image = tf.divide(image, ssd_constants.NORMALIZATION_STD)
+  data_type = images.dtype
+  mean = tf.constant(ssd_constants.NORMALIZATION_MEAN, data_type)
+  std = tf.constant(ssd_constants.NORMALIZATION_STD, data_type)
+  images = tf.divide(tf.subtract(images, mean), std)
 
   mlperf.logger.log(key=mlperf.tags.DATA_NORMALIZATION_MEAN,
                     value=ssd_constants.NORMALIZATION_MEAN)
   mlperf.logger.log(key=mlperf.tags.DATA_NORMALIZATION_STD,
                     value=ssd_constants.NORMALIZATION_STD)
-  return image
+  return images
 
 
 class Encoder(object):
