code: import os
  import random
  import numpy as np
  import pandas as pd
  import tensorflow as tf
  
  from augment import Augment
  
  
  AUTO = tf.data.experimental.AUTOTUNE
  
  
  def set_dataset(task, data_path):
      trainset = pd.read_csv(
          os.path.join(
              data_path, 'imagenet_trainset.csv'
          )).values.tolist()
      trainset = [[os.path.join(data_path, t[0]), t[1]] for t in trainset]
  
      if task == 'lincls':
          valset = pd.read_csv(
              os.path.join(
                  data_path, 'imagenet_valset.csv'
              )).values.tolist()
          valset = [[os.path.join(data_path, t[0]), t[1]] for t in valset]
          return np.array(trainset, dtype='object'), np.array(valset, dtype='object')
  
      return np.array(trainset, dtype='object')
  
  
  class DataLoader:
      def __init__(self, args, mode, datalist, batch_size, num_workers=1, shuffle=True):
          self.args = args
          self.mode = mode
          self.datalist = datalist
          self.batch_size = batch_size
          self.num_workers = num_workers
          self.shuffle = shuffle
  
          self.dataloader = self._dataloader()
  
      def __len__(self):
          return len(self.datalist)
  
      def fetch_dataset(self, path, y=None):
          x = tf.io.read_file(path)
          if y is not None:
              return tf.data.Dataset.from_tensors((x, y))
          return tf.data.Dataset.from_tensors(x)
  
      def augmentation(self, img, shape):
          augset = Augment(self.args, self.mode)
          if self.args.task in ['v1', 'v2']:
              img_list = []
              for _ in range(2): # query, key
                  aug_img = tf.identity(img)
                  if self.args.task == 'v1':
                      aug_img = augset._augmentv1(aug_img, shape) # moco v1
                  else:
                      radius = np.random.choice([3, 5])
                      aug_img = augset._augmentv2(aug_img, shape, (radius, radius)) # moco v2
                  img_list.append(aug_img)
              return img_list
          else:
              return augset._augment_lincls(img, shape)
  
      def dataset_parser(self, value, label=None):
          shape = tf.image.extract_jpeg_shape(value)
          img = tf.io.decode_jpeg(value, channels=3)
          if label is None:
              # moco
              query, key = self.augmentation(img, shape)
              inputs = {'query': query, 'key': key}
              labels = tf.zeros([])
          else:
              # lincls
              inputs = self.augmentation(img, shape)
              labels = tf.one_hot(label, self.args.classes)
          return (inputs, labels)
  
      def shuffle_BN(self, value, labels):
          if self.num_workers > 1:
              pre_shuffle = [(i, value['key'][i]) for i in range(self.batch_size)]
              random.shuffle(pre_shuffle)
              shuffle_idx = []
              value_temp = []
              for vv in pre_shuffle:
                  shuffle_idx.append(vv[0])
                  value_temp.append(tf.expand_dims(vv[1], axis=0))
              value['key'] = tf.concat(value_temp, axis=0)
              unshuffle_idx = np.array(shuffle_idx).argsort().tolist()
              value.update({'unshuffle': unshuffle_idx})
          return (value, labels)
          
      def _dataloader(self):
          self.imglist = self.datalist[:,0].tolist()
          if self.args.task in ['v1', 'v2']:
              dataset = tf.data.Dataset.from_tensor_slices(self.imglist)
          else:
              self.labellist = self.datalist[:,1].tolist()
              dataset = tf.data.Dataset.from_tensor_slices((self.imglist, self.labellist))
  
          dataset = dataset.repeat()
          if self.shuffle:
              dataset = dataset.shuffle(len(self.datalist))
  
          dataset = dataset.interleave(self.fetch_dataset, num_parallel_calls=AUTO)
          dataset = dataset.map(self.dataset_parser, num_parallel_calls=AUTO)
          dataset = dataset.batch(self.batch_size)
          dataset = dataset.prefetch(AUTO)
          if self.args.shuffle_bn and self.args.task in ['v1', 'v2']:
              # only moco
              dataset = dataset.map(self.shuffle_BN, num_parallel_calls=AUTO)
          return dataset
apis: tensorflow.io.decode_jpeg tensorflow.zeros tensorflow.identity tensorflow.image.extract_jpeg_shape tensorflow.expand_dims numpy.random.choice tensorflow.data.Dataset.from_tensors tensorflow.one_hot tensorflow.io.read_file tensorflow.concat numpy.array tensorflow.data.Dataset.from_tensor_slices