简介：这份资源是2023年华中科技大学计算机学院毕业设计项目，主题为基于溯源图的APT攻击检测方法优化，面向网络安全方向的学生与研究人员，适合具备一定机器学习与图神经网络基础、希望深入理解高级持续性威胁检测的读者。项目围绕溯源图构建与检测模型优化展开，涉及数据预处理、特征提取、模型训练与评估等环节，并包含RGAT、GRU等模型实现及StreamSpot、DARPA CADETS等数据集上的实验代码。压缩包共26个文件，以Python源码为主，辅以XML配置、Markdown说明文档及项目配置文件，整体约48KB，结构紧凑，便于快速定位核心算法与实验脚本。目前已有265人学习下载，读者可从中获取完整的毕业设计实现思路、模型代码、实验数据组织方式与结果分析框架，适合作为网络安全检测方向的学习参考或二次开发基础。
1. 溯源图 + APT 检测：这份 HUST 毕设源码到底能跑出什么
APT 攻击检测这件事，最让人头疼的不是没有告警，而是告警太多、线索太碎。一个高级持续性威胁从入口到横向移动再到数据外带，可能横跨几十台主机、几百个进程、上千次文件读写，传统基于单点特征的 IDS 根本串不起来这条链。溯源图（Provenance Graph）的思路就是把系统调用、进程创建、文件访问这些底层审计事件，按因果关系连成一张有向图，攻击者的行为在这张图上会呈现出一条异常路径。这份 2023 年华中科技大学计算机学院的毕业设计，核心就是在这张图上做文章——用 RGAT（Relational Graph Attention Network）和 GRU 组合，对 StreamSpot 和 DARPA TC CADETS 两个数据集做 APT 检测，并且给出了可运行的 PyTorch 实现。
我拿到这个包的第一反应是：文件结构比想象中干净。main.py是入口，streamspot_RGAT.py、darpa_cadets_RGAT_GRU.py、darpa_cadets_RGAT.py分别对应不同数据集和模型组合，model_RGAT.py、model_RGAT_GRU.py、model.py是模型定义，data/下放数据集，doc/里有0416.md和ss.md两份说明。适合谁？如果你正在做网络安全方向的毕设、课程设计，或者想找一个能跑通的图神经网络 + APT 检测的 baseline，这份代码的参考价值很高。但如果你指望解压就能出论文级结果，那得先过数据预处理和依赖版本这两关。
2. 溯源图建模与 RGAT 选型：为什么不是 GCN 或 GAT
2.1 溯源图的图结构特征与 RGAT 的适配逻辑
溯源图不是普通的同构图。节点类型包括进程、文件、socket、管道，边类型包括 fork、exec、read、write、connect，每种边携带的语义完全不同。普通 GCN 把所有边当同一种关系处理，聚合时会把 fork 和 read 混在一起，信息损失很大。GAT 加了注意力机制，能区分邻居的重要性，但仍然是单关系假设。RGAT 的做法是给每种边类型分配独立的注意力头参数，在聚合时按关系类型分组计算注意力权重，这样进程创建和文件读取对目标节点的影响就能被分别建模。
这份代码里model_RGAT.py的实现思路是：先对节点特征做线性变换，然后对每种边类型分别计算注意力系数，最后把多头的输出拼接或平均。节点特征用的是节点类型 embedding + 度数统计，边特征在消息传递时作为额外的注意力偏置。这个设计在 StreamSpot 数据集上比较合适，因为 StreamSpot 的图规模小、边类型少，RGAT 的参数量不会爆炸。
2.2 从原始审计日志到图结构：数据预处理的三个关键步骤
代码包里data/目录下的数据不是直接可用的图结构，需要先做转换。常见做法是：从审计日志中提取 (src, dst, edge_type, timestamp) 四元组，按时间窗口切分，然后构建节点特征矩阵和边索引矩阵。具体步骤：
import torch import numpy as np from collections import defaultdict def build_provenance_graph(events, node_type_map, edge_type_map): """ events: list of (src_id, dst_id, edge_type, timestamp) node_type_map: dict, node_id -> node_type edge_type_map: dict, edge_type_str -> edge_type_id """ node_set = set() edge_list = [] edge_type_list = [] for src, dst, etype, ts in events: node_set.add(src) node_set.add(dst) edge_list.append([src, dst]) edge_type_list.append(edge_type_map[etype]) # 构建节点特征：类型 one-hot + 度数 node_ids = sorted(node_set) id_to_idx = {nid: i for i, nid in enumerate(node_ids)} num_nodes = len(node_ids) num_types = len(set(node_type_map.values())) node_features = np.zeros((num_nodes, num_types + 1)) for nid in node_ids: idx = id_to_idx[nid] ntype = node_type_map.get(nid, 0) node_features[idx, ntype] = 1.0 node_features[idx, -1] = 1.0 # 度数占位，后续填充 # 边索引重映射 edge_index = torch.tensor( [[id_to_idx[s], id_to_idx[d]] for s, d in edge_list], dtype=torch.long ).t().contiguous() edge_type = torch.tensor(edge_type_list, dtype=torch.long) x = torch.tensor(node_features, dtype=torch.float) return x, edge_index, edge_type这段代码的逻辑是先把事件流里的节点去重、编号，然后为每个节点构造类型 one-hot 向量，边索引转成 PyTorch 的edge_index格式（2×E 的矩阵），边类型单独存一个向量。参数说明：node_type_map需要你根据数据集文档手动定义，StreamSpot 里通常是 0=进程、1=文件、2=socket；edge_type_map同理，fork=0、exec=1、read=2、write=3。注意edge_index的转置操作不能漏，否则 RGAT 的propagate会报维度错误。
2.3 RGAT 层的核心参数与调参边界
model_RGAT.py里 RGAT 层的关键参数有三个：in_channels、out_channels、num_relations。in_channels要和节点特征维度对齐，out_channels一般设 64 或 128，num_relations必须等于边类型总数。注意力头数heads建议从 4 开始试，StreamSpot 上 4 头就够了，DARPA CADETS 因为图更大可以上 8 头。dropout 设 0.3 到 0.5 之间，太高会欠拟合，太低在 CADETS 上容易过拟合。学习率用 1e-3 配 Adam，如果 loss 震荡就降到 5e-4。这些参数在main.py里都有对应变量，改的时候注意同步改argparse的默认值。
3. 跑通 StreamSpot 与 DARPA CADETS：从环境配置到训练脚本
3.1 环境依赖与版本对齐
这份代码是 2023 年的，PyTorch 版本建议用 1.13 或 2.0，torch-geometric 用 2.3.x。如果你用最新的 torch-geometric 2.5+，RGATConv的导入路径可能变了，需要改成from torch_geometric.nn import RGATConv并检查num_relations参数名是否一致。常见做法是建一个 conda 环境：
conda create -n hust_apt python=3.9 conda activate hust_apt pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html pip install torch-geometric==2.3.1 pip install numpy pandas scikit-learn tqdm matplotlib注意 torch 和 torch-geometric 的版本要匹配，torch-geometric 2.3.1 对应 torch 1.13 和 2.0 都行，但 2.4+ 要求 torch 2.1 以上。如果你没有 GPU，把+cu117去掉装 CPU 版，但 DARPA CADETS 的图比较大，CPU 训练会慢到让你怀疑人生。
3.2 StreamSpot 训练流程与关键代码
StreamSpot 数据集相对小，适合先跑通验证流程。streamspot_RGAT.py是主训练脚本，核心逻辑是加载图、划分训练测试集、定义 RGAT 模型、训练循环。我一般会先跑 5 个 epoch 看 loss 是否下降：
from model_RGAT import RGAT from torch_geometric.loader import DataLoader import torch.nn.functional as F import torch def train_streamspot(graph_data, num_relations, num_classes, epochs=50): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = RGAT( in_channels=graph_data.x.shape[1], hidden_channels=64, out_channels=num_classes, num_relations=num_relations, heads=4 ).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5) model.train() for epoch in range(epochs): optimizer.zero_grad() out = model(graph_data.x, graph_data.edge_index, graph_data.edge_type) loss = F.nll_loss(out[graph_data.train_mask], graph_data.y[graph_data.train_mask]) loss.backward() optimizer.step() scheduler.step() if epoch % 10 == 0: model.eval() with torch.no_grad(): pred = model(graph_data.x, graph_data.edge_index, graph_data.edge_type).argmax(dim=1) acc = (pred[graph_data.test_mask] == graph_data.y[graph_data.test_mask]).float().mean() print(f'Epoch {epoch}, Loss: {loss.item():.4f}, Test Acc: {acc:.4f}') model.train() return model逻辑说明：train_mask和test_mask是布尔索引，需要在数据加载阶段根据节点划分生成。nll_loss要求模型输出 log_softmax，RGAT 的forward里已经做了。参数说明：hidden_channels设 64 是保守选择，如果你显存够可以上 128；weight_decay用 5e-4 是为了防止过拟合，StreamSpot 上如果 acc 上不去可以降到 1e-4。注意scheduler.step()的位置，放在optimizer.step()之后是对的，但如果你用 PyTorch 1.13 之前的版本，需要改成scheduler.step(epoch)。
3.3 DARPA CADETS 的 GRU 融合与序列建模
DARPA CADETS 的数据是时序的，单纯用 RGAT 做节点分类会丢失时间信息。darpa_cadets_RGAT_GRU.py的做法是把图按时间窗口切成多个快照，每个快照过 RGAT 得到节点嵌入，然后把嵌入序列喂给 GRU 做时序建模，最后用最后一个时间步的输出做分类。这个设计在 APT 检测里比较合理，因为 APT 攻击本身就是多阶段、跨时间的。
关键参数：时间窗口大小window_size建议设 10 到 20 个事件，太小了图结构不完整，太大了 GRU 序列太长梯度容易消失。GRU 的hidden_size和 RGAT 的hidden_channels保持一致，一般是 64 或 128。num_layers设 1 或 2，CADETS 上 2 层 GRU 效果通常比 1 层好，但训练时间翻倍。代码里model_RGAT_GRU.py的forward接收的是(x, edge_index, edge_type, batch)，其中batch是时间步索引，用来把节点嵌入按时间步分组。
4. 避坑与排查：这份毕设代码里最容易翻车的五个地方
4.1 现象：运行main.py报ModuleNotFoundError: No module named 'torch_geometric.nn.conv.rgat_conv'
原因：torch-geometric 版本不对。2.3 之前RGATConv在torch_geometric.nn.conv下，2.3 之后移到了torch_geometric.nn.conv.rgat_conv，但导入方式变了。解决：先pip show torch-geometric看版本，如果是 2.3.x，改成from torch_geometric.nn import RGATConv；如果是 2.4+，需要from torch_geometric.nn.conv.rgat_conv import RGATConv。最稳妥的办法是锁版本到 2.3.1。
4.2 现象：StreamSpot 训练时 loss 一直是nan
原因：节点特征里有全零行，RGAT 的注意力计算对全零输入会出nan。StreamSpot 里有些节点只有出边没有入边，度数特征没填上。解决：在build_provenance_graph里把度数占位那列填上实际度数，或者加一个1e-8的平滑项。另外检查edge_index里有没有自环，有的话在 RGAT 层加add_self_loops=False。
4.3 现象：DARPA CADETS 跑几个 epoch 后显存爆了
原因：CADETS 的图快照数量多，每个快照都保留计算图，显存累积。解决：在训练循环里每个 batch 结束后del out, loss并torch.cuda.empty_cache()，或者把window_size从 20 降到 10。如果还不行，把hidden_channels从 128 降到 64，heads从 8 降到 4。
4.4 现象：测试集准确率很高但实际检测不出 APT
原因：数据划分有问题。StreamSpot 和 CADETS 都是类别不平衡的，正常行为占 90% 以上，如果随机划分测试集，模型只要全预测正常就能拿高 acc。解决：用分层采样StratifiedKFold，或者按时间划分——用前 70% 时间的数据训练，后 30% 测试。代码里main.py的split_data函数默认是随机划分，建议改成时间划分。
4.5 现象：GRU 融合后效果反而比纯 RGAT 差
原因：GRU 的输入序列里混入了 padding 的零向量，GRU 把零向量也当有效输入处理了。解决：在model_RGAT_GRU.py的forward里加pack_padded_sequence，或者手动把 padding 位置的输出 mask 掉。另外检查window_size是否和实际序列长度匹配，如果最后一个窗口不足window_size，要么补齐要么丢弃。
5. 进阶技巧：用注意力权重做攻击路径可视化
跑通训练只是第一步，这份代码真正有价值的地方在于 RGAT 的注意力权重可以拿来做攻击路径的可视化。APT 检测不能只给一个二分类结果，安全分析师需要知道“为什么这个节点被判定为异常”。RGAT 的每一层都会输出注意力系数，把多层注意力按边聚合，就能得到每条边的重要性分数。分数高的边连起来，就是模型认为的攻击路径。
具体做法：在model_RGAT.py的forward里把每层的alpha存下来，训练完后取测试集中预测为异常的节点，回溯它的入边注意力。我一般会写一个后处理脚本：
def extract_attack_path(model, x, edge_index, edge_type, target_node, top_k=10): model.eval() with torch.no_grad(): out, attentions = model(x, edge_index, edge_type, return_attention=True) # 取最后一层注意力，按目标节点的入边排序 last_alpha = attentions[-1] # shape: [num_edges, heads] edge_scores = last_alpha.mean(dim=1) # 多头平均 # 找目标节点的入边 mask = edge_index[1] == target_node in_edges = edge_index[:, mask] in_scores = edge_scores[mask] # 取 top-k topk_idx = torch.topk(in_scores, min(top_k, len(in_scores))).indices path_edges = in_edges[:, topk_idx] return path_edges.cpu().numpy(), in_scores[topk_idx].cpu().numpy()逻辑说明：return_attention=True需要你在模型里加一个分支，把每层的alpha收集到一个列表里返回。edge_scores是多头注意力的平均值，也可以用最大值。mask筛选出指向目标节点的边，topk取分数最高的几条。参数说明：top_k设 10 到 20 比较合适，太少了路径不完整，太多了噪声大。注意edge_index的维度是 2×E，edge_index[1]是目标节点索引。
这个可视化结果可以直接画成子图，用 networkx 或 pyvis 都行。我习惯用 networkx 导出 DOT 格式，然后用 Graphviz 渲染。节点颜色按类型区分，边粗细按注意力分数映射。这样一张图放到论文里，比单纯报一个准确率有说服力得多。
从那以后我每次跑图神经网络做安全检测，都会强制走一遍注意力回溯——模型说异常，我得看到它到底在看哪条边。希望帮到你。
本文还有配套的精品资源，点击获取
