GraphSAGE项目练手

1. # 导包
2. from torch_geometric.datasets import Planetoid
3. from torch_geometric.loader import NeighborLoader
4. from torch_geometric.utils import to_networkx
5. import numpy as np
6. import networkx as nx
7. import matplotlib.pyplot as plt
8. import torch
9. import torch.nn.functional as F
10. from torch_geometric.nn import SAGEConv
11. # 导入PubMed数据集
12. dataset = Planetoid(root='',name='Pubmed')
13. data = dataset[0]
15. # 邻居采样
16. # 使用NeighborLoader 来完成这一任务。
17. # 保留目的节点的10个邻居 和 其邻居的10个邻居， 对60个目的节点进行分组，每16个目的节点为一组
19. # 进行采样
20. train_loader = NeighborLoader(
21.     data,# 数据源
22.     num_neighbors=[5,10], # 每一层采样的邻居采样量，第一层5，第二层10
23.     batch_size=16,
24.     input_nodes=data.train_mask # 60个训练目的节点
25. )
26. # 遍历数据检验
27. # for i,subgraph in enumerate(train_loader):
28. #     print(f'Subgraph{i}:{subgraph}')
30. # 子图可视化
31. # fig = plt.figure(figsize=(16,16))
32. # for idx,(subdata,pos) in enumerate(zip(train_loader,[221,222,223,224])):
33. #     G = to_networkx(subdata,to_undirected=True)
34. #     ax = fig.add_subplot(pos)
35. #     ax.set_title(f'Subgraph{idx},fonts=24')
36. #     plt.axis('off')
37. #     nx.draw_networkx(G,pos=nx.spring_layout(G),with_labels=False,node_color=subdata.y)
38. # plt.show()
40. # 实现准确率评估模型
41. def  accuracy(pre_y,y):
42.     return ((pre_y==y).sum() / len(y)).item()
44. # 定义GraphSAGE
45. class GraphSAGE(torch.nn.Module):
46.     def __init__(self,dim_in,dim_h,dim_out):
47.         super().__init__()
48.         self.sage1 = SAGEConv(dim_in,dim_h)
49.         self.sage2= SAGEConv(dim_h,dim_out)
51.     def forward(self,x,edge_index):
52.         h = self.sage1(x,edge_index)
53.         h = torch.relu(h)
54.         h = F.dropout(h,p=0.5,training=self.training)
55.         h = self.sage2(h,edge_index)
56.         return h
57. # 使用小批量训练，Fit函数要修改为先循环epoch次，然后循环批数据，以在每个批数据上训练epoch次
58.     def fit(self,loader,epochs):
59.         criterion = torch.nn.CrossEntropyLoss()
60.         optimizer = torch.optim.Adam(self.parameters(),lr=0.01)
61.         self.train()
62.         for epoch in range(epochs+1):
63.             total_loss = 0
64.             acc = 0
65.             val_loss = 0
66.             val_acc = 0
67.             for batch in loader:
68.                 optimizer.zero_grad()
69.                 out = self(batch.x, batch.edge_index)
70.                 loss = criterion(out[batch.train_mask], batch.y[batch.train_mask])
71.                 total_loss += loss.item()
72.                 acc += accuracy(out[batch.train_mask].argmax(dim=1), batch.y[batch.train_mask])
73.                 loss.backward()
74.                 optimizer.step()
76.                 # Validation
77.                 val_loss += criterion(out[batch.val_mask], batch.y[batch.val_mask])
78.                 val_acc += accuracy(out[batch.val_mask].argmax(dim=1), batch.y[batch.val_mask])
80.                 if epoch % 20 == 0:
81.                     print(f'Epoch {epoch:>3} | Train Loss: {loss/len(loader):.3f} | Train Acc: {acc/len(loader)*100:>6.2f}% | Val Loss: {val_loss/len(train_loader):.2f} | Val Acc: {val_acc/len(train_loader)*100:.2f}%')
82. @torch.no_grad()
83. def test(self, data):
84.     self.eval()
85.     out = self(data.x, data.edge_index)
86.     acc = accuracy(out.argmax(dim=1)[data.test_mask], data.y[data.test_mask])
87.     return acc
89. # Create GraphSAGE
90. graphsage = GraphSAGE(dataset.num_features, 64, dataset.num_classes)
91. print(graphsage)
93. # Train
94. graphsage.fit(train_loader, 200)

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。