code: def run(args: argparse.ArgumentParser):
      csv_data = defaultdict(list)
  
      if args.write_csv == "prof" and not args.profile:
          warnings.warn(
              "Cannot write profile data to CSV because profiling is "
              "disabled"
          )
  
      if args.device == "xpu":
          try:
              import intel_extension_for_pytorch as ipex
          except ImportError:
              raise RuntimeError(
                  "XPU device requires IPEX to be installed"
              )
  
      if (args.device == "cuda" and not torch.cuda.is_available()) or (
          args.device == "xpu" and not torch.xpu.is_available()
      ):
          raise RuntimeError(f"{args.device.upper()} is not available")
  
      if args.device == "cuda" and args.full_batch:
          raise RuntimeError(
              "CUDA device is not suitable for full batch mode"
          )
  
      device = torch.device(args.device)
  
      print("BENCHMARK STARTS")
      print(f"Running on {args.device.upper()}")
      for dataset_name in args.datasets:
          assert (
              dataset_name in supported_sets.keys()
          ), f"Dataset {dataset_name} isn't supported."
          print(f"Dataset: {dataset_name}")
          load_time = (
              timeit() if args.measure_load_time else nullcontext()
          )
          with load_time:
              result = get_dataset_with_transformation(
                  dataset_name,
                  args.root,
                  args.use_sparse_tensor,
                  args.bf16,
              )
              dataset, num_classes, transformation = result
          data = dataset.to(device)
          hetero = True if dataset_name == "ogbn-mag" else False
          mask = ("paper", None) if dataset_name == "ogbn-mag" else None
          _, _, test_mask = get_split_masks(data, dataset_name)
          degree = None
  
          if hetero and args.cached_loader:
              args.cached_loader = False
              print(
                  "Disabling CachedLoader, not supported in Hetero models"
              )
          if (
              args.num_layers != [1]
              and not hetero
              and args.num_steps != -1
          ):
              raise ValueError(
                  "Layer-wise inference requires `steps=-1`"
              )
  
          if args.device == "cuda":
              amp = torch.cuda.amp.autocast(enabled=False)
          elif args.device == "xpu":
              amp = torch.xpu.amp.autocast(enabled=False)
          else:
              amp = torch.cpu.amp.autocast(enabled=args.bf16)
  
          if args.device == "xpu" and args.warmup < 1:
              print("XPU device requires warmup - setting warmup=1")
              args.warmup = 1
  
          inputs_channels = (
              data["paper"].num_features
              if dataset_name == "ogbn-mag"
              else dataset.num_features
          )
  
          for model_name in args.models:
              if model_name not in supported_sets[dataset_name]:
                  print(
                      f"Configuration of {dataset_name} + {model_name} "
                      f"not supported. Skipping."
                  )
                  continue
              with_loader = not args.full_batch or (
                  model_name == "pna" and degree is None
              )
              print(f"Evaluation bench for {model_name}:")
  
              for batch_size in args.eval_batch_sizes:
                  num_nodes = (
                      data["paper"].num_nodes
                      if hetero
                      else data.num_nodes
                  )
                  sampler = (
                      torch.utils.data.RandomSampler(
                          range(num_nodes),
                          num_samples=args.num_steps * batch_size,
                      )
                      if args.num_steps != -1 and with_loader
                      else None
                  )
                  kwargs = {
                      "batch_size": batch_size,
                      "shuffle": False,
                      "num_workers": args.num_workers,
                  }
                  if not hetero:
                      subgraph_loader = (
                          NeighborLoader(
                              data,
                              num_neighbors=[
                                  -1
                              ],  # layer-wise inference
                              input_nodes=mask,
                              sampler=sampler,
                              **kwargs,
                          )
                          if with_loader
                          else None
                      )
                      if args.evaluate and not args.full_batch:
                          test_loader = NeighborLoader(
                              data,
                              num_neighbors=[
                                  -1
                              ],  # layer-wise inference
                              input_nodes=test_mask,
                              sampler=None,
                              **kwargs,
                          )
  
                  for layers in args.num_layers:
                      num_neighbors = [
                          args.hetero_num_neighbors
                      ] * layers
                      if hetero:
                          # batch-wise inference
                          subgraph_loader = (
                              NeighborLoader(
                                  data,
                                  num_neighbors=num_neighbors,
                                  input_nodes=mask,
                                  sampler=sampler,
                                  **kwargs,
                              )
                              if with_loader
                              else None
                          )
                          if args.evaluate and not args.full_batch:
                              test_loader = NeighborLoader(
                                  data,
                                  num_neighbors=num_neighbors,
                                  input_nodes=test_mask,
                                  sampler=None,
                                  **kwargs,
                              )
  
                      for hidden_channels in args.num_hidden_channels:
                          print(
                              "----------------------------------------------"
                          )
                          print(
                              f"Batch size={batch_size}, "
                              f"Layers amount={layers}, "
                              f"Num_neighbors={num_neighbors}, "
                              f"Hidden features size={hidden_channels}, "
                              f"Sparse tensor={args.use_sparse_tensor}"
                          )
                          params = {
                              "inputs_channels": inputs_channels,
                              "hidden_channels": hidden_channels,
                              "output_channels": num_classes,
                              "num_heads": args.num_heads,
                              "num_layers": layers,
                          }
  
                          if model_name == "pna":
                              if degree is None:
                                  degree = PNAConv.get_degree_histogram(
                                      subgraph_loader
                                  )
                                  print(
                                      f"Calculated degree for {dataset_name}."
                                  )
                              params["degree"] = degree
  
                          model = get_model(
                              model_name,
                              params,
                              metadata=data.metadata()
                              if hetero
                              else None,
                          )
                          model = model.to(device)
                          # TODO: Migrate to ModelHubMixin.
                          if args.ckpt_path:
                              state_dict = torch.load(args.ckpt_path)
                              model.load_state_dict(state_dict)
                          model.eval()
                          if args.device == "xpu":
                              model = ipex.optimize(model)
  
                          # Define context manager parameters:
                          if args.cpu_affinity and with_loader:
                              cpu_affinity = (
                                  subgraph_loader.enable_cpu_affinity(
                                      args.loader_cores
                                  )
                              )
                          else:
                              cpu_affinity = nullcontext()
                          if args.profile and args.device == "xpu":
                              profile = xpu_profile(
                                  args.export_chrome_trace
                              )
                          elif args.profile:
                              profile = torch_profile(
                                  args.export_chrome_trace,
                                  csv_data,
                                  args.write_csv,
                              )
                          else:
                              profile = nullcontext()
                          itt = (
                              emit_itt()
                              if args.vtune_profile
                              else nullcontext()
                          )
  
                          if args.full_batch and args.use_sparse_tensor:
                              data = transformation(data)
  
                          with cpu_affinity, amp, timeit() as time:
                              inference_kwargs = dict(
                                  cache=args.cached_loader
                              )
                              if (
                                  args.reuse_device_for_embeddings
                                  and not hetero
                              ):
                                  inference_kwargs[
                                      "embedding_device"
                                  ] = device
                              for _ in range(args.warmup):
                                  if args.full_batch:
                                      full_batch_inference(model, data)
                                  else:
                                      model.inference(
                                          subgraph_loader,
                                          device,
                                          progress_bar=True,
                                          **inference_kwargs,
                                      )
                              if args.warmup > 0:
                                  time.reset()
                              with itt, profile:
                                  if args.full_batch:
                                      y = full_batch_inference(
                                          model, data
                                      )
                                      if args.evaluate:
                                          mask = data.test_mask
                                          pred = y[mask].argmax(1)
                                          test_acc = (
                                              pred.eq(data.y[mask])
                                              .sum()
                                              .item()
                                              / mask.sum().item()
                                          )
                                          print(
                                              f"Full Batch Test Accuracy: \
                                              {test_acc:.4f}"
                                          )
                                  else:
                                      y = model.inference(
                                          subgraph_loader,
                                          device,
                                          progress_bar=True,
                                          **inference_kwargs,
                                      )
                                      if args.evaluate:
                                          test_acc = test(
                                              model,
                                              test_loader,
                                              device,
                                              hetero,
                                              progress_bar=True,
                                          )
                                          print(
                                              f"Mini Batch Test Accuracy: \
                                              {test_acc:.4f}"
                                          )
  
                          if args.profile and args.export_chrome_trace:
                              rename_profile_file(
                                  model_name,
                                  dataset_name,
                                  str(batch_size),
                                  str(layers),
                                  str(hidden_channels),
                                  str(num_neighbors),
                              )
                          total_time = time.duration
                          if args.num_steps != -1:
                              total_num_samples = (
                                  args.num_steps * batch_size
                              )
                          else:
                              total_num_samples = num_nodes
                          throughput = total_num_samples / total_time
                          latency = (
                              total_time / total_num_samples * 1000
                          )
                          print(
                              f"Throughput: {throughput:.3f} samples/s"
                          )
                          print(f"Latency: {latency:.3f} ms")
  
                          num_records = 1
                          if args.write_csv == "prof":
                              # For profiling with PyTorch, we save the top-5
                              # most time consuming operations. Therefore, the
                              # same data should be entered for each of them.
                              num_records = 5
                          for _ in range(num_records):
                              save_benchmark_data(
                                  csv_data,
                                  batch_size,
                                  layers,
                                  num_neighbors,
                                  hidden_channels,
                                  total_time,
                                  model_name,
                                  dataset_name,
                                  args.use_sparse_tensor,
                              )
      if args.write_csv:
          write_to_csv(csv_data, args.write_csv)
  
file_path: pytorch_geometric/benchmark/inference/inference_benchmark.py
imports: import argparse
  import warnings
  from collections import defaultdict
  from contextlib import nullcontext
  import torch
  from benchmark.utils import emit_itt, get_dataset_with_transformation, get_model, get_split_masks, save_benchmark_data, test, write_to_csv
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import PNAConv
  from torch_geometric.profile import rename_profile_file, timeit, torch_profile, xpu_profile
label_desc: Use image node_id 2 for calling a global function with example usage: run(args) without return types