详解使用Pytorch Geometric实现GraphSAGE模型

我们来详细讲解一下使用 Pytorch Geometric 实现 GraphSAGE 模型的完整攻略。

1. 什么是 GraphSAGE 模型?

GraphSAGE 是一个用于图像分类的模型，其主要思想是对于每一个节点，利用其周围的节点的嵌入向量来产生一个向量来描述该节点。这个向量可以作为分类器的输入。为了实现这个思想，GraphSAGE模型主要包含两个部分：

• 邻居采样: 采样图中与该节点最近的 k 个节点，最终形成一个子图。

• 对每个子图进行嵌入: 根据子图嵌入节点，产生每个节点的嵌入向量。

这样，我们就可以用这些嵌入向量来训练分类器。

2. 使用 Pytorch Geometric 实现 GraphSAGE 模型

2.1 安装 Pytorch Geometric

安装 Pytorch Geometric 可以使用 pip 命令进行安装：

$ pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.9.0+cpu.html

这个过程可能会比较慢，因为需要下载一些依赖包。

2.2 加载数据集

这里我们以 Cora 数据集为例。首先，我们需要加载数据集。

from torch_geometric.datasets import Planetoid

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

2.3 定义 GraphSAGE 模型

接下来，我们定义 GraphSAGE 模型。我们需要定义 GraphSAGE 层以及分类器。

这里我们以两层 GraphSAGE 层为例：

from torch_geometric.nn import SAGEConv

class SAGEModel(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(SAGEModel, self).__init__()

        self.conv1 = SAGEConv(in_channels, hidden_channels)
        self.conv2 = SAGEConv(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 x

2.4 训练模型

接下来，我们需要定义损失函数以及优化器，并使用训练集进行训练。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SAGEModel(dataset.num_features, 16, dataset.num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.CrossEntropyLoss()

def train():
    model.train()

    optimizer.zero_grad()
    out = model(data.x.to(device), data.edge_index.to(device))

    loss = criterion(out[data.train_mask], data.y[data.train_mask].to(device))
    loss.backward()
    optimizer.step()

    return loss

for epoch in range(1, 201):
    loss = train()
    print('Epoch: {:03d}, Loss: {:.5f}'.format(epoch, loss))

2.5 在测试集上评估模型

最后，我们需要在测试数据集上评估模型的性能。

def test():
    model.eval()

    out = model(data.x.to(device), data.edge_index.to(device))

    pred = out.argmax(dim=1)
    acc = pred[data.test_mask].eq(data.y[data.test_mask]).sum().item() / data.test_mask.sum().item()

    return acc

test_acc = test()
print('Test Accuracy: {:.5f}'.format(test_acc))

至此，我们已经完成 GraphSAGE 模型的训练和测试，也就实现了 Pytorch Geometric 上使用 GraphSAGE 模型进行图像分类的示例。

3. 示例说明

3.1 示例 1: 加载其他数据集

我们可以加载其他的数据集进行实验，只需在与 Cora 数据集相同的格式中提供节点特征、节点标签和边数据即可。Pytorch Geometric 提供了许多其他数据集。

from torch_geometric.datasets import DatasetName

dataset = DatasetName(root='/tmp/DatasetName', name='DatasetName')

3.2 示例 2: 使用其他嵌入方法

GraphSAGE 模型中使用的是 Mean Aggregation 的方法进行嵌入，还可以使用其他嵌入方法。例如，我们可以使用 GCN Aggregation 的方法进行嵌入：

from torch_geometric.nn import GCNConv

class GCNModel(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(GCNModel, 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 x

这样，我们可以用这个嵌入方法来训练 GraphSAGE 模型。