如何在图神经网络（GNN）中实现多分类/多标签节点分类？

一、核心思路调整 #

你的场景属于多标签节点分类，和常规单标签多分类的核心区别是：每个节点可以同时属于多个类别（一种药物对应多种疾病），无需用独热编码的单标签逻辑处理，只需调整模型输出和损失函数即可适配现有代码。

二、具体实现方案 #

1. 模型输出层修改 #

基于你参考的GCN代码，只需把输出层维度从“单标签类别数”改成“疾病类别数n”，每个维度对应一种疾病的预测概率。比如原代码中最后一层GCNConv(hidden_channels, num_classes)，这里的num_classes直接设为疾病种类数n即可，不需要对标签做独热编码转换。

2. 损失函数选择：二元交叉熵（BCE） #

完全可以用BCE损失分别判断每个疾病是否匹配，不需要搭建n个独立模型。具体要点：

• 标签直接用形状为[药物数量, n]的0/1矩阵（1表示药物对应该疾病，0表示不对应），无需额外转换
• 优先使用torch.nn.BCEWithLogitsLoss（自带Sigmoid激活，数值稳定性更好），不要用多分类的CrossEntropyLoss

3. 关键代码示例 #

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

# 调整后的GCN模型
class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__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).relu()
        x = self.conv2(x, edge_index)
        return x

# 初始化模型（out_channels设为疾病类别数n）
model = GCN(in_channels=128, hidden_channels=64, out_channels=n)
criterion = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练循环示例
model.train()
for epoch in range(100):
    optimizer.zero_grad()
    out = model(x, edge_index)
    # 标签y是[num_nodes, n]的0/1矩阵
    loss = criterion(out, y.float())
    loss.backward()
    optimizer.step()
    print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')

# 推理阶段
model.eval()
with torch.no_grad():
    out = model(x, edge_index)
    # 用sigmoid转成概率，再设阈值（比如0.5）得到预测标签
    pred_probs = torch.sigmoid(out)
    pred_labels = (pred_probs > 0.5).int()

三、补充优化建议 #

如果疾病类别存在不平衡情况，可以给BCE损失添加类别权重（BCEWithLogitsLoss(pos_weight=weight_tensor)），提升少数类别的预测效果。

内容的提问来源于stack exchange，提问作者knhc12345