code: def benchmark_roi_align():
      def random_boxes(mean_box, stdev, N, maxsize):
          ret = torch.rand(N, 4) * stdev + torch.tensor(
              mean_box, dtype=torch.float
          )
          ret.clamp_(min=0, max=maxsize)
          return ret
  
      def func(
          shape,
          nboxes_per_img,
          sampling_ratio,
          device,
          box_size="large",
      ):
          N, _, H, _ = shape
          input = torch.rand(*shape)
          boxes = []
          batch_idx = []
          for k in range(N):
              if box_size == "large":
                  b = random_boxes(
                      [80, 80, 130, 130], 24, nboxes_per_img, H
                  )
              else:
                  b = random_boxes(
                      [100, 100, 110, 110], 4, nboxes_per_img, H
                  )
              boxes.append(b)
              batch_idx.append(
                  torch.zeros(nboxes_per_img, 1, dtype=torch.float32)
                  + k
              )
          boxes = torch.cat(boxes, axis=0)
          batch_idx = torch.cat(batch_idx, axis=0)
          boxes = torch.cat([batch_idx, boxes], axis=1)
  
          input = input.to(device=device)
          boxes = boxes.to(device=device)
  
          def bench():
              if False and sampling_ratio > 0 and N == 1:
                  # enable to benchmark grid_sample (slower)
                  grid_sample_roi_align(
                      input, boxes[:, 1:], 7, 1.0, sampling_ratio
                  )
              else:
                  roi_align(input, boxes, 7, 1.0, sampling_ratio, True)
              if device == "cuda":
                  torch.cuda.synchronize()
  
          return bench
  
      def gen_args(arg):
          args = []
          for size in ["small", "large"]:
              for ratio in [0, 2]:
                  args.append(copy(arg))
                  args[-1]["sampling_ratio"] = ratio
                  args[-1]["box_size"] = size
          return args
  
      arg = dict(
          shape=(1, 512, 256, 256), nboxes_per_img=512, device="cuda"
      )
      benchmark(
          func,
          "cuda_roialign",
          gen_args(arg),
          num_iters=20,
          warmup_iters=1,
      )
      arg.update({"device": "cpu", "shape": (1, 256, 128, 128)})
      benchmark(
          func,
          "cpu_roialign",
          gen_args(arg),
          num_iters=5,
          warmup_iters=1,
      )
  
file_path: detectron2/tests/layers/test_roi_align.py
imports: import numpy
  import unittest
  from copy import copy
  import cv2
  import torch
  from fvcore.common.benchmark import benchmark
  from torch.nn import functional
  from detectron2.layers.roi_align import ROIAlign, roi_align
label_desc: Use image node_id 2 for calling a global function with example usage: benchmark_roi_align() and returns: ret, bench, args