Pytorch PyG实现EdgePool图分类

Pytorch Geometric(PyG)是一个用于图神经网络(GNN)的Pytorch库。EdgePool是一种PyG中的图池化操作，可以用于图分类任务中。下面是使用PyG实现EdgePool图分类任务的完整攻略。

环境配置

首先需要安装PyTorch和PyG，并使用pip安装以下库：

pip install scikit-learn matplotlib

数据集准备

我们将使用Cora数据集作为示例。Cora是一个由2708个文档组成的引文网络，每个文档由词向量表示。可以在此处下载Cora数据集：

wget https://github.com/kimiyoung/planetoid/raw/master/data/cora.zip
unzip cora.zip

数据集处理

我们将使用PyG的torch_geometric.datasets.Planetoid数据集处理类来加载Cora数据集。同时还需要定义一个DataLoader类以便于我们对数据进行采样和批处理。以下是代码示例：

from torch_geometric.datasets import Planetoid
from torch_geometric.data import DataLoader

dataset = Planetoid(root='cora', name='Cora')
loader = DataLoader(dataset, batch_size=64, shuffle=True)

模型定义

接下来，我们可以定义模型。这里我们将使用三个图卷积层(convolutional layers)和一个EdgePool层，最后使用全连接层将特征图转换为类别概率。以下是代码示例：

from torch.nn import Linear
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, EdgePooling

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(dataset.num_features, 64)
        self.conv2 = GCNConv(64, 64)
        self.pool = EdgePooling(64)
        self.conv3 = GCNConv(64, 64)
        self.fc = Linear(64, dataset.num_classes)

    def forward(self, x, edge_index, edge_attr):
        x = F.relu(self.conv1(x, edge_index))
        x = F.relu(self.conv2(x, edge_index))
        x, edge_index, edge_attr, batch, _ = self.pool(x, edge_index, edge_attr)
        x = F.relu(self.conv3(x, edge_index))
        x = global_max_pool(x, batch)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.fc(x)
        return F.log_softmax(x, dim=1)

模型训练

我们可以使用以下代码块来实例化模型并对模型进行训练：

from torch_geometric.utils import train_test_split_edges

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.NLLLoss()

for epoch in range(1, 201):
    model.train()
    optimizer.zero_grad()
    subgraph = train_test_split_edges(dataset[0])
    subgraph = subgraph.to(device)
    out = model(subgraph.x, subgraph.edge_index, subgraph.edge_attr)
    loss = criterion(out[subgraph.train_mask], subgraph.y[subgraph.train_mask])
    loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        model.eval()
        pred = model(subgraph.x, subgraph.edge_index, subgraph.edge_attr).argmax(dim=1)
        acc = int(pred[subgraph.test_mask].eq(subgraph.y[subgraph.test_mask]).sum()) / int(subgraph.test_mask.sum())
        print(f'Epoch: {epoch:03d}, Rail Loss: {loss:.4f}, Test Acc: {acc:.4f}')

模型评估

在模型训练完毕后，我们可以使用以下代码块来对模型进行评估：

model.eval()
test_subgraph = dataset[0].to(device)
pred = model(test_subgraph.x, test_subgraph.edge_index, test_subgraph.edge_attr).argmax(dim=1)
acc = int(pred[test_subgraph.test_mask].eq(test_subgraph.y[test_subgraph.test_mask]).sum()) / int(test_subgraph.test_mask.sum())
print(f'Test Accuracy: {acc:.4f}')

至此，我们已经完成了使用PyG实现EdgePool图分类任务的完整攻略。