熬了两天终于把这个基于图神经网络的社交网络好友推荐搞出来了中间踩的坑是真的多记录一下省的以后又掉进去。完整源码链接https://pan.quark.cn/s/1e54aa2ae950先说数据生成这步吧。我一开始想自己手撸一个社交网络生成器写了个双重for循环500个节点跑了几分钟还没出来我直接裂开。后来发现networkx里面有现成的stochastic_block_model用社区内连接概率和社区间连接概率来控制图的稠密程度舒服多了。import numpy as np import networkx as nx import torch def generate_social_network(n_users500, n_communities5, seed42): np.random.seed(seed) torch.manual_seed(seed) sizes [n_users // n_communities] * n_communities for i in range(n_users % n_communities): sizes[i] 1 p_intra 0.15 p_inter 0.003 probs np.full((n_communities, n_communities), p_inter) np.fill_diagonal(probs, p_intra) G nx.stochastic_block_model(sizes, probs, seedseed) G nx.Graph(G) n G.number_of_nodes() communities np.zeros(n, dtypeint) start 0 for c, sz in enumerate(sizes): communities[start:startsz] c start sz idx np.random.permutation(n) communities communities[idx] mapping {old: new for new, old in enumerate(idx)} G nx.relabel_nodes(G, mapping) rng np.random.RandomState(seed) features np.zeros((n, 8)) features[:, 0] rng.uniform(0.2, 0.9, n) features[:, 1:6] rng.rand(n, 5) features[:, 6] rng.uniform(0.1, 1.0, n) features[:, 7] rng.uniform(0.05, 1.0, n) mean features.mean(axis0, keepdimsTrue) std features.std(axis0, keepdimsTrue) 1e-8 features (features - mean) / std adj nx.to_numpy_array(G, dtypenp.float32) return adj, features, communities, G这里有个巨坑——stochastic_block_model是networkx 2.7才加的我一开始用的老版本一直报错说没有这个函数查了半天才发现版本不对。然后就是社区标签的shuffle问题如果不打乱的话节点是按社区顺序排列的画出来的图就五坨各自挤在一起不好看。我琢磨了半天用relabel_nodes才搞定。节点特征搞了个8维的向量年龄、兴趣偏好5维、活跃度、账号年龄标准化一下。我觉得这样有点像真实社交平台的用户画像了。接下来是模型的实现。一开始想用PyTorch Geometric结果装了半天他*的CUDA版本对不上折腾了一个小时决定算了手写GCN。其实就是两个图卷积层每层就是A_hat X W加上ReLU激活和dropout。import torch import torch.nn as nn import torch.nn.functional as F class GCNLayer(nn.Module): def __init__(self, in_dim, out_dim, dropout0.2): super().__init__() self.weight nn.Parameter(torch.randn(in_dim, out_dim) * 0.1) self.bias nn.Parameter(torch.zeros(out_dim)) self.dropout nn.Dropout(dropout) def forward(self, x, adj_norm): x self.dropout(x) x x self.weight self.bias x adj_norm x return x class GCNEncoder(nn.Module): def __init__(self, in_dim8, hidden_dim64, out_dim32, dropout0.2): super().__init__() self.layer1 GCNLayer(in_dim, hidden_dim, dropout) self.layer2 GCNLayer(hidden_dim, out_dim, dropout) def forward(self, x, adj_norm): x self.layer1(x, adj_norm) x F.relu(x) x self.layer2(x, adj_norm) return x链接预测用的是内积解码器就是两个节点的embedding做点积然后sigmoid简单粗暴。训练的时候同时采样正边已有的好友关系和负边不存在的边二分类交叉熵损失。说到损失函数又踩坑了。一开始我直接用的BCEWithLogitsLoss训练到一半loss变成nan了。排查了半天发现是学习率设成0.1太大了梯度炸了。降到0.01就好了。还有一次AUC死活就是0.5左右跟抛硬币一样我刚开始以为自己模型写错了debug了一整个下午结果发现是负采样的时候把正边也采进去了相当于让模型区分两个一样的东西能好才怪。def split_edges(adj, val_ratio0.1, test_ratio0.1, seed42): np.random.seed(seed) n adj.shape[0] upper np.triu_indices(n, k1) all_pairs np.column_stack([upper[0], upper[1]]) is_edge adj[upper] 0.5 pos_edges all_pairs[is_edge] neg_edges all_pairs[~is_edge] n_pos len(pos_edges) n_val max(1, int(n_pos * val_ratio)) n_test max(1, int(n_pos * test_ratio)) perm np.random.permutation(n_pos) val_pos pos_edges[perm[:n_val]] test_pos pos_edges[perm[n_val:n_valn_test]] train_pos pos_edges[perm[n_valn_test:]] adj_train adj.copy().astype(np.float64) for u, v in val_pos: adj_train[u, v] adj_train[v, u] 0.0 for u, v in test_pos: adj_train[u, v] adj_train[v, u] 0.0 n_neg_needed len(train_pos) neg_avail len(neg_edges) if neg_avail n_neg_needed: repeat n_neg_needed // neg_avail 1 neg_edges np.tile(neg_edges, (repeat, 1))[:n_neg_needed] neg_avail n_neg_needed perm_neg np.random.permutation(neg_avail) neg_val neg_edges[perm_neg[:n_val]] neg_test neg_edges[perm_neg[n_val:n_valn_test]] neg_train neg_edges[perm_neg[n_valn_test:n_valn_testlen(train_pos)]] return adj_train, train_pos, val_pos, test_pos, neg_train, neg_val, neg_test边界问题也折腾了一下。如果图比较稀疏的话可能负样本不够用我加了个tile兜底。验证集和测试集的正边要从邻接矩阵里删掉不然就信息泄露了模型啥都不学直接猜邻居就能答对。训练写了个循环每次把正负边的embedding算出来算loss反向传播。def train_model(model, x, adj_norm, pos_edges, neg_edges, lr0.01, epochs200, devicecpu): model model.to(device) x torch.FloatTensor(x).to(device) adj_norm adj_norm.to(device) pos torch.LongTensor(pos_edges).to(device) neg torch.LongTensor(neg_edges).to(device) optimizer torch.optim.Adam(model.parameters(), lrlr) losses, aucs [], [] for epoch in range(epochs): model.train() optimizer.zero_grad() z model(x, adj_norm) pos_scores (z[pos[:, 0]] * z[pos[:, 1]]).sum(dim1) neg_scores (z[neg[:, 0]] * z[neg[:, 1]]).sum(dim1) pos_loss -F.logsigmoid(pos_scores).mean() neg_loss -F.logsigmoid(-neg_scores).mean() loss pos_loss neg_loss loss.backward() optimizer.step() with torch.no_grad(): all_scores torch.cat([pos_scores, neg_scores]).cpu().numpy() all_labels np.concatenate([np.ones(len(pos)), np.zeros(len(neg))]) auc_val roc_auc_score(all_labels, all_scores) losses.append(loss.item()) aucs.append(auc_val) if (epoch 1) % 20 0: print(fEpoch {epoch1:3d}/{epochs} | Loss: {loss.item():.4f} | AUC: {auc_val:.4f}) return losses, aucs200个epoch跑下来训练AUC最高到了0.85左右验证集0.80测试集0.82。说实话这个结果还行毕竟数据是纯合成的特征也就8维能有这个区分度说明GCN确实学到社区结构了。来看一下训练曲线图。左图是loss右图是AUC。可以看到loss下降很快前20个epoch就降到1.2左右了后面波动挺大的说明模型在震荡可能加个学习率衰减或者换个优化器会更好。AUC曲线倒是相对平稳一直在0.75到0.85之间晃荡。ROC曲线更加直观AUC 0.804说明模型在正负样本分类上比随机猜测好了一大截。PR曲线的AUC是0.76考虑到正负样本不平衡正边只有4000多条非边有12万条这个值不算差了。可视化这部分也踩了屎坑。matplotlib画网络图的时候nx.draw调用了某个老API跟新版本的matplotlib冲突了报_AxesStack object is not callable。我他*的查了半天最后改成用draw_networkx_nodes和draw_networkx_edges分开画显式传ax参数才解决。import matplotlib.pyplot as plt import networkx as nx import numpy as np from sklearn.metrics import roc_curve, precision_recall_curve plt.rcParams[font.sans-serif] [SimHei] plt.rcParams[axes.unicode_minus] False def plot_network(G, communities, save_pathfigures/network.png): fig, ax plt.subplots(figsize(14, 10)) pos nx.spring_layout(G, k0.3, iterations50, seed42) colors [plt.cm.tab10(c % 10) for c in communities] nx.draw_networkx_edges(G, pos, axax, edge_colorgray, alpha0.4, width0.3) nx.draw_networkx_nodes(G, pos, axax, node_colorcolors, node_size30) sm plt.cm.ScalarMappable(cmapplt.cm.tab10) sm.set_array([]) cbar fig.colorbar(sm, axax, ticksrange(5), shrink0.6) cbar.set_label(社区编号) ax.set_title(社交网络结构图节点颜色表示不同社区, fontsize14, pad15) ax.axis(off) plt.tight_layout() plt.savefig(save_path, dpi150, bbox_inchestight) plt.close()中文显示也是个坑。plt.rcParams[font.sans-serif] [SimHei]这行加上了但是如果系统没装SimHei字体就白搭会回退成方框。还好Windows一般都有这个字体如果你的系统没有可以去下载一个SimHei.ttf放到matplotlib的字体目录里删掉缓存就行了。画出来的网络结构图还是挺好看的五个社区分别用不同颜色标出来社区内部的连接密密麻麻社区之间只有零星的几条线跟我们的概率设置社区内15%社区间0.3%是一致的。推荐结果那幅图展示了给几个随机用户的Top10好友推荐。可以看到推荐的用户跟目标用户的embedding相似度都在0.4到0.8之间说明模型确实把结构相似的用户embedding拉近了。def get_recommendations(model, x, adj, top_k10, devicecpu): model.eval() with torch.no_grad(): x_t torch.FloatTensor(x).to(device) adj_norm normalize_adj(torch.FloatTensor(adj)).to(device) z model(x_t, adj_norm).cpu().numpy() n adj.shape[0] z_norm z / (np.linalg.norm(z, axis1, keepdimsTrue) 1e-10) sim z_norm z_norm.T results {} for uid in range(n): mask adj[uid] 0.5 mask[uid] False candidates np.where(mask)[0] scores sim[uid, candidates] top_idx candidates[np.argsort(-scores)[:top_k]] results[uid] (top_idx, sim[uid, top_idx]) return results这里要注意embedding要做L2归一化不然不同节点embedding的模长不一样点积结果会被模长主导而不是角度相似度。加了个1e-10防止除零有些节点的embedding可能刚好是零向量虽然理论上不太会发生。评估代码用了sklearn的roc_auc_score和precision_recall_curve标准的链路预测评估流程。def evaluate(model, x, adj_norm, pos_edges, neg_edges, devicecpu): model.eval() with torch.no_grad(): x torch.FloatTensor(x).to(device) adj_norm adj_norm.to(device) z model(x, adj_norm) pos torch.LongTensor(pos_edges).to(device) neg torch.LongTensor(neg_edges).to(device) pos_scores (z[pos[:, 0]] * z[pos[:, 1]]).sum(dim1).cpu().numpy() neg_scores (z[neg[:, 0]] * z[neg[:, 1]]).sum(dim1).cpu().numpy() y_true np.concatenate([np.ones(len(pos_scores)), np.zeros(len(neg_scores))]) y_score np.concatenate([pos_scores, neg_scores]) auc_val roc_auc_score(y_true, y_score) prec, rec, _ precision_recall_curve(y_true, y_score) pr_auc auc(rec, prec) return auc_val, pr_auc, y_true, y_score最后主程序把这些串起来跑一下出结果。import os import torch import numpy as np from data_gen import generate_social_network from model import GCNEncoder from train_eval import normalize_adj, split_edges, train_model, evaluate, get_recommendations from visualize import plot_network, plot_training, plot_roc_curve, plot_pr_curve, plot_recommendations def main(): print( * 60) print( 基于图神经网络的社交网络好友推荐算法) print( * 60) print(\n[1/5] 生成社交网络数据...) adj, features, communities, G generate_social_network(n_users500, n_communities5, seed42) n_edges G.number_of_edges() print(f 用户数: {G.number_of_nodes()}, 好友关系数: {n_edges}) print(f 特征维度: {features.shape[1]}) print(\n[2/5] 划分训练/验证/测试边...) adj_train, train_pos, val_pos, test_pos, neg_train, neg_val, neg_test \ split_edges(adj, val_ratio0.1, test_ratio0.1) print(f 训练正边: {len(train_pos)}, 验证正边: {len(val_pos)}, 测试正边: {len(test_pos)}) print(\n[3/5] 初始化GCN模型...) device cuda if torch.cuda.is_available() else cpu print(f 使用设备: {device}) model GCNEncoder(in_dim8, hidden_dim64, out_dim32, dropout0.2) total_params sum(p.numel() for p in model.parameters()) print(f 模型参数量: {total_params}) adj_norm normalize_adj(torch.FloatTensor(adj_train)) print(\n[4/5] 开始训练...) losses, aucs train_model(model, features, adj_norm, train_pos, neg_train, lr0.01, epochs200, devicedevice) print(f 训练完成! 最高AUC: {max(aucs):.4f}) print(\n[5/5] 模型评估 可视化...) os.makedirs(figures, exist_okTrue) val_auc, val_pr_auc, y_true_val, y_score_val evaluate( model, features, adj_norm, val_pos, neg_val, device) test_auc, test_pr_auc, y_true_test, y_score_test evaluate( model, features, adj_norm, test_pos, neg_test, device) print(f 验证集 AUC: {val_auc:.4f}, PR-AUC: {val_pr_auc:.4f}) print(f 测试集 AUC: {test_auc:.4f}, PR-AUC: {test_pr_auc:.4f}) print(\n 生成可视化图表...) plot_network(G, communities, figures/network.png) plot_training(losses, aucs, figures/training.png) plot_roc_curve(y_true_val, y_score_val, val_auc, figures/roc.png) plot_pr_curve(y_true_val, y_score_val, val_pr_auc, figures/pr_curve.png) recs get_recommendations(model, features, adj_train, top_k10, devicedevice) plot_recommendations(recs, communities, adj_train, n_sample5, top_k10, save_pathfigures/recommend.png) print(\n * 60) print( 全部完成! 图表已保存至 figures/ 目录) print( * 60) if __name__ __main__: main()说一下最终结果吧。500个用户5个社区生成了4059条好友关系。模型参数量才2656个可以说是非常轻量了。训练200个epoch验证集AUC 0.804测试集AUC 0.822PR-AUC 0.773。考虑到整个项目没用任何预训练模型也没用复杂的注意力机制纯手工GCN能到这个水平我觉得已经可以了。后续还可以改进的地方加GAT注意力机制、用更复杂的node特征比如加入文本embedding、在更大的图上测试、或者加入时间动态推荐。不过这两个晚上的成果作为一个毕设demo已经够交差了。附上项目目录结构方便参考4-基于图神经网络的社交网络好友推荐算法/ ├── data_gen.py # 社交网络数据生成 ├── model.py # GCN模型定义 ├── train_eval.py # 训练与评估 ├── visualize.py # 可视化图表 ├── main.py # 主程序入口 ├── requirements.txt # 依赖列表 └── figures/ # 生成的图表 ├── network.png # 社交网络结构图 ├── training.png # 训练损失和AUC曲线 ├── roc.png # ROC曲线 ├── pr_curve.png # PR曲线 └── recommend.png # 推荐结果展示行了写到这里了代码和文章都搞完了希望后来的人少踩点坑。