code: import argparse
  import os
  import os.path as osp
  import pickle
  import shutil
  import tempfile
  
  import mmcv
  import torch
  import torch.distributed as dist
  from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
  from mmcv.runner import get_dist_info, init_dist, load_checkpoint
  
  from mmdet.core import wrap_fp16_model
  from mmdet.datasets import build_dataloader, build_dataset
  from mmdet.models import build_detector
  
  
  def single_gpu_test(model, data_loader, show=False):
      model.eval()
      results = []
      dataset = data_loader.dataset
      prog_bar = mmcv.ProgressBar(len(dataset))
      for i, data in enumerate(data_loader):
          with torch.no_grad():
              result = model(return_loss=False, rescale=not show, **data)
          results.append(result)
  
          if show:
              model.module.show_result(data, result)
  
          batch_size = data['img'][0].size(0)
          for _ in range(batch_size):
              prog_bar.update()
      return results
  
  
  def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False):
      """Test model with multiple gpus.
  
      This method tests model with multiple gpus and collects the results
      under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
      it encodes results to gpu tensors and use gpu communication for results
      collection. On cpu mode it saves the results on different gpus to 'tmpdir'
      and collects them by the rank 0 worker.
  
      Args:
          model (nn.Module): Model to be tested.
          data_loader (nn.Dataloader): Pytorch data loader.
          tmpdir (str): Path of directory to save the temporary results from
              different gpus under cpu mode.
          gpu_collect (bool): Option to use either gpu or cpu to collect results.
  
      Returns:
          list: The prediction results.
      """
      model.eval()
      results = []
      dataset = data_loader.dataset
      rank, world_size = get_dist_info()
      if rank == 0:
          prog_bar = mmcv.ProgressBar(len(dataset))
      for i, data in enumerate(data_loader):
          with torch.no_grad():
              result = model(return_loss=False, rescale=True, **data)
          results.append(result)
  
          if rank == 0:
              batch_size = data['img'][0].size(0)
              for _ in range(batch_size * world_size):
                  prog_bar.update()
  
      # collect results from all ranks
      if gpu_collect:
          results = collect_results_gpu(results, len(dataset))
      else:
          results = collect_results_cpu(results, len(dataset), tmpdir)
      return results
  
  
  def collect_results_cpu(result_part, size, tmpdir=None):
      rank, world_size = get_dist_info()
      # create a tmp dir if it is not specified
      if tmpdir is None:
          MAX_LEN = 512
          # 32 is whitespace
          dir_tensor = torch.full((MAX_LEN, ),
                                  32,
                                  dtype=torch.uint8,
                                  device='cuda')
          if rank == 0:
              tmpdir = tempfile.mkdtemp()
              tmpdir = torch.tensor(
                  bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda')
              dir_tensor[:len(tmpdir)] = tmpdir
          dist.broadcast(dir_tensor, 0)
          tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip()
      else:
          mmcv.mkdir_or_exist(tmpdir)
      # dump the part result to the dir
      mmcv.dump(result_part, osp.join(tmpdir, 'part_{}.pkl'.format(rank)))
      dist.barrier()
      # collect all parts
      if rank != 0:
          return None
      else:
          # load results of all parts from tmp dir
          part_list = []
          for i in range(world_size):
              part_file = osp.join(tmpdir, 'part_{}.pkl'.format(i))
              part_list.append(mmcv.load(part_file))
          # sort the results
          ordered_results = []
          for res in zip(*part_list):
              ordered_results.extend(list(res))
          # the dataloader may pad some samples
          ordered_results = ordered_results[:size]
          # remove tmp dir
          shutil.rmtree(tmpdir)
          return ordered_results
  
  
  def collect_results_gpu(result_part, size):
      rank, world_size = get_dist_info()
      # dump result part to tensor with pickle
      part_tensor = torch.tensor(
          bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
      # gather all result part tensor shape
      shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
      shape_list = [shape_tensor.clone() for _ in range(world_size)]
      dist.all_gather(shape_list, shape_tensor)
      # padding result part tensor to max length
      shape_max = torch.tensor(shape_list).max()
      part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
      part_send[:shape_tensor[0]] = part_tensor
      part_recv_list = [
          part_tensor.new_zeros(shape_max) for _ in range(world_size)
      ]
      # gather all result part
      dist.all_gather(part_recv_list, part_send)
  
      if rank == 0:
          part_list = []
          for recv, shape in zip(part_recv_list, shape_list):
              part_list.append(
                  pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))
          # sort the results
          ordered_results = []
          for res in zip(*part_list):
              ordered_results.extend(list(res))
          # the dataloader may pad some samples
          ordered_results = ordered_results[:size]
          return ordered_results
  
  
  class MultipleKVAction(argparse.Action):
      """
      argparse action to split an argument into KEY=VALUE form
      on the first = and append to a dictionary.
      """
  
      def _is_int(self, val):
          try:
              _ = int(val)
              return True
          except Exception:
              return False
  
      def _is_float(self, val):
          try:
              _ = float(val)
              return True
          except Exception:
              return False
  
      def _is_bool(self, val):
          return val.lower() in ['true', 'false']
  
      def __call__(self, parser, namespace, values, option_string=None):
          options = {}
          for val in values:
              parts = val.split('=')
              key = parts[0].strip()
              if len(parts) > 2:
                  val = '='.join(parts[1:])
              else:
                  val = parts[1].strip()
              # try parsing val to bool/int/float first
              if self._is_bool(val):
                  import json
                  val = json.loads(val.lower())
              elif self._is_int(val):
                  val = int(val)
              elif self._is_float(val):
                  val = float(val)
              options[key] = val
          setattr(namespace, self.dest, options)
  
  
  def parse_args():
      parser = argparse.ArgumentParser(
          description='MMDet test (and eval) a model')
      parser.add_argument('config', help='test config file path')
      parser.add_argument('checkpoint', help='checkpoint file')
      parser.add_argument('--out', help='output result file in pickle format')
      parser.add_argument(
          '--eval',
          type=str,
          nargs='+',
          help='evaluation metrics, which depends on the dataset, e.g., "bbox",'
          ' "segm", "proposal" for COCO, and "mAP", "recall" for PASCAL VOC')
      parser.add_argument('--show', action='store_true', help='show results')
      parser.add_argument(
          '--gpu_collect',
          action='store_true',
          help='whether to use gpu to collect results.')
      parser.add_argument(
          '--tmpdir',
          help='tmp directory used for collecting results from multiple '
          'workers, available when gpu_collect is not specified')
      parser.add_argument(
          '--options', nargs='+', action=MultipleKVAction, help='custom options')
      parser.add_argument(
          '--launcher',
          choices=['none', 'pytorch', 'slurm', 'mpi'],
          default='none',
          help='job launcher')
      parser.add_argument('--local_rank', type=int, default=0)
      args = parser.parse_args()
      if 'LOCAL_RANK' not in os.environ:
          os.environ['LOCAL_RANK'] = str(args.local_rank)
      return args
  
  
  def main():
      args = parse_args()
  
      assert args.out or args.eval or args.show, \
          ('Please specify at least one operation (save or eval or show the '
           'results) with the argument "--out", "--eval" or "--show"')
  
      if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
          raise ValueError('The output file must be a pkl file.')
  
      cfg = mmcv.Config.fromfile(args.config)
      # set cudnn_benchmark
      if cfg.get('cudnn_benchmark', False):
          torch.backends.cudnn.benchmark = True
      cfg.model.pretrained = None
      cfg.data.test.test_mode = True
  
      # init distributed env first, since logger depends on the dist info.
      if args.launcher == 'none':
          distributed = False
      else:
          distributed = True
          init_dist(args.launcher, **cfg.dist_params)
  
      # build the dataloader
      # TODO: support multiple images per gpu (only minor changes are needed)
      dataset = build_dataset(cfg.data.test)
      data_loader = build_dataloader(
          dataset,
          imgs_per_gpu=1,
          workers_per_gpu=cfg.data.workers_per_gpu,
          dist=distributed,
          shuffle=False)
  
      # build the model and load checkpoint
      model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)
      fp16_cfg = cfg.get('fp16', None)
      if fp16_cfg is not None:
          wrap_fp16_model(model)
      checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
      # old versions did not save class info in checkpoints, this walkaround is
      # for backward compatibility
      if 'CLASSES' in checkpoint['meta']:
          model.CLASSES = checkpoint['meta']['CLASSES']
      else:
          model.CLASSES = dataset.CLASSES
  
      if not distributed:
          model = MMDataParallel(model, device_ids=[0])
          outputs = single_gpu_test(model, data_loader, args.show)
      else:
          model = MMDistributedDataParallel(
              model.cuda(),
              device_ids=[torch.cuda.current_device()],
              broadcast_buffers=False)
          outputs = multi_gpu_test(model, data_loader, args.tmpdir,
                                   args.gpu_collect)
  
      rank, _ = get_dist_info()
      if rank == 0:
          if args.out:
              print('\nwriting results to {}'.format(args.out))
              mmcv.dump(outputs, args.out)
          if args.eval:
              kwargs = {} if args.options is None else args.options
              dataset.evaluate(outputs, args.eval, **kwargs)
  
  
  if __name__ == '__main__':
      main()
  
apis: torch.distributed.all_gather torch.distributed.broadcast torch.no_grad torch.tensor torch.full torch.cuda.current_device torch.distributed.barrier torch.zeros