Pytorch 如何创建图神经网络数据集（Pytorch Geometric）

在本文中，我们将介绍如何使用Pytorch Geometric库创建图神经网络（Graph Neural Network, GNN）的数据集。Pytorch Geometric是一个专门用于处理图数据的PyTorch扩展库，它提供了一些方便的工具和函数来处理和操作图数据。

阅读更多：Pytorch 教程

什么是图神经网络？

图神经网络是一种用于处理图数据的神经网络模型。与传统的神经网络模型不同，图神经网络能够处理非结构化的图数据，例如社交网络、知识图谱等。图神经网络通过学习节点和边之间的关系来进行信息传递和特征提取。

创建图数据集

PyTorch Geometric提供了创建各种图数据集的功能。下面我们将介绍如何创建一个简单的图数据集。

首先，我们需要导入必要的库和模块：

import torch
from torch_geometric.data import Dataset, Data

接下来，我们需要定义一个自定义的数据集类来创建数据集。我们可以继承PyTorch Geometric的Dataset类，并实现其中的方法：

class CustomGraphDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(CustomGraphDataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        # 返回原始文件的文件名列表，如果没有原始文件则返回空列表
        return []

    @property
    def processed_file_names(self):
        # 返回预处理文件的文件名列表，如果没有预处理文件则返回空列表
        return ['data.pt']

    def download(self):
        # 如果需要，可以在此处下载原始数据

    def process(self):
        # 在此处进行数据的预处理，并保存到self.processed_dir目录下
        data = Data(...)  # 根据需要创建图数据
        torch.save(data, self.processed_paths[0])

    def len(self):
        # 返回数据集中图的数量
        return 1

    def get(self, idx):
        # 从数据集中取出一个图数据
        data = torch.load(self.processed_paths[0])
        return data

在上面的代码中，我们定义了一个名为CustomGraphDataset的自定义数据集类。我们需要实现raw_file_names、processed_file_names、download、process、len和get这几个方法。

raw_file_names方法用于返回原始文件的文件名列表，如果我们的数据集有原始文件，可以在这里返回文件名列表，如果没有原始文件，可以返回空列表。
processed_file_names方法用于返回预处理文件的文件名列表，如果我们的数据集有预处理文件，可以在这里返回文件名列表，如果没有预处理文件，可以返回空列表。
download方法用于下载原始数据，如果我们的数据集需要下载原始数据，可以在这里实现下载逻辑。
process方法用于进行数据的预处理，我们需要在这里创建图数据，并保存到self.processed_dir目录下。
len方法用于返回数据集中图的数量。
get方法用于从数据集中取出一个图数据，返回一个Data对象。

完成了自定义的数据集类后，我们可以创建一个实例来使用这个数据集了：

dataset = CustomGraphDataset(root='data')
data = dataset[0]

print(data)

上面的代码中，我们首先创建了一个CustomGraphDataset的实例，并指定了数据集的根目录为’data’。然后，我们通过索引的方式获取第一个图数据，并将其保存到data变量中。最后，我们打印出这个图数据。

数据集示例

接下来，让我们来看一个实际的例子，如何使用PyTorch Geometric库创建一个图数据集。

假设我们要创建一个用于分类任务的图数据集。每个图由一组节点和它们之间的边组成，每个节点都带有一些特征。我们需要为每个节点和边定义特征，并为每个节点和边指定一个目标标签。

首先，我们导入需要的库和模块：

import torch
import torch.nn.functional as F
from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader
from torch_geometric.nn import GCNConv, global_add_pool

接下来，我们可以从PyTorch Geometric提供的预定义图数据集中加载数据，例如TUDataset。TUDataset是一个常用的图分类数据集，包含了很多图和相应的分类标签。

dataset = TUDataset(root='data', name='ENZYMES', pre_transform=None)

在上述代码中，我们创建了一个名为ENZYMES的图分类数据集，并将数据集保存到根目录为'data'的位置。我们可以通过设置pre_transform参数来应用图数据的预处理，例如对节点特征进行标准化等。

接下来，我们需要定义一个图神经网络模型来处理图数据。这里我们使用GCNConv层作为图卷积层，并结合global_add_pool函数对全局特征进行池化。

class GCN(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(num_features, 128)
        self.conv2 = GCNConv(128, num_classes)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        x = global_add_pool(x, batch)
        return F.log_softmax(x, dim=1)

model = GCN(dataset.num_features, dataset.num_classes)

在上述代码中，我们定义了一个名为GCN的图神经网络模型，其中包含了两个GCNConv层。在forward方法中，我们定义了前向传播的操作，先通过第一个GCNConv层获取节点的局部特征，然后使用ReLU激活函数进行非线性变换，再通过第二个GCNConv层获取节点的全局特征，并使用global_add_pool函数对全局特征进行池化。最后，使用log_softmax函数将模型的输出转化为概率。

接下来，我们需要定义一些训练参数，并将数据集划分为训练集、验证集和测试集：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = dataset.shuffle()
train_dataset = dataset[:150]
val_dataset = dataset[150:200]
test_dataset = dataset[200:]
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

在上述代码中，我们首先检查是否有可用的GPU，如果有则使用GPU进行训练，否则使用CPU。然后，我们使用shuffle函数将数据集打乱，确保训练数据和验证数据的随机性。接着，我们将数据集划分为训练集、验证集和测试集，并使用DataLoader将数据集分批加载。

最后，我们可以定义一些训练和评估的函数，并进行模型训练和评估：

def train(model, loader, optimizer):
    model.train()
    for data in loader:
        data = data.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, data.y)
        loss.backward()
        optimizer.step()

def evaluate(model, loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data in loader:
            data = data.to(device)
            output = model(data)
            _, predicted = torch.max(output, 1)
            correct += (predicted == data.y).sum().item()
            total += data.num_nodes

    accuracy = correct / total
    return accuracy

optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
num_epochs = 100

model.to(device)
best_val_accuracy = 0.0

for epoch in range(1, num_epochs + 1):
    train(model, train_loader, optimizer)
    train_accuracy = evaluate(model, train_loader)
    val_accuracy = evaluate(model, val_loader)
    print(f'Epoch: {epoch}, Train Accuracy: {train_accuracy}, Val Accuracy: {val_accuracy}')

    if val_accuracy > best_val_accuracy:
        best_val_accuracy = val_accuracy
        torch.save(model.state_dict(), 'best_model.pt')

test_accuracy = evaluate(model, test_loader)
print(f'Test Accuracy: {test_accuracy}')

在上述代码中，我们首先定义了train和evalute函数分别用于模型的训练和评估。在训练函数中，我们首先将模型设置为训练模式，然后将每个小批量的数据移动到设备上，在优化器的帮助下计算模型的输出和损失，并进行反向传播和权重更新。在评估函数中，我们将模型设置为评估模式，然后使用无梯度的上下文计算模型在验证集或测试集上的准确率。

然后，我们定义了优化器和训练的一些超参数，将模型移动到设备上，并初始化最佳验证准确率为0。

接下来，我们通过多个训练周期来训练模型，并在每个周期后计算训练集和验证集的准确率。如果验证准确率提升，则保存当前模型的参数。

最后，我们使用测试集评估训练得到的模型，并输出测试准确率。

总结

本文介绍了如何使用PyTorch Geometric库创建图神经网络数据集。我们首先了解了图神经网络的概念，并详细介绍了如何使用PyTorch Geometric库创建自定义的图数据集。然后，我们以一个分类任务为例，演示了如何使用PyTorch Geometric库加载预定义的图分类数据集，并使用图神经网络模型进行训练和评估。

希望本文对于学习PyTorch Geometric库和图神经网络的使用有所帮助。通过掌握创建图数据集的技巧，我们可以更好地处理和分析图数据，并在图神经网络任务中取得更好的效果。