title: How to reduce tensorflow dataset input pipeline host device (cpu) time (currently ~40%)?
body: <p>I am trying to replicate the resnet18 paper. Before running this on the full Image Net dataset on disk, I'm doing some evaluation runs with the publicly available <code>imagenette/320px</code> dataset from TFDS (much much smaller subset of imagenet with 10 classes, already in <code>.tfrecord</code> format._</p>
  <p><strong>Note</strong>: the full notebook to do training and tracing is available here: <a href="https://colab.research.google.com/gist/ntakouris/1a17cf3148a3d775fde85cac2bfd52a4/resnet18_baseline.ipynb" rel="nofollow noreferrer">resnet18_baseline.ipynb</a> Just switch to a <code>GPU</code> runtime and run all the cells. It's already set-up with tensorboard profiling on the second batch. (You can use TPU as well, but some <code>keras.layers.experimental.preprocessing</code> layers do not support TPU ops yet and you have to enable soft device placement. Please use a GPU).</p>
  <h2>Input Operations</h2>
  <ol>
  <li>read images from the input dataset. These images usually have got different dimensions and we need some crop function because input tensors can <strong>not</strong> have different dimensions for batching. Therefore, for training I use random crop and for testing/validation datasets a center crop.</li>
  </ol>
  <pre class="lang-py prettyprint-override"><code>random_crop_layer = keras.layers.experimental.preprocessing.RandomCrop(224, 224)
  center_crop_layer = keras.layers.experimental.preprocessing.CenterCrop(224, 224)
  
  @tf.function(experimental_relax_shapes=True) # avoid retracing
  def train_crop_fn(x, y):
    return random_crop_layer(x), y
  
  @tf.function(experimental_relax_shapes=True)
  def eval_crop_fn(x, y):
    return center_crop_layer(x), y
  </code></pre>
  <ol start="2">
  <li>Perform some simple preprocessing/augmentations to the input data. These include rescaling to 0-1 and also scaling based on mean and stdev of the rgb colours on imagenet. Also, random</li>
  </ol>
  <pre class="lang-py prettyprint-override"><code>rescaling_layer = keras.layers.experimental.preprocessing.Rescaling(1./255)
  
  train_preproc = keras.Sequential([
    rescaling_layer
  ])
  
  # from https://github.com/tensorflow/models/blob/master/official/vision/image_classification/preprocessing.py
  # Calculated from the ImageNet training set
  MEAN_RGB = (0.485 , 0.456, 0.406)
  STDDEV_RGB = (0.229, 0.224, 0.225)
  
  @tf.function
  def z_score_scale(x):
    return (x - MEAN_RGB) / STDDEV_RGB
  
  @tf.function
  def train_preproc_fn(x, y):
    return z_score_scale(train_preproc(x)), y
  
  @tf.function
  def eval_preproc_fn(x, y):
    return z_score_scale(eval_preproc(x)), y
  </code></pre>
  <h2>Input Pipeline</h2>
  <pre class="lang-py prettyprint-override"><code>def get_input_pipeline(input_ds, bs, crop_fn, augmentation_fn):
    ret_ds = (
        input_ds
        .batch(1) # pre-crop are different dimensions and can't be batched
        .map(crop_fn, 
            num_parallel_calls=tf.data.experimental.AUTOTUNE)
        .unbatch()
        .batch(bs)
        .map(augmentation_fn, # augmentations can be batched though.
            num_parallel_calls=tf.data.experimental.AUTOTUNE)
    )
  
    return ret_ds
  
  # dataset loading
  def load_imagenette():
    train_ds, ds_info = tfds.load('imagenette/320px', split='train', as_supervised=True, with_info=True)
    valid_ds = tfds.load('imagenette/320px', split='validation', as_supervised=True)
  
    return train_ds, valid_ds, valid_ds, ds_info.features['label'].num_classes
  
  # pipeline construction
  train_ds, valid_ds, test_ds, num_classes = load_imagenette()
  
  # datasets used for training (notice that I use prefetch here)
  train_samples = get_input_pipeline(train_ds, BS, train_crop_fn, train_preproc_fn).prefetch(tf.data.experimental.AUTOTUNE)
  valid_samples = get_input_pipeline(valid_ds, BS, eval_crop_fn, eval_preproc_fn).prefetch(tf.data.experimental.AUTOTUNE)
  test_samples = get_input_pipeline(test_ds, BS, eval_crop_fn, eval_preproc_fn).prefetch(tf.data.experimental.AUTOTUNE)
  
  </code></pre>
  <h2>Training and Profiling</h2>
  <p>I use tensorboard profiler to check the second batch size and I get a warning that this is highly input bound, with about 40% of processing wasted on inputs.</p>
  <p>For a classic resnet18 model, you can drive the batch size up to 768 without getting a OOM error, which is what I use. A single step with bs 256 takes about 2-3 seconds.</p>
  <p>I also get a warning that <code>on_train_batch_size_end</code> is slow, at around ~1.5 seconds, compared to the 1s batch time.</p>
  <p>The model training code is very simple keras:</p>
  <pre class="lang-py prettyprint-override"><code>model.fit(
          train_samples,
          validation_data=valid_samples,
          epochs=100,
          batch_size=BS,
          use_multiprocessing=True
          callbacks=[tensorboard_callback, model_checkpoint_callback, early_stop_callback, reduce_lr_callback]
  )
  </code></pre>
  <p>and the callbacks are specified as:</p>
  <pre class="lang-py prettyprint-override"><code>log_dir = os.path.join(os.getcwd(), 'logs')
  tensorboard_callback = TensorBoard(log_dir=log_dir, update_freq=&quot;epoch&quot;, profile_batch=2)
  
  reduce_lr_callback = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=5, min_lr=0.001, verbose=1)
  
  model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath='model.{epoch:02d}-{val_loss:.4f}.h5',
                                                monitor='val_loss',
                                                verbose=1,
                                                save_best_only=True)
  
  early_stop_callback = keras.callbacks.EarlyStopping(monitor='val_loss', patience=15)
  
  </code></pre>
  <p>Lastly, here are some sample tensorboard profiling screenshots. I can't figure out how to make this run faster:</p>
  <p><a href="https://i.stack.imgur.com/svfiJ.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/svfiJ.png" alt="input pipeline analysis" /></a></p>
  <p><a href="https://i.stack.imgur.com/lnWNH.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/lnWNH.png" alt="tensorboard profiling tip" /></a></p>