torch_geometric 笔记： 数据集Cora &简易 GNN

Posted 2021-10-24 UQI-LIUWJ

1 获取数据集

该数据集用于semi-supervised的节点分类任务

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/tmp/Cora', name='Cora')

dataset.num_classes
#7
#节点一共七个类

dataset.num_features
#1433
#每个点1433个特征

len(dataset)
#1
#只有一张图

dataset[0].is_undirected()
#True

dataset[0]
#Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])

'''
edge_index=[2, 10556]————这张图有10556/2=5278条无向边
x=[2708, 1433]————这张图有2708个点，每个点1433条边
y=[2708]——每个节点的标签
'''



dataset[0]['train_mask']
#tensor([ True,  True,  True,  ..., False, False, False])
#train_mask：2708维向量，训练集的mask向量，标识哪些节点属于训练集。
#val_mask：2708维向量，验证集的mask向量，标识哪些节点属于验证集。
#test_mask：2708维向量，测试集的mask向量，表示哪些节点属于测试集。

2 简易GCN

2.1 torch_geometric.nn中有的模型

 

在torch_geometric.nn — pytorch_geometric 2.0.1 documentation (pytorch-geometric.readthedocs.io)

列举了torch_geometric.nn中有的模型

2.2 简易模型

2.2.1 导入库

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

2.2.2 设计模型

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        #两层GCN，输入是每个节点的num_node_features维特征，输出是16维向量
        self.conv2 = GCNConv(16, dataset.num_classes)
        #两层GCN，输入是16维向量，输出是点有的类别数

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        #GCN1
        '''
        forward(
            x: torch.Tensor, 
            edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor],
            edge_weight: Optional[torch.Tensor] = None) 
        → torch.Tensor
        '''
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)

2.2.3 训练模型

model = Net()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
#优化函数

loss_func=F.nll_loss

model.train()
for epoch in range(200):
    optimizer.zero_grad()
#清空上一步残余的参数更新值

    out = model(data)

    loss = loss_func(out[data.train_mask], data.y[data.train_mask])
#计算误差

    loss.backward()
#清空上一步残余的参数更新值

    optimizer.step()
 #将参数更新值施加到net的parameters上

 2.2.4 测试模型

model.eval()
_, pred = model(data).max(dim=1)
#预测结果

correct = int(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / int(data.test_mask.sum())
print('Accuracy: {:.4f}'.format(acc))
#Accuracy: 0.8080