text: from paddle.io import Sampler
  import paddle.distributed as dist
  
  import math
  import random
  import numpy as np
  
  from ppcls import data
  
  
  class MultiScaleSampler(Sampler):
      def __init__(self,
                   data_source,
                   scales,
                   first_bs,
                   divided_factor=32,
                   is_training=True,
                   seed=None):
          """
              multi scale samper
              Args:
                  data_source(dataset)
                  scales(list): several scales for image resolution
                  first_bs(int): batch size for the first scale in scales
                  divided_factor(int): ImageNet models down-sample images by a factor, ensure that width and height dimensions are multiples are multiple of devided_factor.
                  is_training(boolean): mode 
          """
          # min. and max. spatial dimensions
          self.data_source = data_source
          self.n_data_samples = len(self.data_source)
  
          if isinstance(scales[0], tuple):
              width_dims = [i[0] for i in scales]
              height_dims = [i[1] for i in scales]
          elif isinstance(scales[0], int):
              width_dims = scales
              height_dims = scales
          base_im_w = width_dims[0]
          base_im_h = height_dims[0]
          base_batch_size = first_bs
  
          # Get the GPU and node related information
          num_replicas = dist.get_world_size()
          rank = dist.get_rank()
          # adjust the total samples to avoid batch dropping
          num_samples_per_replica = int(
              math.ceil(self.n_data_samples * 1.0 / num_replicas))
          img_indices = [idx for idx in range(self.n_data_samples)]
  
          self.shuffle = False
          if is_training:
              # compute the spatial dimensions and corresponding batch size
              # ImageNet models down-sample images by a factor of 32.
              # Ensure that width and height dimensions are multiples are multiple of 32.
              width_dims = [
                  int((w // divided_factor) * divided_factor) for w in width_dims
              ]
              height_dims = [
                  int((h // divided_factor) * divided_factor)
                  for h in height_dims
              ]
  
              img_batch_pairs = list()
              base_elements = base_im_w * base_im_h * base_batch_size
              for (h, w) in zip(height_dims, width_dims):
                  batch_size = int(max(1, (base_elements / (h * w))))
                  img_batch_pairs.append((w, h, batch_size))
              self.img_batch_pairs = img_batch_pairs
              self.shuffle = True
          else:
              self.img_batch_pairs = [(base_im_w, base_im_h, base_batch_size)]
  
          self.img_indices = img_indices
          self.n_samples_per_replica = num_samples_per_replica
          self.epoch = 0
          self.rank = rank
          self.num_replicas = num_replicas
          self.seed = seed
          self.batch_list = []
          self.current = 0
          indices_rank_i = self.img_indices[self.rank:len(self.img_indices):
                                            self.num_replicas]
          while self.current < self.n_samples_per_replica:
              curr_w, curr_h, curr_bsz = random.choice(self.img_batch_pairs)
  
              end_index = min(self.current + curr_bsz,
                              self.n_samples_per_replica)
  
              batch_ids = indices_rank_i[self.current:end_index]
              n_batch_samples = len(batch_ids)
              if n_batch_samples != curr_bsz:
                  batch_ids += indices_rank_i[:(curr_bsz - n_batch_samples)]
              self.current += curr_bsz
  
              if len(batch_ids) > 0:
                  batch = [curr_w, curr_h, len(batch_ids)]
                  self.batch_list.append(batch)
          self.length = len(self.batch_list)
  
      def __iter__(self):
          if self.shuffle:
              if self.seed is not None:
                  random.seed(self.seed)
              else:
                  random.seed(self.epoch)
              random.shuffle(self.img_indices)
              random.shuffle(self.img_batch_pairs)
              indices_rank_i = self.img_indices[self.rank:len(self.img_indices):
                                                self.num_replicas]
          else:
              indices_rank_i = self.img_indices[self.rank:len(self.img_indices):
                                                self.num_replicas]
  
          start_index = 0
          for batch_tuple in self.batch_list:
              curr_w, curr_h, curr_bsz = batch_tuple
              end_index = min(start_index + curr_bsz, self.n_samples_per_replica)
              batch_ids = indices_rank_i[start_index:end_index]
              n_batch_samples = len(batch_ids)
              if n_batch_samples != curr_bsz:
                  batch_ids += indices_rank_i[:(curr_bsz - n_batch_samples)]
              start_index += curr_bsz
  
              if len(batch_ids) > 0:
                  batch = [(curr_w, curr_h, b_id) for b_id in batch_ids]
                  yield batch
  
      def set_epoch(self, epoch: int):
          self.epoch = epoch
  
      def __len__(self):
          return self.length
  