title: TF Dataset from Keras Sequence Class
body: <p>I thought I would share something that took me a while to figure out: easily wrapping an existing Keras Sequence Class with a TF Dataset object. After following tutorials and migrating from TF 1.X and Keras to TF 2.X I finally figured out how to do it with minimal code. Hopefully I'm not the only one who struggled with this and others will find this helpful :)</p>
  
  <p>A few assumptions:</p>
  
  <ul>
  <li><p>Sequence class loads data and labels</p></li>
  <li><p>Labels have the same shape (apart from channels) as the source data (i.e. this is something I use for training U-Nets)</p></li>
  <li><p>Data format is channels last</p></li>
  </ul>
  
  <pre class="lang-py prettyprint-override"><code>import tensorflow as tf
  
  def DatasetFromSequenceClass(sequenceClass, stepsPerEpoch, nEpochs, batchSize, dims=[512,512,3], n_classes=2, data_type=tf.float32, label_type=tf.float32):
      # eager execution wrapper
      def DatasetFromSequenceClassEagerContext(func):
          def DatasetFromSequenceClassEagerContextWrapper(batchIndexTensor):
              # Use a tf.py_function to prevent auto-graph from compiling the method
              tensors = tf.py_function(
                  func,
                  inp=[batchIndexTensor],
                  Tout=[data_type, label_type]
              )
  
              # set the shape of the tensors - assuming channels last
              tensors[0].set_shape([batchSize, dims[0], dims[1], dims[2]])   # [samples, height, width, nChannels]
              tensors[1].set_shape([batchSize, dims[0], dims[1], n_classes]) # [samples, height, width, nClasses for one hot]
              return tensors
          return DatasetFromSequenceClassEagerContextWrapper
  
      # TF dataset wrapper that indexes our sequence class
      @DatasetFromSequenceClassEagerContext
      def LoadBatchFromSequenceClass(batchIndexTensor):
          # get our index as numpy value - we can use .numpy() because we have wrapped our function
          batchIndex = batchIndexTensor.numpy()
  
          # zero-based index for what batch of data to load; i.e. goes to 0 at stepsPerEpoch and starts cound over
          zeroBatch = batchIndex % stepsPerEpoch
  
          # load data
          data, labels = sequenceClass[zeroBatch]
  
          # convert to tensors and return
          return tf.convert_to_tensor(data), tf.convert_to_tensor(labels)
  
      # create our data set for how many total steps of training we have
      dataset = tf.data.Dataset.range(stepsPerEpoch*nEpochs)
  
      # return dataset using map to load our batches of data, use TF to specify number of parallel calls
      return dataset.map(LoadBatchFromSequenceClass, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  </code></pre>
  
  <p>With that function, you can then update your training to look something like this:</p>
  
  <pre class="lang-py prettyprint-override"><code># load our data as tensorflow datasets
  training = DatasetFromSequenceClass(trainingSequence, training_steps, nEpochs, batchSize, dims=shp, n_classes=nClasses)
  validation = DatasetFromSequenceClass(validationSequence, validation_steps, nEpochs, batchSize, dims=shp, n_classes=nClasses)
  
  # train
  model_object.fit(training,
      steps_per_epoch=training_steps,
      validation_data=validation,
      validation_steps=validation_steps,
      epochs=nEpochs,
      callbacks=callbacks,
      verbose=1)
  </code></pre>
  
  <p>From here there are lots of other options for the Dataset API (like prefetch), but this should be a good starting point.</p>