Input Datasets for CIFAR10

cifar10_dataset.py

def extract(example):
  data = tf.io.parse_example(
    example,
    # Schema of the example.
    {
      'image': tf.io.FixedLenFeature(shape=(32, 32, 3), dtype=tf.float32),
      'label': tf.io.FixedLenFeature(shape=(), dtype=tf.int64)
    }
  )
  return data['image'], data['label']

def get_dataset(filename):
  return tf.data.TFRecordDataset([GCS_PATH_FOR_DATA + filename]).
    map(extract, num_parallel_calls=tf.data.experimental.AUTOTUNE).
    shuffle(1024).
    batch(128).
    cache().
    prefetch(tf.data.experimental.AUTOTUNE)
 
train_dataset = get_dataset('train.tfrecord')
val_dataset = get_dataset('val.tfrecord')
test_dataset = get_dataset('test.tfrecord')