GNN模型训练

系列文章目录

第一章 图神经网络GNN实践入门
第二章 GNN模型输入数据处理方法

---

---

前言

本篇是基于处理好的数据集，通过定义模型的网络结构，将数据放入模型中进行训练。

• GNN网络结构定义模块
• 获取全局特征
• 模型训练与总结

---

提示：以下是本篇文章正文内容，下面案例可供参考

一、模型结构

API文档解释：

• https://pytorchgeometric.readthedocs.io/en/latest/generated/torch_geometric.nn.pool.TopKPooling.html

TokPooling流程如下图所示：

• 相当于对图进行剪枝操作，选择分低的节点剔除掉，然后再重新组成一个新的图

二、代码实现

1.模型定义部分

embed_dim = 128
import torch
from torch_geometric.nn import TopKPooling, SAGEConv
from torch_geometric.nn import global_mean_pool as gap, global_max_pool
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

• embed_dim = 128：定义了嵌入维度，表示节点特征的维度。
• device 定义了使用的设备，优先选择CUDA（GPU），如果不可用则使用CPU。

2.神经网络模型定义

代码如下：

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        self.conv1 = SAGEConv(embed_dim, 128)
        self.pool1 = TopKPooling(128, ratio=0.8)
        self.conv2 = SAGEConv(128, 128)
        self.pool2 = TopKPooling(128, ratio=0.8)
        self.conv3 = SAGEConv(128, 128)
        self.pool3 = TopKPooling(128, ratio=0.8)
        self.item_embedding = torch.nn.Embedding(num_embeddings=df.item_id.max() + 10, embedding_dim=embed_dim)
        self.lin1 = torch.nn.Linear(128, 128)
        self.lin2 = torch.nn.Linear(128, 64)
        self.lin3 = torch.nn.Linear(64, 1)
        self.bn1 = torch.nn.BatchNorm1d(128)
        self.bn2 = torch.nn.BatchNorm1d(64)
        self.act1 = torch.nn.ReLU()
        self.act2 = torch.nn.ReLU()

• SAGEConv：一种图卷积网络层（Graph Convolutional Layer），用于提取节点的特征。
• TopKPooling：一种池化层，用于对图进行下采样。
• item_embedding：嵌入层，将节点ID映射到高维特征向量。
• Linear 层：全连接层，用于分类任务。
• BatchNorm1d 和 ReLU 激活函数：用于加速训练和提高模型性能。

3、前向传播方法

def forward(self, data):
    x, edge_index, batch = data.x, data.edge_index, data.batch
    print(x)
    x = self.item_embedding(x).squeeze(1)
    print('item_embedding', x.shape)
    print(f'edge_index: {data.edge_index}')
    print(f'edge_index max value: {data.edge_index.max()}')
    print(f'num_nodes: {data.num_nodes}')
    print(f'node feature matrix shape: {data.x.shape}')

• data 包含图的节点特征矩阵 x，边索引 edge_index 和批次信息 batch。
• item_embedding 将节点特征嵌入到高维空间，并通过 squeeze(1) 去掉多余的维度。

    x = F.relu(self.conv1(x, edge_index))
    print('conv1', x.shape)
    x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, None, batch)
    print('self.pool1', x.shape)
    print('self.pool1', edge_index)
    print('self.pool1', batch)
    x1 = gap(x, batch)
    print('x1', x1.shape)

• 使用 SAGEConv 和 ReLU 激活函数对节点特征进行卷积。
• 使用 TopKPooling 对图进行下采样。
• gap（全局平均池化）将每个图的节点特征聚合成一个特征向量。

    x = F.relu(self.conv2(x, edge_index))
    x, edge_index, _, batch, _, _ = self.pool2(x, edge_index, None, batch)
    print('self.pool2', x.shape)
    print('self.pool2', edge_index)
    print('self.pool2', batch)
    x2 = gap(x, batch)
    print('x2', x2.shape)

• 重复上述步骤，对卷积和池化后的特征进行进一步处理，得到 x2。

    x = F.relu(self.conv3(x, edge_index))
    x, edge_index, _, batch, _, _ = self.pool3(x, edge_index, None, batch)
    print('self.pool3', x.shape)
    print('self.pool3', edge_index)
    print('self.pool3', batch)
    x3 = gap(x, batch)
    print('x3', x3.shape)

• 再次重复卷积和池化步骤，得到 x3。

    x = x1 + x2 + x3

    x = self.lin1(x)
    x = self.act1(x)
    x = self.lin2(x)
    x = self.act2(x)
    x = F.dropout(x, p=0.5, training=self.training)

    x = torch.sigmoid(self.lin3(x)).squeeze(1)
    print('sigmoid', x.shape)
    return x

• 将 x1, x2, x3 相加，得到最终的特征向量。
• 通过全连接层、激活函数和dropout进行进一步处理。
• 最后一层使用 sigmoid 激活函数得到输出，并通过 squeeze(1) 去掉多余的维度。

4、模型训练

from torch_geometric.loader import DataLoader

def train():
    model.train()

	#初始化一个变量 loss_all，用于累积每个批次的损失值。
    loss_all = 0
    #对训练数据加载器中的每个数据批次进行迭代。
    for data in train_loader:
        data = data.to(device)
        #对优化器进行零梯度操作，以确保之前计算的梯度不会影响当前的计算。
        optimizer.zero_grad()
        #将数据批次输入到模型中进行前向传播，得到模型的输出。
        output = model(data)
        label = data.y
        #使用预定义的损失函数 crit 计算模型输出与标签之间的损失值。
        loss = crit(output, label)
        #执行反向传播，计算模型参数的梯度。
        loss.backward()
        # 累加当前批次的损失值
        loss_all += data.num_graphs * loss.item()
        #根据计算得到的梯度更新模型的参数。
        optimizer.step()
    return loss_all / len(dataset)

• loss.item() 返回的是单个样本的损失值，所以乘以当前批次图的数量 data.num_graphs 可以得到整个批次的损失值。

model = Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = 0.001)
crit = torch.nn.BCELoss()
train_loader = DataLoader(dataset, batch_size=64)
for epoch in range(10):
    print('epoch', epoch)
    loss = train()
    print(loss)

• optimizer = torch.optim.Adam(model.parameters(), lr=0.001): 定义优化器，这里使用 Adam 优化器，用于调整模型参数以最小化损失函数。model.parameters() 提供了模型中需要更新的参数，而 lr=0.001 指定了学习率。
• crit = torch.nn.BCELoss(): 定义损失函数，这里使用了二分类交叉熵损失函数 BCELoss()，用于衡量模型输出与真实标签之间的差异。
• train_loader = DataLoader(dataset, batch_size=64): 创建一个数据加载器 DataLoader，用于批量加载训练数据。batch_size=64 指定了每个批次中包含的样本数量。

5、模型性能评估

• 该函数的目的是在给定的数据加载器上评估模型的性能，并返回 ROC AUC 分数作为评估指标。

from sklearn.metrics import roc_auc_score

def evaluate(loader, model):
    model.eval()

    prediction = []
    labels = []

    with torch.no_grad():
        for data in loader:
            data = data.to(device)
            pred = model(data)

            label = data.y.detach()
            prediction.append(pred)
            labels.append(label)
        prediction = np.hstack(prediction)
        labels = np.hstack(labels)

        return roc_auc_score(labels, prediction)

---

总结

以上就是真实数据集放入图神经网络训练的全部过程，今后可以将这个思想应用于自己领域的数据集中。