基于图卷积神经网络GCN的二部图链路预测方法实现思路和完整代码【可用于疾病-靶点、miRNA-疾病等相关预测】

本文解决的问题

基本假设：二部图只有邻接矩阵，没有节点特征，并且进行链路预测，有部分链路未知。
如果你有初始节点特征，直接换掉即可

实现思路

这段代码主要是用于构建一个基于图卷积神经网络（GCN）的二分类模型，实现对给定邻接矩阵的节点间关系进行分类。具体实现步骤如下：
加载邻接矩阵数据，并将邻接矩阵转换为张量形式。
根据邻接矩阵生成正负样本的边索引。
自己创建节点特征x, 和结果y。
将数据随机分为训练、验证和测试集。
从正负样本中随机选取相同数量的边，形成平衡的训练集。
创建一个GCN模型，用于学习节点特征和边信息之间的关系。
通过训练得到最优的模型参数，并在测试集上评估模型性能。

结果

代码全览

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.utils import to_undirected
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score

# 加载邻接矩阵数据
adj_matrix = pd.read_csv('m_d.csv', header=None, index_col=None).values

# 将邻接矩阵转换为张量，这里要将列索引加上总节点数，使得列节点的索引不与行节点重复
num_nodes = adj_matrix.shape[0] + adj_matrix.shape[1]
edge_index_pos = np.column_stack(np.argwhere(adj_matrix != 0))
edge_index_pos[:, 1] += adj_matrix.shape[0]
edge_index_pos = torch.tensor(edge_index_pos, dtype=torch.long)

# 获取不相关连接 为了构建训练数据
edge_index_neg = np.column_stack(np.argwhere(adj_matrix == 0))
edge_index_neg[:, 1] += adj_matrix.shape[0]
edge_index_neg = torch.tensor(edge_index_neg, dtype=torch.long)

# 获取平衡样本
num_pos_edges_number = edge_index_pos.shape[1]
selected_neg_edge_indices = torch.randint(high=edge_index_neg.shape[1], size=(num_pos_edges_number,), dtype=torch.long)
edge_index_neg_selected = edge_index_neg[:, selected_neg_edge_indices]

edg_index_all = torch.cat((edge_index_pos, edge_index_neg_selected), dim=1)

# 创建数据
x = torch.ones((num_nodes, 1)) # 没有节点特征，所以设置为1
y = torch.cat((torch.ones((edge_index_pos.shape[1], 1)),
               torch.zeros((edge_index_neg_selected.shape[1], 1))), dim=0) # 将所有y值设置为1,0

# 将邻接矩阵转换为无向图（会将每条边复制成两条相反的边）
edge_index = to_undirected(edge_index_pos)

# 将数据拆分为训练、验证和测试集
idx = np.arange(y.shape[0])
np.random.shuffle(idx)
train_idx = idx[:int(0.8 * len(idx))] # 用前80%的样本作为训练集
test_idx = idx[int(0.8 * len(idx)):] # 用后20%的样本作为测试集


# 创建一个GCN模型
class GCNII(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(GCNII, self).__init__()
        self.conv1 = GCNConv(in_channels, 2 * out_channels, cached=True)
        self.conv2 = GCNConv(2 * out_channels, out_channels)
        self.fc = torch.nn.Linear(out_channels*2, 1)

    def forward(self, x, edge_index, edge_index_pos, edge_index_neg, edge_weight=None,param={}):
        x = F.relu(self.conv1(x, edge_index, edge_weight))
        x = F.relu(self.conv2(x, edge_index, edge_weight))

        # 按照边的顺序将节点特征向量重新排列
        edg_index_all = torch.cat((edge_index_pos, edge_index_neg), dim=1)
        Em, Ed = self.pro_data(x, edg_index_all)  # 筛选数据 获得源节点特征和目标节点

        # 将源节点特征和目标节点特征拼接起来
        x = torch.cat((Em, Ed),dim=1)
        x = self.fc(x)

        return x

    def pro_data(self, x, edg_index):
        m_index = edg_index[0]
        d_index = edg_index[1]
        Em = torch.index_select(x, 0, m_index)  # 沿着x1的第0维选出m_index
        Ed = torch.index_select(x, 0, d_index)
        return Em, Ed


# 模型构建
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCNII(in_channels=x.shape[1], out_channels=16).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=0.0005)

# 训练模型
model.train()
for epoch in range(1, 201):
    optimizer.zero_grad()
    out = model(x, edge_index, edge_index_pos=edge_index_pos, edge_index_neg=edge_index_neg_selected,
                param={'m_number': adj_matrix.shape[0], 'd_number': adj_matrix.shape[1]})

    # 使用train数据进行训练
    loss = F.binary_cross_entropy_with_logits(out[train_idx], y[train_idx].float())
    loss.backward()
    optimizer.step()
    loss = loss.item()
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')

# 模型验证
model.eval()
with torch.no_grad():
    # 获得所有数据
    out = model(x, edge_index, edge_index_pos=edge_index_pos, edge_index_neg=edge_index_neg_selected,
                param={'m_number':adj_matrix.shape[0],'d_number':adj_matrix.shape[1]})

    # 提取验证集数据
    out_pred = out[test_idx]
    y_pred = y[test_idx]

    # 计算AUC
    auc = roc_auc_score(y_pred, out_pred)
    print('AUC:', auc)

代码详解

以下是更详细的代码解释：

首先，导入需要用到的库：

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.utils import to_undirected
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score

然后，加载邻接矩阵数据并将其转换为张量形式：

adj_matrix = pd.read_csv('m_d.csv', header=None, index_col=None).values

num_nodes = adj_matrix.shape[0] + adj_matrix.shape[1]
edge_index_pos = np.column_stack(np.argwhere(adj_matrix != 0))
edge_index_pos[:, 1] += adj_matrix.shape[0]
edge_index_pos = torch.tensor(edge_index_pos, dtype=torch.long)

edge_index_neg = np.column_stack(np.argwhere(adj_matrix == 0))
edge_index_neg[:, 1] += adj_matrix.shape[0]
edge_index_neg = torch.tensor(edge_index_neg, dtype=torch.long)

这里的 adj_matrix 是一个邻接矩阵，表示节点之间的连接情况。edge_index_pos 表示有边相连的节点对的索引，edge_index_neg 则表示没有边相连的节点对的索引。

接下来，从正负样本中随机选取相同数量的边，形成平衡的训练集：

num_pos_edges_number = edge_index_pos.shape[1]
selected_neg_edge_indices = torch.randint(high=edge_index_neg.shape[1], size=(num_pos_edges_number,), dtype=torch.long)
edge_index_neg_selected = edge_index_neg[:, selected_neg_edge_indices]

edg_index_all = torch.cat((edge_index_pos, edge_index_neg_selected), dim=1)

这里使用 torch.randint 函数从负样本索引中随机抽取与正样本相同数量的索引，以达到正负样本平衡的目的。

然后，创建数据集，并将邻接矩阵转换为无向图（会将每条边复制成两条相反的边）：

x = torch.ones((num_nodes, 1))
y = torch.cat((torch.ones((edge_index_pos.shape[1], 1)),
               torch.zeros((edge_index_neg_selected.shape[1], 1))), dim=0)

edge_index = to_undirected(edge_index_pos)

x 是节点特征矩阵，这里没有节点特征，所以全部设为 1。y 表示分类标签，其中正样本标签为 1，负样本标签为 0。edge_index 是无向图上的边索引。

接着，将数据随机分为训练、验证和测试集：
这里使用 numpy 库的 random.shuffle 函数将数据索引随机打乱，并按照 8:2 的比例划分为训练集和测试集。

idx = np.arange(y.shape[0])
np.random.shuffle(idx)
train_idx = idx[:int(0.8 * len(idx))]
test_idx = idx[int(0.8 * len(idx)):]

然后，创建一个 GCN 模型：
GCNII 是一个继承自 torch.nn.Module 的类，它包含了两层 GCN 卷积 (self.conv1, self.conv2) 和一个全连接层 (self.fc)。在前向传播函数中，首先对节点特征进行两层 GCN 卷积操作，然后将边索引按照顺序重新排列，接着通过 pro_data 函数筛选出源节点和目标节点的特征向量，并将其拼接在一起，最后通过全连接层得到输出。

class GCNII(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(GCNII, self).__init__()
        self.conv1 = GCNConv(in_channels, 2 * out_channels, cached=True)
        self.conv2 = GCNConv(2 * out_channels, out_channels)
        self.fc = torch.nn.Linear(out_channels*2, 1)

    def forward(self, x, edge_index, edge_index_pos, edge_index_neg, edge_weight=None,param={}):
        """前向传播函数"""
        # 对节点特征进行两层 GCN 卷积操作
        x = F.relu(self.conv1(x, edge_index, edge_weight))
        x = F.relu(self.conv2(x, edge_index, edge_weight))

        # 筛选数据，获取源节点特征和目标节点特征
        edg_index_all = torch.cat((edge_index_pos, edge_index_neg), dim=1)
        Em, Ed = self.pro_data(x, edg_index_all)

        # 将源节点特征和目标节点特征拼接起来，并通过全连接层得到输出
        x = torch.cat((Em, Ed),dim=1)
        x = self.fc(x)

        return x

    def pro_data(self, x, edg_index):
        """筛选数据，获取源节点特征和目标节点特征"""
        m_index = edg_index[0]
        d_index = edg_index[1]
        Em = torch.index_select(x, 0, m_index)
        Ed = torch.index_select(x, 0, d_index)
        return Em, Ed

接下来，构建模型并开始训练：
这里首先判断是否有可用的 GPU，然后创建一个 GCNII 模型，并使用 Adam 优化器对模型参数进行训练。在训练过程中，计算模型输出和分类标签之间的二元交叉熵损失，并调用 backward 函数计算梯度并更新模型参数。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCNII(in_channels=x.shape[1], out_channels=16).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=0.0005)

model.train()
for epoch in range(1, 201):
    optimizer.zero_grad()

    # 计算输出，使用 train_idx 进行训练
    out = model(x, edge_index, edge_index_pos=edge_index_pos, edge_index_neg=edge_index_neg_selected,
                param={'m_number': adj_matrix.shape[0], 'd_number': adj_matrix.shape[1]})
    loss = F.binary_cross_entropy_with_logits(out[train_idx], y[train_idx].float())
    
    # 反向传播计算梯度并更新参数
    loss.backward()
    optimizer.step()
    loss = loss.item()
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')

最后，在测试集上评估模型性能：

model.eval()
with torch.no_grad():
    # 计算模型输出
    out = model(x, edge_index, edge_index_pos=edge_index_pos, edge_index_neg=edge_index_neg_selected,
                param={'m_number':adj_matrix.shape[0],'d_number':adj_matrix.shape[1]})

    # 提取验证集数据并计算 AUC
    out_pred = out[test_idx]
    y_pred = y[test_idx]
    auc = roc_auc_score(y_pred, out_pred)
    print('AUC:', auc)