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

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条有向边
x=[2708, 1433]————这张图有2708个点，每个点1433个特征
y=[2708]——每个节点的标签（一共有7个类）
'''



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

 

1.1 cora 数据集  

cora数据集的点表示的是机器学习的论文， 这些论文的选择方式使得在最终的语料库中每篇论文都引用或被至少另一篇论文引用。

全语料库有2708篇论文。我们得到了一个大小为 1433 个唯一词的词汇表。 所有文档频率小于 10 的单词都被删除。

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