Repository: pyg-team/pytorch_geometric
Pull request #7750 (2023-07-15 22:47:20)

Replace `F.dropout` with `nn.Dropout` when used with `to_hetero`
Contributes to #7745.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e9eaa0d29..9bcc1c405 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added possibility to run inference benchmarks on XPU device ([#7705](https://github.com/pyg-team/pytorch_geometric/pull/7705))
 - Added `HeteroData` support in `to_networkx` ([#7713](https://github.com/pyg-team/pytorch_geometric/pull/7713))
 - Added `FlopsCount` support via `fvcore` ([#7693](https://github.com/pyg-team/pytorch_geometric/pull/7693))
 - Added back support for PyTorch >= 1.11.0 ([#7656](https://github.com/pyg-team/pytorch_geometric/pull/7656))
diff --git a/benchmark/inference/inference_benchmark.py b/benchmark/inference/inference_benchmark.py
index 380c1d53f..22d6dd7bf 100644
--- a/benchmark/inference/inference_benchmark.py
+++ b/benchmark/inference/inference_benchmark.py
@@ -16,7 +16,12 @@ from benchmark.utils import (
 )
 from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import PNAConv
-from torch_geometric.profile import rename_profile_file, timeit, torch_profile
+from torch_geometric.profile import (
+    rename_profile_file,
+    timeit,
+    torch_profile,
+    xpu_profile,
+)
 
 supported_sets = {
     'ogbn-mag': ['rgat', 'rgcn'],
@@ -42,11 +47,23 @@ def run(args: argparse.ArgumentParser):
         warnings.warn("Cannot write profile data to CSV because profiling is "
                       "disabled")
 
-    # cuda device is not suitable for full batch mode
-    device = torch.device(
-        'cuda' if not args.full_batch and torch.cuda.is_available() else 'cpu')
+    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."
@@ -66,11 +83,17 @@ def run(args: argparse.ArgumentParser):
         if args.num_layers != [1] and not hetero and args.num_steps != -1:
             raise ValueError("Layer-wise inference requires `steps=-1`")
 
-        if torch.cuda.is_available():
+        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
@@ -163,6 +186,8 @@ def run(args: argparse.ArgumentParser):
                             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:
@@ -170,9 +195,13 @@ def run(args: argparse.ArgumentParser):
                                 args.loader_cores)
                         else:
                             cpu_affinity = nullcontext()
-                        profile = torch_profile(
-                            args.export_chrome_trace, csv_data,
-                            args.write_csv) if args.profile else 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()
 
@@ -256,6 +285,8 @@ 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('--datasets', nargs='+',
         default=['ogbn-mag', 'ogbn-products', 'Reddit'], type=str)
     add('--use-sparse-tensor', action='store_true',
diff --git a/examples/pytorch_lightning/relational_gnn.py b/examples/pytorch_lightning/relational_gnn.py
index 042ebdf60..fa7b264c7 100644
--- a/examples/pytorch_lightning/relational_gnn.py
+++ b/examples/pytorch_lightning/relational_gnn.py
@@ -21,7 +21,7 @@ class GNN(torch.nn.Module):
     def __init__(self, hidden_channels: int, out_channels: int,
                  dropout: float):
         super().__init__()
-        self.dropout = dropout
+        self.dropout = torch.nn.Dropout(p=dropout)
 
         self.conv1 = SAGEConv((-1, -1), hidden_channels)
         self.conv2 = SAGEConv((-1, -1), hidden_channels)
@@ -29,9 +29,9 @@ class GNN(torch.nn.Module):
 
     def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
         x = self.conv1(x, edge_index).relu()
-        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.dropout(x)
         x = self.conv2(x, edge_index).relu()
-        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.dropout(x)
         return self.lin(x)
 
 
diff --git a/torch_geometric/profile/__init__.py b/torch_geometric/profile/__init__.py
index 9f7340979..a72d953e6 100644
--- a/torch_geometric/profile/__init__.py
+++ b/torch_geometric/profile/__init__.py
@@ -4,6 +4,7 @@ from .profile import (
     print_time_total,
     rename_profile_file,
     torch_profile,
+    xpu_profile,
 )
 from .utils import count_parameters
 from .utils import get_model_size
@@ -21,6 +22,7 @@ __all__ = [
     'print_time_total',
     'rename_profile_file',
     'torch_profile',
+    'xpu_profile',
     'count_parameters',
     'get_model_size',
     'get_data_size',
diff --git a/torch_geometric/profile/profile.py b/torch_geometric/profile/profile.py
index fc8f0a887..748d7f57b 100644
--- a/torch_geometric/profile/profile.py
+++ b/torch_geometric/profile/profile.py
@@ -265,6 +265,15 @@ def torch_profile(export_chrome_trace=True, csv_data=None, write_csv=None):
         save_profile_data(csv_data, events, use_cuda)
 
 
+@contextmanager
+def xpu_profile(export_chrome_trace=True):
+    with torch.autograd.profiler_legacy.profile(use_xpu=True) as profile:
+        yield
+    print(profile.key_averages().table(sort_by='self_xpu_time_total'))
+    if export_chrome_trace:
+        profile.export_chrome_trace('timeline.json')
+
+
 def format_prof_time(time):
     # Profile time is in micro seconds, so format it appropriately:
     return round(time / 1e6, 3)


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index 4b4fd3ee1..2271fb59d 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -339,12 +339,13 @@ def test_data_share_memory():
 
     for data in data_list:
         assert not data.x.is_shared()
+        assert torch.all(data.x == 0.0)
 
     mp.spawn(run, args=(data_list, ), nprocs=4, join=True)
 
     for data in data_list:
         assert data.x.is_shared()
-        assert torch.allclose(data.x, torch.full((8, ), 4.))
+        assert torch.all(data.x > 0.0)
 
 
 def test_data_setter_properties():
diff --git a/test/loader/test_imbalanced_sampler.py b/test/loader/test_imbalanced_sampler.py
index 28fa7c50d..eb042765f 100644
--- a/test/loader/test_imbalanced_sampler.py
+++ b/test/loader/test_imbalanced_sampler.py
@@ -3,7 +3,7 @@ from typing import List
 import torch
 
 from torch_geometric.data import Data
-from torch_geometric.datasets import FakeDataset
+from torch_geometric.datasets import FakeDataset, FakeHeteroDataset
 from torch_geometric.loader import (
     DataLoader,
     ImbalancedSampler,
@@ -92,3 +92,25 @@ def test_neighbor_loader_with_imbalanced_sampler():
                             num_neighbors=[-1])
 
     assert torch.allclose(y, torch.cat([batch.y for batch in loader]))
+
+
+@onlyNeighborSampler
+def test_hetero_neighbor_loader_with_imbalanced_sampler():
+    torch.manual_seed(12345)
+    data = FakeHeteroDataset(num_classes=2)[0]
+
+    loader = NeighborLoader(
+        data,
+        batch_size=100,
+        input_nodes='v0',
+        num_neighbors=[-1],
+        sampler=ImbalancedSampler(data['v0'].y),
+    )
+
+    y = torch.cat([batch['v0'].y[:batch['v0'].batch_size] for batch in loader])
+
+    histogram = y.bincount()
+    prob = histogram / histogram.sum()
+
+    assert histogram.sum() == data['v0'].num_nodes
+    assert prob.min() > 0.4 and prob.max() < 0.6
diff --git a/test/profile/test_profile.py b/test/profile/test_profile.py
index c4b5c38f3..b64d1125f 100644
--- a/test/profile/test_profile.py
+++ b/test/profile/test_profile.py
@@ -1,6 +1,7 @@
 import os.path
 import warnings
 
+import pytest
 import torch
 import torch.nn.functional as F
 
@@ -11,11 +12,12 @@ from torch_geometric.profile import (
     rename_profile_file,
     timeit,
 )
-from torch_geometric.profile.profile import torch_profile
+from torch_geometric.profile.profile import torch_profile, xpu_profile
 from torch_geometric.testing import (
     onlyCUDA,
     onlyLinux,
     onlyOnline,
+    onlyXPU,
     withCUDA,
     withPackage,
 )
@@ -105,3 +107,30 @@ def test_torch_profile(capfd, get_dataset, device):
     rename_profile_file('test_profile')
     assert os.path.exists('profile-test_profile.json')
     os.remove('profile-test_profile.json')
+
+
+@onlyXPU
+@onlyOnline
+@pytest.mark.parametrize('export_chrome_trace', [False, True])
+def test_xpu_profile(capfd, get_dataset, export_chrome_trace):
+    dataset = get_dataset(name='Cora')
+    device = torch.device('xpu')
+    data = dataset[0].to(device)
+    model = GraphSAGE(dataset.num_features, hidden_channels=64, num_layers=3,
+                      out_channels=dataset.num_classes).to(device)
+
+    with xpu_profile(export_chrome_trace):
+        model(data.x, data.edge_index)
+
+    out, _ = capfd.readouterr()
+    assert 'Self CPU' in out
+    if data.x.is_xpu:
+        assert 'Self XPU' in out
+
+    f_name = 'timeline.json'
+    f_exists = os.path.exists(f_name)
+    if not export_chrome_trace:
+        assert not f_exists
+    else:
+        assert f_exists
+        os.remove(f_name)
diff --git a/torch_geometric/testing/__init__.py b/torch_geometric/testing/__init__.py
index 83f982041..6a108d42f 100644
--- a/torch_geometric/testing/__init__.py
+++ b/torch_geometric/testing/__init__.py
@@ -4,6 +4,7 @@ from .decorators import (
     onlyLinux,
     onlyPython,
     onlyCUDA,
+    onlyXPU,
     onlyOnline,
     onlyGraphviz,
     onlyNeighborSampler,
@@ -22,6 +23,7 @@ __all__ = [
     'onlyLinux',
     'onlyPython',
     'onlyCUDA',
+    'onlyXPU',
     'onlyOnline',
     'onlyGraphviz',
     'onlyNeighborSampler',
diff --git a/torch_geometric/testing/decorators.py b/torch_geometric/testing/decorators.py
index b62625fa6..8c5e35c16 100644
--- a/torch_geometric/testing/decorators.py
+++ b/torch_geometric/testing/decorators.py
@@ -59,6 +59,20 @@ def onlyCUDA(func: Callable) -> Callable:
     )(func)
 
 
+def onlyXPU(func: Callable) -> Callable:
+    r"""A decorator to skip tests if XPU is not found."""
+    import pytest
+    try:
+        import intel_extension_for_pytorch as ipex
+        xpu_available = ipex.xpu.is_available()
+    except ImportError:
+        xpu_available = False
+    return pytest.mark.skipif(
+        not xpu_available,
+        reason="XPU not available",
+    )(func)
+
+
 def onlyOnline(func: Callable):
     r"""A decorator to skip tests if there exists no connection to the
     internet."""