Handle data format and graph compile properly for better GPU performa…

…nce (tensorflow#6013)

* Handle data format in Keras ResNet model properly for better performance on GPU; Compile only the training graph when skip_eval flag is True

* Added data format fix to Keras Cifar model; Removed unnecessary import

* Add a comment to the skip_eval flag per Priya's request

official/resnet/keras/keras_cifar_main.py

@@ -108,6 +108,8 @@ def run(flags_obj):
  per_device_batch_size = distribution_utils.per_device_batch_size(
  flags_obj.batch_size, flags_core.get_num_gpus(flags_obj))
 
+ tf.keras.backend.set_image_data_format(flags_obj.data_format)
+
  if flags_obj.use_synthetic_data:
  input_fn = keras_common.get_synth_input_fn(
  height=cifar_main.HEIGHT,
@@ -160,6 +162,7 @@ def run(flags_obj):
 
  validation_data = eval_input_dataset
  if flags_obj.skip_eval:
+ tf.keras.backend.set_learning_phase(1)
  num_eval_steps = None
  validation_data = None
 

official/resnet/keras/keras_imagenet_main.py

@@ -95,6 +95,8 @@ def run(flags_obj):
  raise ValueError('dtype fp16 is not supported in Keras. Use the default '
  'value(fp32).')
 
+ tf.keras.backend.set_image_data_format(flags_obj.data_format)
+
  per_device_batch_size = distribution_utils.per_device_batch_size(
  flags_obj.batch_size, flags_core.get_num_gpus(flags_obj))
 
@@ -149,6 +151,10 @@ def run(flags_obj):
 
  validation_data = eval_input_dataset
  if flags_obj.skip_eval:
+ # Only build the training graph. This reduces memory usage introduced by
+ # control flow ops in layers that have different implementations for
+ # training and inference (e.g., batch norm).
+ tf.keras.backend.set_learning_phase(1)
  num_eval_steps = None
  validation_data = None
 

official/resnet/keras/resnet_cifar_model.py

@@ -187,16 +187,18 @@ def resnet56(classes=100, training=None):
  Returns:
  A Keras model instance.
  """
- # Determine proper input shape
+ input_shape = (32, 32, 3)
+ img_input = layers.Input(shape=input_shape)
+
  if backend.image_data_format() == 'channels_first':
- input_shape = (3, 32, 32)
+ x = layers.Lambda(lambda x: backend.permute_dimensions(x, (0, 3, 1, 2)),
+ name='transpose')(img_input)
  bn_axis = 1
  else: # channel_last
- input_shape = (32, 32, 3)
+ x = img_input
  bn_axis = 3
 
- img_input = layers.Input(shape=input_shape)
- x = tf.keras.layers.ZeroPadding2D(padding=(1, 1), name='conv1_pad')(img_input)
+ x = tf.keras.layers.ZeroPadding2D(padding=(1, 1), name='conv1_pad')(x)
  x = tf.keras.layers.Conv2D(16, (3, 3),
  strides=(1, 1),
  padding='valid',

official/resnet/keras/resnet_model.py

@@ -190,16 +190,18 @@ def resnet50(num_classes):
  Returns:
  A Keras model instance.
  """
- # Determine proper input shape
+ input_shape = (224, 224, 3)
+ img_input = layers.Input(shape=input_shape)
+
  if backend.image_data_format() == 'channels_first':
- input_shape = (3, 224, 224)
+ x = layers.Lambda(lambda x: backend.permute_dimensions(x, (0, 3, 1, 2)),
+ name='transpose')(img_input)
  bn_axis = 1
- else:
- input_shape = (224, 224, 3)
+ else: # channels_last
+ x = img_input
  bn_axis = 3
 
- img_input = layers.Input(shape=input_shape)
- x = layers.ZeroPadding2D(padding=(3, 3), name='conv1_pad')(img_input)
+ x = layers.ZeroPadding2D(padding=(3, 3), name='conv1_pad')(x)
  x = layers.Conv2D(64, (7, 7),
  strides=(2, 2),
  padding='valid',