Pytorch 如何将Cora数据集分割以仅在训练集上训练GCN模型

在本文中，我们将介绍如何使用PyTorch将Cora数据集分割为训练集和测试集，并且仅在训练集上训练Graph Convolutional Network（GCN）模型。Cora是一个常用的图数据集，包含了2708个科学出版物的分类信息，作为学术领域图神经网络的常用数据集之一。

阅读更多：Pytorch 教程

什么是Cora数据集？

Cora数据集是一个文献引用网络数据集，包含了由2708篇论文组成的图结构，每篇论文由一个特征向量表示，特征维度为1433。此外，每篇论文还被标记了7个不同的类别，包括case-based, genetic algorithms, neural networks等。Cora数据集用于研究节点分类、图表示学习等任务。

数据集分割

在训练GCN模型时，通常将Cora数据集以一定比例划分为训练集、验证集和测试集。在本文中，我们只关注训练集部分，将其作为我们GCN模型训练的数据。

首先，我们需要下载Cora数据集。可以通过以下方式下载：

import torch
from torch_geometric.datasets import Planetoid

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

下载完成后，我们可以将Cora数据集划分为训练集和测试集。以下是一个将数据集划分为训练集和测试集的示例代码：

from torch_geometric.data import DataLoader
from torch_geometric.utils import train_test_split_edges

data = dataset[0]

# 划分数据集为训练集和测试集
data.train_mask = torch.zeros(data.num_nodes, dtype=torch.bool)
data.train_mask[:data.num_nodes//2] = 1  # 手动指定前一半节点为训练集

# 将训练集和测试集划分为边集合
data = train_test_split_edges(data)

# 创建数据加载器
loader = DataLoader(data, batch_size=batch_size, shuffle=True)

这段代码首先将数据集的训练集节点标记为True，然后使用train_test_split_edges函数将训练集和测试集划分为边集合。最后，我们使用DataLoader创建一个数据加载器。

GCN模型训练

接下来，我们将使用划分好的训练集数据来训练GCN模型。

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

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

model = GCN(in_channels=dataset.num_features, hidden_channels=16, out_channels=dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.train_edge_index)
    loss = F.nll_loss(out[data.train_mask], data.train_y)
    loss.backward()
    optimizer.step()

for epoch in range(200):
    train()

以上代码定义了一个两层GCN模型，并使用随机梯度下降优化器进行训练。在每个epoch中，我们都调用train函数进行训练。

总结

本文介绍了如何使用PyTorch将Cora数据集分割为训练集和测试集，并且仅在训练集上训练GCN模型。首先，我们下载了Cora数据集，并将其划分为训练集和测试集。然后，我们定义了一个GCN模型，并使用训练集数据进行训练。通过本文的指导，希望读者能够更好地理解如何使用PyTorch处理和训练图数据集。