Merge pull request #58 from MooreThreads/add_geometric

feat(pytorch-geometric): add a GNN train demo for geometric

Author: caizhi-mt

pytorch/pytorch_geometric/README.md

@@ -0,0 +1,5 @@
+
+测试脚本，下载数据集默认会从github上下载，可能耗时较长。
+```bash
+python test_train_GNN_demo.py
+```

pytorch/pytorch_geometric/test_train_GNN_demo.py

@@ -0,0 +1,35 @@
+import torch
+import torch_musa
+from torch import Tensor
+from torch_geometric.nn import GCNConv
+from torch_geometric.datasets import Planetoid
+import torch.nn.functional as F
+
+dataset = Planetoid(root='.', name='Cora')
+
+class GCN(torch.nn.Module):
+    def __init__(self, in_channels, hidden_channels, out_channels):
+        super().__init__()
+        self.conv1 = GCNConv(in_channels, hidden_channels)
+        self.conv2 = GCNConv(hidden_channels, out_channels)
+
+    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
+        # x: Node feature matrix of shape [num_nodes, in_channels]
+        # edge_index: Graph connectivity matrix of shape [2, num_edges]
+        x = self.conv1(x, edge_index).relu()
+        x = self.conv2(x, edge_index)
+        return x
+
+model = GCN(dataset.num_features, 16, dataset.num_classes).to('musa')
+
+data = dataset[0]
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+for epoch in range(200):
+    pred = model(data.x.to('musa'), data.edge_index.to('musa'))
+    loss = F.cross_entropy(pred[data.train_mask], data.y[data.train_mask].to('musa'))
+
+    # Backpropagation
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()