text: import copy
  import random
  from collections import defaultdict
  from torch.utils.data.sampler import Sampler, RandomSampler, SequentialSampler
  
  
  class RandomDomainSampler(Sampler):
      """Random domain sampler.
  
      This sampler randomly samples N domains each with K
      images to form a minibatch.
      """
  
      def __init__(self, data_source, batch_size, n_domain):
          self.data_source = data_source
  
          # Keep track of image indices for each domain
          self.domain_dict = defaultdict(list)
          for i, item in enumerate(data_source):
              self.domain_dict[item.domain].append(i)
          self.domains = list(self.domain_dict.keys())
  
          # Make sure each domain has equal number of images
          if n_domain is None or n_domain <= 0:
              n_domain = len(self.domains)
          assert batch_size % n_domain == 0
          self.n_img_per_domain = batch_size // n_domain
  
          self.batch_size = batch_size
          # n_domain denotes number of domains sampled in a minibatch
          self.n_domain = n_domain
  
          # Estimate the number of images that will be used in each epoch
          self.length = 0
          for _, idxs in self.domain_dict.items():
              n_img = len(idxs)
              n_img -= n_img % self.n_img_per_domain
              self.length += n_img
  
      def __iter__(self):
          domain_dict = copy.deepcopy(self.domain_dict)
          final_idxs = []
          stop_sampling = False
  
          while not stop_sampling:
              selected_domains = random.sample(self.domains, self.n_domain)
  
              for domain in selected_domains:
                  idxs = domain_dict[domain]
                  selected_idxs = random.sample(idxs, self.n_img_per_domain)
                  final_idxs.extend(selected_idxs)
  
                  for idx in selected_idxs:
                      domain_dict[domain].remove(idx)
  
                  remaining = len(domain_dict[domain])
                  if remaining < self.n_img_per_domain:
                      stop_sampling = True
  
          return iter(final_idxs)
  
      def __len__(self):
          return self.length
  
  
  def build_sampler(
      sampler_type, cfg=None, data_source=None, batch_size=32, n_domain=0
  ):
      if sampler_type == 'RandomSampler':
          return RandomSampler(data_source)
  
      elif sampler_type == 'SequentialSampler':
          return SequentialSampler(data_source)
  
      elif sampler_type == 'RandomDomainSampler':
          return RandomDomainSampler(data_source, batch_size, n_domain)
  
      else:
          raise ValueError('Unknown sampler type: {}'.format(sampler_type))
  