Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions pytorch/pytorch_geometric/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@

测试脚本,下载数据集默认会从github上下载,可能耗时较长。
```bash
python test_train_GNN_demo.py
```
35 changes: 35 additions & 0 deletions pytorch/pytorch_geometric/test_train_GNN_demo.py
Original file line number Diff line number Diff line change
@@ -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()