text: import horovod.tensorflow as hvd
  import os
  import tensorflow as tf
  from preprocessing import resnet_preprocessing, imagenet_preprocessing, darknet_preprocessing
  import functools
  
  def create_dataset(data_dir, batch_size, preprocessing='resnet', validation=False):
      filenames = [os.path.join(data_dir, i) for i in os.listdir(data_dir)]
      data = tf.data.TFRecordDataset(filenames).shard(hvd.size(), hvd.rank())
      if not validation:
          parse_fn = functools.partial(parse_train, preprocessing=preprocessing)
          data = data.shuffle(buffer_size=1000)
          data = data.map(parse_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      else:
          parse_fn = functools.partial(parse_validation, preprocessing=preprocessing)
          data = data.map(parse_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      # we drop remainder because we want same sized batches - XLA and because of allreduce being used to calculate
      # accuracy - validation accuracy may be slightly different than computing on all of validation data
      data = data.batch(batch_size, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)
      return data
  
  
  @tf.function
  def parse(record, is_training, preprocessing): 
      features = {'image/encoded': tf.io.FixedLenFeature((), tf.string),
                  'image/class/label': tf.io.FixedLenFeature((), tf.int64),
                  'image/object/bbox/xmin': tf.io.VarLenFeature(dtype=tf.float32),
                  'image/object/bbox/ymin': tf.io.VarLenFeature(dtype=tf.float32),
                  'image/object/bbox/xmax': tf.io.VarLenFeature(dtype=tf.float32),
                  'image/object/bbox/ymax': tf.io.VarLenFeature(dtype=tf.float32),
                  }
      parsed = tf.io.parse_single_example(record, features)
      image_bytes = tf.reshape(parsed['image/encoded'], shape=[])
      # bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4])
      bbox = tf.stack([parsed['image/object/bbox/%s' % x].values for x in ['ymin', 'xmin', 'ymax', 'xmax']])
      bbox = tf.transpose(tf.expand_dims(bbox, 0), [0, 2, 1])
      if preprocessing == 'resnet':
          augmenter = None # augment.AutoAugment()
          image = resnet_preprocessing.preprocess_image(image_bytes, bbox, 224, 224, 3, is_training=is_training)
      elif preprocessing == 'imagenet': # used by hrnet
          image = imagenet_preprocessing.preprocess_image(image_bytes, bbox, 224, 224, 3, is_training=is_training)
      elif preprocessing == 'darknet':
          image = darknet_preprocessing.preprocess_image(image_bytes, bbox, 256, 256, 3, is_training=is_training)
  
  
      label = tf.cast(parsed['image/class/label'] - 1, tf.int32)
      one_hot_label = tf.one_hot(label, depth=1000, dtype=tf.float32)
      return image, one_hot_label
  
  
  def parse_train(record, preprocessing):
      return parse(record, is_training=True, preprocessing=preprocessing)
  
  
  def parse_validation(record, preprocessing):
      return parse(record, is_training=False, preprocessing=preprocessing)
  
  