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.io import fs 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, ) supported_sets = { 'ogbn-mag': ['rgat', 'rgcn'], 'ogbn-products': ['edge_cnn', 'gat', 'gcn', 'pna', 'sage'], 'Reddit': ['edge_cnn', 'gat', 'gcn', 'pna', 'sage'], } @torch.no_grad() def full_batch_inference(model, data): model.eval() if hasattr(data, 'adj_t'): edge_index = data.adj_t else: edge_index = data.edge_index return model(data.x, edge_index) 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.amp.autocast('cuda', 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 = fs.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) if __name__ == '__main__': argparser = argparse.ArgumentParser('GNN inference benchmark') add = argparser.add_argument add('--device', choices=['cpu', 'cuda', 'xpu'], default='cpu', help='Device to run benchmark on') add('--reuse-device-for-embeddings', action='store_true', help='Use the same device for embeddings as specified in "--device"') add('--datasets', nargs='+', default=['ogbn-mag', 'ogbn-products', 'Reddit'], type=str) add('--use-sparse-tensor', action='store_true', help='use torch_sparse.SparseTensor as graph storage format') add('--models', nargs='+', default=['edge_cnn', 'gat', 'gcn', 'pna', 'rgat', 'rgcn'], type=str) add('--root', default='../../data', type=str, help='relative path to look for the datasets') add('--eval-batch-sizes', nargs='+', default=[512, 1024, 2048, 4096, 8192], type=int) add('--num-layers', nargs='+', default=[2, 3], type=int) add('--num-hidden-channels', nargs='+', default=[64, 128, 256], type=int) add('--num-heads', default=2, type=int, help='number of hidden attention heads, applies only for gat and rgat') add('--hetero-num-neighbors', default=10, type=int, help='number of neighbors to sample per layer for hetero workloads') add('--num-workers', default=0, type=int) add('--num-steps', default=-1, type=int, help='number of steps, -1 means iterating through all the data') add('--warmup', default=1, type=int) add('--profile', action='store_true') add('--vtune-profile', action='store_true') add('--bf16', action='store_true') add('--cpu-affinity', action='store_true', help='Use DataLoader affinitzation.') add('--loader-cores', nargs='+', default=[], type=int, help="List of CPU core IDs to use for DataLoader workers") add('--measure-load-time', action='store_true') add('--full-batch', action='store_true', help='Use full batch mode') add('--evaluate', action='store_true') add('--ckpt_path', type=str, help='Checkpoint path for loading a model') add('--write-csv', choices=[None, 'bench', 'prof'], default=None, help='Write benchmark or PyTorch profile data to CSV') add('--export-chrome-trace', default=True, type=bool, help='Export chrome trace file. Works only with PyTorch profiler') add('--cached-loader', action='store_true', help='Use CachedLoader') run(argparser.parse_args())