text: # --------------------------------------------------------
  # Swin Transformer
  # Copyright (c) 2021 Microsoft
  # Licensed under The MIT License [see LICENSE for details]
  # Written by Ze Liu
  # --------------------------------------------------------
  # Copyright (c) 2021-2022, NVIDIA CORPORATION.  All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import os
  import torch
  import numpy as np
  import torch.distributed as dist
  from torchvision import datasets, transforms
  from timm.data import create_transform
  
  try:
      from torchvision.transforms import InterpolationMode
  
  
      def _pil_interp(method):
          if method == 'bicubic':
              return InterpolationMode.BICUBIC
          elif method == 'lanczos':
              return InterpolationMode.LANCZOS
          elif method == 'hamming':
              return InterpolationMode.HAMMING
          else:
              # default bilinear, do we want to allow nearest?
              return InterpolationMode.BILINEAR
  except:
      from timm.data.transforms import _pil_interp
  
  def build_val_loader(config):
      config.freeze()
      dataset_val, _ = build_dataset(is_train=False, config=config)
      # print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build val dataset")
  
      indices = np.arange(0, len(dataset_val), 1)
      if config.TEST.SEQUENTIAL:
          sampler_val = torch.utils.data.SequentialSampler(dataset_val)
      else:
          sampler_val = torch.utils.data.distributed.DistributedSampler(
              dataset_val, shuffle=False
          )
  
      data_loader_val = torch.utils.data.DataLoader(
          dataset_val, sampler=sampler_val,
          batch_size=config.DATA.BATCH_SIZE,
          shuffle=False,
          num_workers=config.DATA.NUM_WORKERS,
          pin_memory=config.DATA.PIN_MEMORY,
          drop_last=False
      )
      
      return dataset_val, data_loader_val
  
  def build_loader(config, args):
      config.defrost()
      dataset_train, config.MODEL.NUM_CLASSES = build_dataset(is_train=True, config=config)
      config.freeze()
      print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build train dataset")
      dataset_val, _ = build_dataset(is_train=False, config=config)
      print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()} successfully build val dataset")
  
      num_tasks = dist.get_world_size()
      global_rank = dist.get_rank()
      
      sampler_train = torch.utils.data.DistributedSampler(
          dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
      )
  
      if config.TEST.SEQUENTIAL:
          sampler_val = torch.utils.data.SequentialSampler(dataset_val)
      else:
          sampler_val = torch.utils.data.distributed.DistributedSampler(
              dataset_val, shuffle=False
          )
  
      data_loader_train = torch.utils.data.DataLoader(
          dataset_train, sampler=sampler_train,
          batch_size=args.calib_batchsz if args.calib else config.DATA.BATCH_SIZE,
          num_workers=config.DATA.NUM_WORKERS,
          pin_memory=config.DATA.PIN_MEMORY,
          drop_last=True,
      )
  
      data_loader_val = torch.utils.data.DataLoader(
          dataset_val, sampler=sampler_val,
          batch_size=config.DATA.BATCH_SIZE,
          shuffle=False,
          num_workers=config.DATA.NUM_WORKERS,
          pin_memory=config.DATA.PIN_MEMORY,
          drop_last=False
      )
  
      return dataset_train, dataset_val, data_loader_train, data_loader_val
  
  
  def build_dataset(is_train, config):
      transform = build_transform(is_train, config)
      if config.DATA.DATASET == 'imagenet':
          prefix = 'train' if is_train else 'val'
          root = os.path.join(config.DATA.DATA_PATH, prefix)
          dataset = datasets.ImageFolder(root, transform=transform)
          nb_classes = 1000
      elif config.DATA.DATASET == 'imagenet22K':
          raise NotImplementedError("Imagenet-22K will come soon.")
      else:
          raise NotImplementedError("We only support ImageNet Now.")
  
      return dataset, nb_classes
  
  
  def build_transform(is_train, config):
      resize_im = config.DATA.IMG_SIZE > 32
      if is_train:
          # this should always dispatch to transforms_imagenet_train
          transform = create_transform(
              input_size=config.DATA.IMG_SIZE,
              is_training=True,
              color_jitter=config.AUG.COLOR_JITTER if config.AUG.COLOR_JITTER > 0 else None,
              auto_augment=config.AUG.AUTO_AUGMENT if config.AUG.AUTO_AUGMENT != 'none' else None,
              re_prob=config.AUG.REPROB,
              re_mode=config.AUG.REMODE,
              re_count=config.AUG.RECOUNT,
              interpolation=config.DATA.INTERPOLATION,
          )
          if not resize_im:
              # replace RandomResizedCropAndInterpolation with
              # RandomCrop
              transform.transforms[0] = transforms.RandomCrop(config.DATA.IMG_SIZE, padding=4)
          return transform
  
      t = []
      if resize_im:
          if config.TEST.CROP:
              size = int((256 / 224) * config.DATA.IMG_SIZE)
              t.append(
                  transforms.Resize(size, interpolation=_pil_interp(config.DATA.INTERPOLATION)),
                  # to maintain same ratio w.r.t. 224 images
              )
              t.append(transforms.CenterCrop(config.DATA.IMG_SIZE))
          else:
              t.append(
                  transforms.Resize((config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
                                    interpolation=_pil_interp(config.DATA.INTERPOLATION))
              )
  
      t.append(transforms.ToTensor())
      t.append(transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)))
      return transforms.Compose(t)
  