Add XPU support for hetero examples (#9439)

Add XPU support for hetero examples --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Matthias Fey <[email protected]>
pyg-team · Jun 24, 2024 · e0ab5b6 · e0ab5b6
1 parent cd36c86
commit e0ab5b6
Show file tree

Hide file tree

Showing 9 changed files with 53 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -48,7 +48,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Made `MessagePassing` interface thread-safe ([#9001](https://github.com/pyg-team/pytorch_geometric/pull/9001))
 - Breaking Change: Added support for `EdgeIndex` in `cugraph` GNN layers ([#8938](https://github.com/pyg-team/pytorch_geometric/pull/8937))
 - Added the `dim` arg to `torch.cross` calls ([#8918](https://github.com/pyg-team/pytorch_geometric/pull/8918))
-- Added XPU support to basic GNN examples ([#9421](https://github.com/pyg-team/pytorch_geometric/pull/9001))
+- Added XPU support to basic GNN examples ([#9421](https://github.com/pyg-team/pytorch_geometric/pull/9421), [#9439](https://github.com/pyg-team/pytorch_geometric/pull/9439))
 
 ### Deprecated
 

diff --git a/examples/hetero/bipartite_sage_unsup.py b/examples/hetero/bipartite_sage_unsup.py
@@ -8,13 +8,19 @@
 from sklearn.metrics import roc_auc_score
 from torch.nn import Embedding, Linear
 
+import torch_geometric
 import torch_geometric.transforms as T
 from torch_geometric.datasets import Taobao
 from torch_geometric.loader import LinkNeighborLoader
 from torch_geometric.nn import SAGEConv
 from torch_geometric.utils.convert import to_scipy_sparse_matrix
 
-device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/Taobao')
 
 dataset = Taobao(path)

diff --git a/examples/hetero/dmgi_unsup.py b/examples/hetero/dmgi_unsup.py
@@ -9,6 +9,7 @@
 from sklearn.linear_model import LogisticRegression
 from torch.optim import Adam
 
+import torch_geometric
 import torch_geometric.transforms as T
 from torch_geometric.datasets import IMDB
 from torch_geometric.nn import GCNConv
@@ -74,7 +75,13 @@ def loss(self, pos_hs, neg_hs, summaries):
 
 model = DMGI(data['movie'].num_nodes, data['movie'].x.size(-1),
              out_channels=64, num_relations=len(data.edge_types))
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 data, model = data.to(device), model.to(device)
 
 optimizer = Adam(model.parameters(), lr=0.0005, weight_decay=0.0001)

diff --git a/examples/hetero/han_imdb.py b/examples/hetero/han_imdb.py
@@ -5,6 +5,7 @@
 import torch.nn.functional as F
 from torch import nn
 
+import torch_geometric
 import torch_geometric.transforms as T
 from torch_geometric.datasets import IMDB
 from torch_geometric.nn import HANConv
@@ -16,7 +17,6 @@
                            drop_unconnected_node_types=True)
 dataset = IMDB(path, transform=transform)
 data = dataset[0]
-print(data)
 
 
 class HAN(nn.Module):
@@ -34,7 +34,12 @@ def forward(self, x_dict, edge_index_dict):
 
 
 model = HAN(in_channels=-1, out_channels=3)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 data, model = data.to(device), model.to(device)
 
 with torch.no_grad():  # Initialize lazy modules.

diff --git a/examples/hetero/hetero_conv_dblp.py b/examples/hetero/hetero_conv_dblp.py
@@ -3,6 +3,7 @@
 import torch
 import torch.nn.functional as F
 
+import torch_geometric
 import torch_geometric.transforms as T
 from torch_geometric.datasets import DBLP
 from torch_geometric.nn import HeteroConv, Linear, SAGEConv
@@ -11,7 +12,6 @@
 # We initialize conference node features with a single one-vector as feature:
 dataset = DBLP(path, transform=T.Constant(node_types='conference'))
 data = dataset[0]
-print(data)
 
 
 class HeteroGNN(torch.nn.Module):
@@ -37,7 +37,12 @@ def forward(self, x_dict, edge_index_dict):
 
 model = HeteroGNN(data.metadata(), hidden_channels=64, out_channels=4,
                   num_layers=2)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 data, model = data.to(device), model.to(device)
 
 with torch.no_grad():  # Initialize lazy modules.

diff --git a/examples/hetero/hetero_link_pred.py b/examples/hetero/hetero_link_pred.py
@@ -5,6 +5,7 @@
 import torch.nn.functional as F
 from torch.nn import Linear
 
+import torch_geometric
 import torch_geometric.transforms as T
 from torch_geometric.datasets import MovieLens
 from torch_geometric.nn import SAGEConv, to_hetero
@@ -14,12 +15,7 @@
                     help='Whether to use weighted MSE loss.')
 args = parser.parse_args()
 
-if torch.cuda.is_available():
-    device = torch.device('cuda')
-elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
-    device = torch.device('mps')
-else:
-    device = torch.device('cpu')
+device = torch_geometric.device('auto')
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/MovieLens')
 dataset = MovieLens(path, model_name='all-MiniLM-L6-v2')

diff --git a/examples/hetero/hgt_dblp.py b/examples/hetero/hgt_dblp.py
@@ -3,6 +3,7 @@
 import torch
 import torch.nn.functional as F
 
+import torch_geometric
 import torch_geometric.transforms as T
 from torch_geometric.datasets import DBLP
 from torch_geometric.nn import HGTConv, Linear
@@ -11,7 +12,6 @@
 # We initialize conference node features with a single one-vector as feature:
 dataset = DBLP(path, transform=T.Constant(node_types='conference'))
 data = dataset[0]
-print(data)
 
 
 class HGT(torch.nn.Module):
@@ -43,7 +43,12 @@ def forward(self, x_dict, edge_index_dict):
 
 
 model = HGT(hidden_channels=64, out_channels=4, num_heads=2, num_layers=1)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 data, model = data.to(device), model.to(device)
 
 with torch.no_grad():  # Initialize lazy modules.

diff --git a/examples/hetero/hierarchical_sage.py b/examples/hetero/hierarchical_sage.py
@@ -4,18 +4,23 @@
 import torch.nn.functional as F
 from tqdm import tqdm
 
+import torch_geometric
 import torch_geometric.transforms as T
 from torch_geometric.datasets import OGB_MAG
 from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import HeteroConv, Linear, SAGEConv
 from torch_geometric.utils import trim_to_layer
 
 parser = argparse.ArgumentParser()
-parser.add_argument('--device', type=str, default='cuda')
 parser.add_argument('--use-sparse-tensor', action='store_true')
 args = parser.parse_args()
 
-device = args.device if torch.cuda.is_available() else 'cpu'
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 
 transforms = [T.ToUndirected(merge=True)]
 if args.use_sparse_tensor:

diff --git a/examples/hetero/metapath2vec.py b/examples/hetero/metapath2vec.py
@@ -4,6 +4,7 @@
 
 import torch
 
+import torch_geometric
 from torch_geometric.datasets import AMiner
 from torch_geometric.nn import MetaPath2Vec
 
@@ -18,7 +19,12 @@
     ('paper', 'written_by', 'author'),
 ]
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 model = MetaPath2Vec(data.edge_index_dict, embedding_dim=128,
                      metapath=metapath, walk_length=50, context_size=7,
                      walks_per_node=5, num_negative_samples=5,