Add XPU support to hetero examples

examples/hetero/bipartite_sage_unsup.py

@@ -13,8 +13,14 @@
 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')
+from torch_geometric.testing.device import is_xpu_avaliable
+
+if torch.cuda.is_available():
+    device = torch.device('cuda:0')
+elif is_xpu_avaliable():
+    device = torch.device('xpu:0')
+else:
+    device = torch.device('cpu')
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/Taobao')
 
 dataset = Taobao(path)

examples/hetero/dmgi_unsup.py

@@ -12,6 +12,7 @@
 import torch_geometric.transforms as T
 from torch_geometric.datasets import IMDB
 from torch_geometric.nn import GCNConv
+from torch_geometric.testing.device import is_xpu_avaliable
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/IMDB')
 dataset = IMDB(path)
@@ -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 is_xpu_avaliable():
+    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)

examples/hetero/han_imdb.py

@@ -8,6 +8,7 @@
 import torch_geometric.transforms as T
 from torch_geometric.datasets import IMDB
 from torch_geometric.nn import HANConv
+from torch_geometric.testing.device import is_xpu_avaliable
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/IMDB')
 metapaths = [[('movie', 'actor'), ('actor', 'movie')],
@@ -34,7 +35,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 is_xpu_avaliable():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 data, model = data.to(device), model.to(device)
 
 with torch.no_grad():  # Initialize lazy modules.

examples/hetero/hetero_conv_dblp.py

@@ -6,6 +6,7 @@
 import torch_geometric.transforms as T
 from torch_geometric.datasets import DBLP
 from torch_geometric.nn import HeteroConv, Linear, SAGEConv
+from torch_geometric.testing.device import is_xpu_avaliable
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/DBLP')
 # We initialize conference node features with a single one-vector as feature:
@@ -37,7 +38,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 is_xpu_avaliable():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 data, model = data.to(device), model.to(device)
 
 with torch.no_grad():  # Initialize lazy modules.

examples/hetero/hetero_link_pred.py

@@ -8,6 +8,7 @@
 import torch_geometric.transforms as T
 from torch_geometric.datasets import MovieLens
 from torch_geometric.nn import SAGEConv, to_hetero
+from torch_geometric.testing.device import is_xpu_avaliable
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--use_weighted_loss', action='store_true',
@@ -18,6 +19,8 @@
     device = torch.device('cuda')
 elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
     device = torch.device('mps')
+elif is_xpu_avaliable():
+    device = torch.device('xpu')
 else:
     device = torch.device('cpu')
 

examples/hetero/hgt_dblp.py

@@ -6,6 +6,7 @@
 import torch_geometric.transforms as T
 from torch_geometric.datasets import DBLP
 from torch_geometric.nn import HGTConv, Linear
+from torch_geometric.testing.device import is_xpu_avaliable
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/DBLP')
 # We initialize conference node features with a single one-vector as feature:
@@ -43,7 +44,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 is_xpu_avaliable():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 data, model = data.to(device), model.to(device)
 
 with torch.no_grad():  # Initialize lazy modules.

examples/hetero/hierarchical_sage.py

@@ -15,7 +15,12 @@
 parser.add_argument('--use-sparse-tensor', action='store_true')
 args = parser.parse_args()
 
-device = args.device if torch.cuda.is_available() else 'cpu'
+if args.device=='cuda' and torch.cuda.is_available():
+    device = 'cuda'
+elif args.device=='xpu' and is_xpu_avaliable():
+    device = 'xpu'
+else:
+    device = 'cpu'
 
 transforms = [T.ToUndirected(merge=True)]
 if args.use_sparse_tensor:

examples/hetero/metapath2vec.py

@@ -6,6 +6,7 @@
 
 from torch_geometric.datasets import AMiner
 from torch_geometric.nn import MetaPath2Vec
+from torch_geometric.testing.device import is_xpu_avaliable
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/AMiner')
 dataset = AMiner(path)
@@ -18,7 +19,12 @@
     ('paper', 'written_by', 'author'),
 ]
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+if torch.cuda.is_available():
+    device = 'cuda'
+elif is_xpu_avaliable():
+    device = 'xpu'
+else:
+    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,