parameters: (data_dir,
                            split,
                            data_aug=True,
                            batch_size=100,
                            epochs=None,
                            shuffle=True,
                            data_format="channels_first",
                            num_parallel_calls=12,
                            prefetch=0,
                            div255=True,
                            dtype=tf.float32)
return_statement: return dataset
docstring: Returns a tf.train.Dataset object from reading tfrecords.
  
    Args:
        data_dir: Directory of tfrecords
        split: "train", "validation", or "test"
        data_aug: Apply data augmentation if True
        batch_size: Batch size of dataset object
        epochs: Number of epochs to repeat the dataset; default `None` means
            repeating indefinitely
        shuffle: Shuffle the dataset if True
        data_format: `channels_first` or `channels_last`
        num_parallel_calls: Number of threads for dataset preprocess
        prefetch: Buffer size for prefetch
        div255: Divide the images by 255 if True
        dtype: Data type of images
    Returns:
        A tf.train.Dataset object
  
    Raises:
        ValueError: Unknown split
docstring_summary: Returns a tf.train.Dataset object from reading tfrecords.
docstring_tokens: Returns a tf . train . Dataset object from reading tfrecords .
function: def get_ds_from_tfrecords(data_dir,
                            split,
                            data_aug=True,
                            batch_size=100,
                            epochs=None,
                            shuffle=True,
                            data_format="channels_first",
                            num_parallel_calls=12,
                            prefetch=0,
                            div255=True,
                            dtype=tf.float32):
    """Returns a tf.train.Dataset object from reading tfrecords.
  
    Args:
        data_dir: Directory of tfrecords
        split: "train", "validation", or "test"
        data_aug: Apply data augmentation if True
        batch_size: Batch size of dataset object
        epochs: Number of epochs to repeat the dataset; default `None` means
            repeating indefinitely
        shuffle: Shuffle the dataset if True
        data_format: `channels_first` or `channels_last`
        num_parallel_calls: Number of threads for dataset preprocess
        prefetch: Buffer size for prefetch
        div255: Divide the images by 255 if True
        dtype: Data type of images
    Returns:
        A tf.train.Dataset object
  
    Raises:
        ValueError: Unknown split
    """
  
    if split not in ["train", "validation", "test", "train_all"]:
      raise ValueError("Unknown split {}".format(split))
  
    def _parser(serialized_example):
      """Parses a single tf.Example into image and label tensors."""
      features = tf.parse_single_example(
          serialized_example,
          features={
              "image": tf.FixedLenFeature([], tf.string),
              "label": tf.FixedLenFeature([], tf.int64),
          })
      image = tf.decode_raw(features["image"], tf.uint8)
      # Initially reshaping to [H, W, C] does not work
      image = tf.reshape(image, [NUM_CHANNEL, IMAGE_HEIGHT, IMAGE_WIDTH])
      # This is needed for `tf.image.resize_image_with_crop_or_pad`
      image = tf.transpose(image, [1, 2, 0])
  
      image = tf.cast(image, dtype)
      label = tf.cast(features["label"], tf.int32)
  
      if data_aug:
        image = tf.image.resize_image_with_crop_or_pad(image, IMAGE_HEIGHT + 4,
                                                       IMAGE_WIDTH + 4)
        image = tf.random_crop(image, [IMAGE_HEIGHT, IMAGE_WIDTH, NUM_CHANNEL])
        image = tf.image.random_flip_left_right(image)
  
      if data_format == "channels_first":
        image = tf.transpose(image, [2, 0, 1])
  
      if div255:
        image /= 255.
  
      return image, label
  
    filename = os.path.join(data_dir, split + ".tfrecords")
    dataset = tf.data.TFRecordDataset(filename)
    dataset = dataset.repeat(epochs)
    dataset = dataset.map(_parser, num_parallel_calls=num_parallel_calls)
    dataset = dataset.prefetch(prefetch)
  
    if shuffle:
      # Find the right size according to the split
      size = {
          "train": 40000,
          "validation": 10000,
          "test": 10000,
          "train_all": 50000
      }[split]
      dataset = dataset.shuffle(size)
  
    dataset = dataset.batch(batch_size, drop_remainder=True)
  
    return dataset
function_tokens: def get_ds_from_tfrecords ( data_dir , split , data_aug = True , batch_size = 100 , epochs = None , shuffle = True , data_format = "channels_first" , num_parallel_calls = 12 , prefetch = 0 , div255 = True , dtype = tf . float32 ) : if split not in [ "train" , "validation" , "test" , "train_all" ] : raise ValueError ( "Unknown split {}" . format ( split ) ) def _parser ( serialized_example ) : """Parses a single tf.Example into image and label tensors.""" features = tf . parse_single_example ( serialized_example , features = { "image" : tf . FixedLenFeature ( [ ] , tf . string ) , "label" : tf . FixedLenFeature ( [ ] , tf . int64 ) , } ) image = tf . decode_raw ( features [ "image" ] , tf . uint8 ) # Initially reshaping to [H, W, C] does not work image = tf . reshape ( image , [ NUM_CHANNEL , IMAGE_HEIGHT , IMAGE_WIDTH ] ) # This is needed for `tf.image.resize_image_with_crop_or_pad` image = tf . transpose ( image , [ 1 , 2 , 0 ] ) image = tf . cast ( image , dtype ) label = tf . cast ( features [ "label" ] , tf . int32 ) if data_aug : image = tf . image . resize_image_with_crop_or_pad ( image , IMAGE_HEIGHT + 4 , IMAGE_WIDTH + 4 ) image = tf . random_crop ( image , [ IMAGE_HEIGHT , IMAGE_WIDTH , NUM_CHANNEL ] ) image = tf . image . random_flip_left_right ( image ) if data_format == "channels_first" : image = tf . transpose ( image , [ 2 , 0 , 1 ] ) if div255 : image /= 255. return image , label filename = os . path . join ( data_dir , split + ".tfrecords" ) dataset = tf . data . TFRecordDataset ( filename ) dataset = dataset . repeat ( epochs ) dataset = dataset . map ( _parser , num_parallel_calls = num_parallel_calls ) dataset = dataset . prefetch ( prefetch ) if shuffle : # Find the right size according to the split size = { "train" : 40000 , "validation" : 10000 , "test" : 10000 , "train_all" : 50000 } [ split ] dataset = dataset . shuffle ( size ) dataset = dataset . batch ( batch_size , drop_remainder = True ) return dataset