DGL类库——用GCN 跑下CORA数据集的节点分类任务

内容纲要

import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.data import CoraGraphDataset
from torch.optim import Adam

# 加载 CORA 数据集
data = CoraGraphDataset()
g = data[0]

# 为了简便，我们直接使用图中的特征和标签
features = g.ndata['feat']
labels = g.ndata['label']
train_mask = g.ndata['train_mask']
val_mask = g.ndata['val_mask']
test_mask = g.ndata['test_mask']

n_features = features.shape[1]
n_labels = len(torch.unique(labels))

# 定义 GCN 层
class GCNLayer(nn.Module):
    def __init__(self, in_feats, out_feats):
        super(GCNLayer, self).__init__()
        self.linear = nn.Linear(in_feats, out_feats)

    def forward(self, g, feature):
        with g.local_scope():
            g.ndata['h'] = feature
            g.update_all(dgl.function.copy_u('h', 'm'),
                         dgl.function.mean('m', 'h'))
            h = g.ndata['h']
            return self.linear(h)

# 定义 GCN 模型
class GCN(nn.Module):
    def __init__(self, in_feats, hidden_size, num_classes):
        super(GCN, self).__init__()
        self.gcn1 = GCNLayer(in_feats, hidden_size)
        self.gcn2 = GCNLayer(hidden_size, num_classes)

    def forward(self, g, features):
        x = F.relu(self.gcn1(g, features))
        x = self.gcn2(g, x)
        return x

# 初始化模型
model = GCN(n_features, 16, n_labels)
model.load_state_dict(torch.load('/Users/davidlee/Downloads/dgl.pth'))
# 定义优化器
optimizer = Adam(model.parameters(), lr=0.01)

# 定义损失函数
loss_fn = nn.CrossEntropyLoss()

# 训练模型
epochs = 200
for epoch in range(epochs):
    model.train()
    # 前向传播
    logits = model(g, features)
    # 计算损失值
    loss = loss_fn(logits[train_mask], labels[train_mask])

    # 后向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 每隔一段时间输出训练信息
    if epoch % 10 == 0:
        # 计算验证集上的准确率
        model.eval()
        with torch.no_grad():
            logits = model(g, features)
            _, predicted = torch.max(logits[val_mask], 1)
            correct = (predicted == labels[val_mask]).sum().item()
            accuracy = correct / val_mask.sum().item()
        print(f'Epoch {epoch}, Loss: {loss.item():.4f}, Val Acc: {accuracy:.4f}')

# 测试模型
model.eval()
with torch.no_grad():
    logits = model(g, features)
    _, predicted = torch.max(logits[test_mask], 1)
    correct = (predicted == labels[test_mask]).sum().item()
    accuracy = correct / test_mask.sum().item()
print(f'Test Accuracy: {accuracy:.4f}')
torch.save(model.state_dict(), '/Users/davidlee/Downloads/dgl.pth')

这个脚本是用于在CORA数据集上训练和测试一个简单的图卷积网络（GCN）模型，逐步解释如下：

导入必要的库：
- import dgl: 导入DGL（Deep Graph Library），一个用于图神经网络的Python库。
- import torch: 导入PyTorch，一个流行的深度学习框架。
- import torch.nn as nn 和 import torch.nn.functional as F: 导入PyTorch中的神经网络模块和函数。
- from dgl.data import CoraGraphDataset: 从DGL库中导入CoraGraphDataset类，用于加载CORA数据集。
- from torch.optim import Adam: 导入Adam优化器。
加载CORA数据集：
- data = CoraGraphDataset(): 加载CORA数据集。
- g = data[0]: 获取数据集中的图。
提取图数据的特征和标签：
- 特征(features), 标签(labels), 训练掩码(train_mask), 验证掩码(val_mask), 测试掩码(test_mask)。这些用于在训练、验证和测试过程中索引相应的节点。
获取特征和类别的数量：
- 通过特征的形状和标签的唯一值计算。
定义GCN层：GCNLayer类扩展了nn.Module，使用线性变换和图的邻居聚合来实现图卷积。
定义GCN模型：GCN类包含两个GCNLayer层，使用ReLU激活函数连接。
模型初始化和状态加载：
- 初始化GCN模型，设置隐藏层大小为16，并加载预先训练的模型状态（如果有）。
定义优化器和损失函数：
- 使用Adam优化器和交叉熵损失函数。
模型训练：
- 迭代训练模型200轮，每轮包括前向传播、损失计算、后向传播和参数更新。每10轮输出一次训练信息。
模型验证：
- 在每个输出步骤，计算并打印验证集上的准确率。
模型测试：
- 在训练完成后，评估模型在测试集上的性能，并输出测试准确率。
保存模型状态：
- 训练结束后，保存模型的状态，以便将来使用或进一步训练。

发表回复取消回复