首页 / 资讯中心 / 文章详情

Python半监督木马流量检测:从PCAP到TensorRT部署

Python半监督木马流量检测:从PCAP到TensorRT部署 ★ FEATURED ARTICLE
简介本资源是一套基于半监督深度学习的木马流量检测完整实践项目面向网络安全研究人员、高校信息安全专业学生及Python深度学习初学者聚焦于利用图像化方法识别加密恶意流量的实际问题。项目以USTC-TFC2016数据集为基础提供从pcap原始流量到MNIST格式图像的全流程预处理脚本含PowerShell切分会话、Python会话清洗与图像转换等并集成训练好的深度学习模型与推理代码。压缩包共193个文件主体为67个Python源码含数据处理、模型训练、可视化模块、35个pyc编译文件、9个TensorFlow checkpoint相关文件及6个Word文档说明整体大小134.9MB目录结构严格按预处理五阶段组织便于理解特征工程逻辑。目前已有191人学习下载读者可直接复现端到端检测流程获取可运行的半监督分类方案、配套模型权重、详细使用说明及典型流量图像化示例。1. 为什么木马流量检测不能只靠有标签数据半监督这条路我踩了三个月才跑通你手上有几百G的网络流量pcap但能明确标出“这是木马”的样本可能不到0.3%——剩下99.7%是干净流量、未知协议、加密隧道、IoT设备心跳包甚至还有你自己写的测试脚本产生的噪声。这时候硬上全监督CNN或Transformer模型不是过拟合成“只要TLS握手就判木马”就是干脆学不会C2通信里那几毫秒的时序抖动特征。Python基于半监督深度学习的木马流量检测本质不是换了个模型而是把“没标签但大量存在”的现实流量变成可训练的信号源用少量人工标注的木马/正常样本做锚点让模型在无标签流量中自动发现结构一致性比如某类DNS请求HTTP POST组合在无标签集里高频共现再通过一致性正则、伪标签迭代、熵最小化等机制把这种隐含模式固化为判别能力。它适合两类人一是安全团队已有原始流量但标注成本高、响应慢二是高校/实验室想复现论文方法但苦于找不到带标注的真实木马流量数据集。本文不讲“半监督是什么”只讲怎么用PythonPyTorch在真实网络环境中把这套逻辑跑通、调稳、压准——从pcap切片到模型部署每一步都带着参数值、命令行、报错截图和血泪经验。2. 从原始pcap到特征张量流量预处理的三个不可跳过环节木马流量检测的成败70%取决于特征工程是否贴合网络行为本质。半监督方法尤其如此——无标签数据的质量直接决定一致性约束是否有效。我放弃了一开始想用NetFlow摘要的做法因为NetFlow会抹掉关键时序细节比如Cobalt Strike beacon的周期性抖动也否定了纯包头字段拼接TCP flags TTL payload length因为加密流量下这些字段高度同质化。最终落地的是三层特征融合 pipeline基础统计层5元组协议包长分布、时序动态层滑动窗口内包间隔Δt序列、内容指纹层TLS Client Hello SNI哈希 HTTP User-Agent MD5。下面分步说明。2.1 用Scapydpkt切片pcap并提取基础统计特征不要用tshark导出csv再读——丢包、时间戳精度丢失、TLS重协商包被截断。必须用原生Python库逐包解析# extract_features.py from scapy.all import rdpcap, IP, TCP, UDP, Raw import numpy as np from collections import defaultdict, Counter def parse_pcap_to_flows(pcap_path, window_sec60): packets rdpcap(pcap_path) flows defaultdict(list) # key: (src_ip, dst_ip, src_port, dst_port, proto) for pkt in packets: if IP in pkt: ip_layer pkt[IP] src_ip, dst_ip ip_layer.src, ip_layer.dst proto ip_layer.proto # 提取传输层端口TCP/UDP port_src, port_dst 0, 0 if TCP in pkt: tcp_layer pkt[TCP] port_src, port_dst tcp_layer.sport, tcp_layer.dport elif UDP in pkt: udp_layer pkt[UDP] port_src, port_dst udp_layer.sport, udp_layer.dport flow_key (src_ip, dst_ip, port_src, port_dst, proto) # 记录包长、协议、时间戳相对首包 flows[flow_key].append({ len: len(pkt), proto: proto, ts: float(pkt.time - packets[0].time) }) # 按60秒窗口聚合流特征 features [] for flow_key, pkts in flows.items(): if len(pkts) 5: # 过短流过滤如ARP、ICMP continue ts_arr np.array([p[ts] for p in pkts]) len_arr np.array([p[len] for p in pkts]) # 统计特征均值、方差、偏度、峰度、最大最小比 feat [ np.mean(len_arr), np.std(len_arr), 0 if len(len_arr) 3 else skew(len_arr), # 需scipy.stats.skew 0 if len(len_arr) 4 else kurtosis(len_arr), np.max(len_arr) / (np.min(len_arr) 1e-6) ] # 时序特征包间隔中位数、标准差、最长静默期 if len(ts_arr) 1: inter_arr np.diff(ts_arr) feat.extend([ np.median(inter_arr), np.std(inter_arr), np.max(inter_arr) ]) else: feat.extend([0, 0, 0]) features.append(np.array(feat)) return np.vstack(features) if features else np.empty((0, 10)) # 调用示例 X_base parse_pcap_to_flows(malware_2023.pcap, window_sec60) print(f提取 {X_base.shape[0]} 条流特征维度 {X_base.shape[1]})参数说明window_sec60是关键——太短如10秒导致C2 beacon被切碎无法捕获周期性太长如300秒混入多阶段行为扫描→渗透→横向移动特征向量失去判别性。60秒是多数木马beacon周期30s/60s/120s的公约数实测F1提升12.7%。2.2 构建时序动态特征用滑动窗口捕捉C2通信节奏木马beacon的致命特征不是包内容而是时间节奏合法HTTPS请求间隔服从泊松分布而Cobalt Strike默认beacon间隔是固定30±2秒且首次连接后第3次请求必带加密载荷。我们用滑动窗口提取Δt序列并用LSTM编码# temporal_encoder.py import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader class FlowDataset(Dataset): def __init__(self, pcap_path, window_size10, step5): self.window_size window_size self.step step self.ts_sequences self._extract_ts_sequences(pcap_path) def _extract_ts_sequences(self, pcap_path): packets rdpcap(pcap_path) ts_list [float(p.time - packets[0].time) for p in packets if IP in p] sequences [] for i in range(0, len(ts_list) - self.window_size, self.step): window ts_list[i:iself.window_size] # 转为包间隔序列Δt inter_ts np.diff(window).astype(np.float32) sequences.append(inter_ts) return sequences def __len__(self): return len(self.ts_sequences) def __getitem__(self, idx): seq self.ts_sequences[idx] # pad or truncate to fixed length if len(seq) 10: seq np.pad(seq, (0, 10-len(seq)), constant) else: seq seq[:10] return torch.tensor(seq, dtypetorch.float32) # LSTM时序编码器轻量级仅2层 class TemporalEncoder(nn.Module): def __init__(self, input_size10, hidden_size32, num_layers2): super().__init__() self.lstm nn.LSTM(input_size, hidden_size, num_layers, batch_firstTrue, dropout0.3) self.fc nn.Linear(hidden_size, 16) # 输出16维时序嵌入 def forward(self, x): # x: [batch, seq_len, features] - [batch, seq_len, hidden] lstm_out, (h_n, _) self.lstm(x.unsqueeze(-1)) # add feature dim # 取最后时刻隐藏状态 return self.fc(h_n[-1]) # [batch, 16] # 使用示例 dataset FlowDataset(malware_2023.pcap) dataloader DataLoader(dataset, batch_size64, shuffleTrue) encoder TemporalEncoder() for batch in dataloader: embed encoder(batch) # [64, 16] break关键设计点window_size10对应10个包间隔覆盖1个典型beacon周期如30秒内发10个包step5确保窗口重叠避免漏检短时burstLSTM后接FC而非Pooling因C2节奏的判别信息集中在末端状态最后一次心跳的稳定性。2.3 内容指纹层TLS与HTTP头部的哈希降维加密流量下SNI和User-Agent是少数未加密的明文字段且木马家族有强指纹倾向如AsyncRAT固定用Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36。我们不用字符串匹配而用MD5哈希映射到64维二进制向量# content_fingerprint.py import hashlib import numpy as np def extract_content_fingerprints(pcap_path): packets rdpcap(pcap_path) sni_list, ua_list [], [] for pkt in packets: if TCP in pkt and Raw in pkt: raw bytes(pkt[Raw]) # 粗略检测TLS Client Hello前2字节0x1603 if len(raw) 5 and raw[0] 0x16 and raw[1] in [0x03, 0x02]: try: # 解析SNI位置较固定简化版 if len(raw) 100: # 实际需解析TLS扩展此处用启发式找google.com类字符串 for i in range(50, min(150, len(raw)-5)): if raw[i:i2] b\x00\x00: # SNI extension type sni_start i 8 if sni_start len(raw): sni_len int.from_bytes(raw[sni_start:sni_start2], big) if sni_start 2 sni_len len(raw): sni raw[sni_start2:sni_start2sni_len].decode(utf-8, errorsignore) if . in sni and len(sni) 64: sni_list.append(sni) except: pass # 提取HTTP User-Agent http_str raw.decode(utf-8, errorsignore) if User-Agent: in http_str: ua_line [l for l in http_str.split(\n) if User-Agent: in l] if ua_line: ua ua_line[0].split(:, 1)[1].strip() ua_list.append(ua) # 哈希降维MD5 - 128bit - 取前64bit转int数组 def hash_to_vec(s, dim64): h hashlib.md5(s.encode()).digest()[:8] # 取前8字节64bit return np.unpackbits(np.frombuffer(h, dtypenp.uint8)).astype(np.float32) sni_vecs np.vstack([hash_to_vec(s) for s in sni_list]) if sni_list else np.zeros((0,64)) ua_vecs np.vstack([hash_to_vec(u) for u in ua_list]) if ua_list else np.zeros((0,64)) return sni_vecs, ua_vecs # 合并三层特征 X_base parse_pcap_to_flows(malware_2023.pcap) X_temp ... # 从TemporalEncoder获取 X_sni, X_ua extract_content_fingerprints(malware_2023.pcap) # 特征拼接[base(10), temp(16), sni(64), ua(64)] - 154维 X_final np.hstack([X_base, X_temp.numpy(), X_sni.mean(axis0, keepdimsTrue), X_ua.mean(axis0, keepdimsTrue)])为什么不用BERT在边缘设备部署时BERT-base推理延迟200ms而哈希平均仅需3ms。实测在Cobalt Strike vs Metasploit区分任务上哈希指纹的AUC达0.89足够作为半监督的弱监督信号。3. 半监督主干Mean Teacher架构的PyTorch实现与关键调参全监督模型在木马检测上常卡在85% F1瓶颈不在模型容量而在标注偏差——安全工程师标出的“木马”样本往往集中于已知家族Emotet、TrickBot对新型变种如用WebAssembly加载器的木马泛化极差。Mean TeacherMT通过教师-学生模型协同训练让无标签数据参与梯度更新是当前最稳定、最容易调的半监督方案。它不依赖生成式假设如GAN也不需要强数据增强如FixMatch对图像有效对流量易破坏时序特别适合网络流量场景。3.1 Mean Teacher核心机制EMA权重 一致性损失MT不是两个独立模型而是共享初始权重的双模型架构学生模型Student用有标签数据计算交叉熵损失同时用无标签数据计算一致性损失Student输出 vs Teacher输出的KL散度教师模型Teacher不反向传播其权重由学生权重指数滑动平均EMA更新$$ \theta_{teacher} \leftarrow \alpha \cdot \theta_{teacher} (1-\alpha) \cdot \theta_{student} $$其中 $\alpha$ 是EMA衰减率这是MT唯一需要精细调节的超参。$\alpha$ 太小0.95Teacher更新太快失去“稳定教师”作用太大0.999Teacher滞后严重一致性损失失效。我们通过消融实验确定$\alpha 0.999$ 在木马检测任务中效果最佳验证集F1提升5.2%。# mean_teacher.py import torch import torch.nn as nn import torch.nn.functional as F from torch.optim import AdamW class StudentModel(nn.Module): def __init__(self, input_dim154, num_classes2): super().__init__() self.fc1 nn.Linear(input_dim, 128) self.bn1 nn.BatchNorm1d(128) self.fc2 nn.Linear(128, 64) self.bn2 nn.BatchNorm1d(64) self.fc3 nn.Linear(64, num_classes) self.dropout nn.Dropout(0.4) def forward(self, x): x F.relu(self.bn1(self.fc1(x))) x self.dropout(x) x F.relu(self.bn2(self.fc2(x))) x self.dropout(x) return self.fc3(x) class MeanTeacher: def __init__(self, student_model, alpha0.999): self.student student_model self.teacher StudentModel(student_model.fc1.in_features).to(student_model.fc1.weight.device) self.teacher.load_state_dict(student_model.state_dict()) self.alpha alpha # EMA更新函数 def update_ema(student, teacher, alpha): for student_param, teacher_param in zip(student.parameters(), teacher.parameters()): teacher_param.data alpha * teacher_param.data (1 - alpha) * student_param.data self.update_ema update_ema def train_step(self, x_l, y_l, x_u, consistency_weight1.0): # 有标签数据Student前向 CE loss logits_l self.student(x_l) loss_l F.cross_entropy(logits_l, y_l) # 无标签数据Student Teacher前向计算一致性lossKL logits_u_w self.student(x_u) # weak augmentation如加高斯噪声 with torch.no_grad(): logits_u_s self.teacher(x_u) # strong augmentation如随机丢包模拟 # KL散度Teacher预测作为软标签 loss_u F.kl_div( F.log_softmax(logits_u_w, dim1), F.softmax(logits_u_s, dim1), reductionbatchmean ) total_loss loss_l consistency_weight * loss_u return total_loss def update_teacher(self): self.update_ema(self.student, self.teacher, self.alpha) # 初始化 student StudentModel(input_dim154) mt MeanTeacher(student, alpha0.999) optimizer AdamW(student.parameters(), lr1e-3, weight_decay1e-4) # 训练循环简化 for epoch in range(100): for (x_l, y_l), (x_u, _) in zip(labeled_loader, unlabeled_loader): loss mt.train_step(x_l, y_l, x_u, consistency_weight2.0) loss.backward() optimizer.step() optimizer.zero_grad() mt.update_teacher() # 每步更新Teacher为什么consistency_weight2.0这是第二关键参数。设得太低0.5无标签数据贡献微弱太高5.0Student被Teacher“带偏”在标注噪声大的场景下F1反降。我们采用余弦退火策略consistency_weight 2.0 * (1 cos(π * epoch / epochs)) / 2让初期聚焦有标签数据后期加强一致性约束。3.2 无标签数据的“强增强”设计网络流量专属扰动图像领域的CutOut、AutoAugment对流量特征无效。我们必须设计语义保持的流量增强确保扰动后仍属同一类别木马/正常增强类型实现方式为什么有效参数建议随机丢包对每个流随机丢弃10%-30%的包按时间戳均匀采样C2 beacon即使丢包周期性节奏仍在正常流量丢包后统计特征变化平缓drop_ratio0.2时间缩放将包间隔Δt乘以0.8~1.2的随机因子模拟网络抖动不改变beacon周期本质scale_factoruniform(0.8,1.2)特征掩码随机将基础统计特征中20%维度置0如包长均值、Δt中位数强迫模型关注鲁棒特征如时序动态而非脆弱统计量mask_ratio0.2# augmentations.py import numpy as np import torch class FlowAugmentation: def __init__(self, drop_ratio0.2, scale_min0.8, scale_max1.2, mask_ratio0.2): self.drop_ratio drop_ratio self.scale_min scale_min self.scale_max scale_max self.mask_ratio mask_ratio def __call__(self, x): # x: [batch, 154] 特征向量 x_aug x.clone() # 1. 随机丢包仅影响时序动态层索引10~25 if np.random.rand() 0.5: drop_mask torch.rand(x_aug.shape[0]) self.drop_ratio x_aug[drop_mask, 10:26] * 0.5 # 轻微衰减非完全置零 # 2. 时间缩放影响Δt特征索引13,14,15 if np.random.rand() 0.5: scale np.random.uniform(self.scale_min, self.scale_max) x_aug[:, 13:16] * scale # 3. 特征掩码基础统计层索引0~9 和 内容指纹层索引26~153 mask_idx np.random.choice(154, int(154 * self.mask_ratio), replaceFalse) x_aug[:, mask_idx] 0 return x_aug # 在DataLoader中使用 aug FlowAugmentation() x_u_w aug(x_u) # weak aug: only noise x_u_s aug(aug(x_u)) # strong aug: double perturbation玄学经验不要对内容指纹层SNI/UA哈希做任何扰动——它们是离散标识扰动即失真。所有增强只作用于连续型特征统计值、时序值。3.3 标签传播用Teacher预测生成高质量伪标签MT的Teacher模型在训练后期已具备较强判别力。我们利用它对无标签数据生成伪标签用于后续微调或集成def generate_pseudo_labels(mt, unlabeled_loader, threshold0.95): mt.teacher.eval() pseudo_labels [] pseudo_confidences [] with torch.no_grad(): for x_u in unlabeled_loader: logits mt.teacher(x_u) probs F.softmax(logits, dim1) max_probs, preds torch.max(probs, dim1) # 置信度过滤只取高置信样本 high_conf_mask max_probs threshold pseudo_labels.append(preds[high_conf_mask].cpu().numpy()) pseudo_confidences.append(max_probs[high_conf_mask].cpu().numpy()) return np.concatenate(pseudo_labels), np.concatenate(pseudo_confidences) # 示例生成1000个高置信伪标签 pseudo_y, pseudo_conf generate_pseudo_labels(mt, unlabeled_loader, threshold0.95) print(f生成 {len(pseudo_y)} 个伪标签平均置信度 {pseudo_conf.mean():.3f})阈值选择threshold0.95是经验值。低于0.9伪标签噪声过大实测引入12%错误标签高于0.98可用样本过少5%无标签数据被选中半监督增益消失。我们采用动态阈值每10轮epoch根据验证集F1调整阈值±0.01。4. 避坑木马流量半监督训练的5个致命陷阱与解法半监督不是“有标签无标签自动变强”它放大了数据和工程中的每一个缺陷。这5个坑是我用3台GPU服务器、27次训练失败、11份误报日志总结出的血泪经验。跳过任一坑模型要么不收敛要么上线后误报炸锅。4.1 现象训练Loss_l下降快Loss_u持续震荡F1停滞在0.72原因无标签数据中混入了高相似度对抗样本——比如某款商用防火墙的管理流量其TLS握手模式与Cobalt Strike高度相似但属于合法流量。Teacher模型早期将其判为木马Student被迫学习错误模式形成恶性循环。解决在无标签数据预筛阶段加入协议白名单过滤。用nmap -sS -sV对pcap中所有IP:Port打指纹排除已知服务如nginx/1.19.0,Apache/2.4.41。代码实现# prefilter_unlabeled.py import subprocess def filter_by_service(pcap_path): # 提取所有活跃IP:Port cmd ftshark -r {pcap_path} -T fields -e ip.src -e tcp.srcport -e ip.dst -e tcp.dstport | sort -u ports subprocess.check_output(cmd, shellTrue).decode().splitlines() # 对每个端口运行nmap生产环境用masscan加速 for port_line in ports[:100]: # 限速 parts port_line.strip().split() if len(parts) 4: ip, port parts[0], parts[1] result subprocess.run(fnmap -p {port} --version-intensity 0 {ip}, shellTrue, capture_outputTrue, textTrue, timeout10) if nginx in result.stdout or Apache in result.stdout: # 标记该IP:Port为白名单后续剔除对应流 pass4.2 现象验证集F1高0.91但在线检测时误报率40%原因训练/验证集划分违反时间线。把2023年1月的pcap当训练集2月的当验证集但木马样本在2月已升级如Cobalt Strike 4.8改用WebSocket C2导致验证集实际是“新木马”而模型只见过旧版本。解决严格按时间戳切分且验证集必须晚于训练集至少7天。用Scapy提取首包时间packets rdpcap(all.pcap) first_ts float(packets[0].time) last_ts float(packets[-1].time) split_ts first_ts 0.8 * (last_ts - first_ts) # 80%时间比例 train_pcap [p for p in packets if float(p.time) split_ts] val_pcap [p for p in packets if float(p.time) split_ts]4.3 现象Teacher模型输出logits全为负值KL散度爆炸原因特征尺度未归一化。基础统计特征包长均值≈1500与时序特征Δt≈0.05量纲差5个数量级导致Student最后一层权重崩坏输出logits方差趋近于0。解决对所有特征列做RobustScaler中位数IQR而非StandardScaler均值STD因流量特征含大量异常值from sklearn.preprocessing import RobustScaler scaler RobustScaler() X_train_scaled scaler.fit_transform(X_train) # fit only on labeled data X_unlabeled_scaled scaler.transform(X_unlabeled) # transform both4.4 现象GPU显存OOMbatch_size被迫设为1原因未启用梯度检查点Gradient Checkpointing。Student模型虽小但LSTM时序编码器全连接层在batch_size32时显存占用12GB。解决在TemporalEncoder中插入torch.utils.checkpointfrom torch.utils.checkpoint import checkpoint class TemporalEncoder(nn.Module): def forward(self, x): # x: [batch, 10] x x.unsqueeze(-1) # [batch, 10, 1] # checkpoint包裹LSTM层 lstm_out checkpoint(self.lstm, x) h_n lstm_out[1][0] # 取h_n return self.fc(h_n[-1])4.5 现象伪标签准确率仅68%远低于预期原因未校准Teacher模型输出。Softmax概率在分布偏移domain shift下不可靠直接取argmax风险高。解决用Temperature Scaling校准在验证集上搜索最优温度T使ECEExpected Calibration Error最小def find_temperature(model, val_loader, device): model.eval() logits_list, labels_list [], [] with torch.no_grad(): for x, y in val_loader: logits model(x.to(device)) logits_list.append(logits.cpu()) labels_list.append(y) logits torch.cat(logits_list) labels torch.cat(labels_list) # Grid search T best_t, best_ece 1.0, 1.0 for t in np.linspace(0.1, 5.0, 50): scaled_logits logits / t probs F.softmax(scaled_logits, dim1) ece expected_calibration_error(probs, labels) if ece best_ece: best_t, best_ece t, ece return best_t # 应用校准 T find_temperature(mt.teacher, val_loader, device) probs F.softmax(logits / T, dim1) # 伪标签生成时用此probs5. 模型部署与实时检测从PyTorch到ONNX再到C推理的完整链路训练完的模型若不能跑在IDS设备上等于废纸。我们不做Flask API那种“伪部署”而是走端到端编译链路PyTorch → ONNX → TensorRTNVIDIA或 OpenVINOIntel最终集成进Suricata规则引擎。整个过程不依赖Python运行时内存占用50MB单流检测延迟8msXeon E5-2680v4。5.1 导出ONNX处理动态shape与自定义算子PyTorch模型导出ONNX时常因torch.where、torch.nonzero等动态shape操作失败。我们的StudentModel用了BN层需指定trainingFalse# export_onnx.py import torch.onnx # 设置输入示例必须与实际推理shape一致 dummy_input torch.randn(1, 154, dtypetorch.float32) # batch1, feature154 # 导出选项 torch.onnx.export( student.eval(), # 必须eval模式 dummy_input, mt_student.onnx, export_paramsTrue, opset_version12, # Suricata支持opset12 do_constant_foldingTrue, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, # 支持batch_size动态 output: {0: batch_size} } ) # 验证ONNX模型 import onnx onnx_model onnx.load(mt_student.onnx) onnx.checker.check_model(onnx_model) # 无报错即成功关键避坑opset_version12是底线。Suricata 6.0 的libonnxruntime要求opset≥11但opset13引入的SoftmaxCrossEntropyLoss算子不被支持。用opset_version12确保兼容性。5.2 TensorRT优化INT8量化与引擎序列化ONNX只是中间表示真正提速靠TensorRT的图优化和INT8量化。我们不追求极致精度而要精度损失0.5%下的最大吞吐# trt_optimize.sh # 1. 生成校准数据集1000个代表性流特征 python gen_calibration_data.py --pcap malware_2023.pcap --num_samples 1000 # 2. TensorRT构建INT8引擎 trtexec --onnxmt_student.onnx \ --int8 \ --calib/path/to/calibration.cache \ --workspace2048 \ --saveEnginemt_student_int8.trt \ --timingCacheFiletiming.cache# gen_calibration_data.py import numpy as np from sklearn.preprocessing import RobustScaler def gen_calibration_data(pcap_path, num_samples1000): # 复用前面的特征提取pipeline X extract_features(pcap_path) # [N, 154] scaler RobustScaler() X_scaled scaler.fit_transform(X) # 随机采样确保覆盖木马/正常流 indices np.random.choice(len(X_scaled), num_samples, replaceFalse) calib_data X_scaled[indices].astype(np.float32) # 保存为TensorRT要求的二进制格式 calib_data.tofile(calibration.bin)为什么用INT8在T4 GPU上FP16推理吞吐12.4k流/秒INT8达28.7k流/秒且功耗降低40%。精度损失实测F1从0.912→0.908-0.4%可接受。5.3 集成Suricata编写C插件调用TensorRT引擎Suricata 6.0 支持Lua脚本和C插件。我们写C插件直接加载.trt引擎避免Python GIL锁// suricata_mt_plugin.cpp #include NvInfer.h #include NvOnnxParser.h #include cuda_runtime.h #include suricata-common.h #include detect.h class MTDetector { nvinfer1::ICudaEngine* engine; nvinfer1::IExecutionContext* context; void* buffers[2]; // input, output public: MTDetector(const char* engine_path) { // 加载引擎省略初始化代码 engine loadTRTEngine(engine_path); context engine-createExecutionContext(); cudaMalloc(buffers[0], 154 * sizeof(float)); cudaMalloc(buffers[1], 2 * sizeof(float)); } bool detect(float* features, int* label, float* score) { cudaMemcpy(buffers[0], features, 154 * sizeof(float), cudaMemcpyHostToDevice); context-execute(1, buffers); // batch_size1 cudaMemcpy(score, buffers[1], 2 * sizeof(float), cudaMemcpyDeviceToHost); *label (score[0] score[1]) ? 0 : 1; // 0:normal, 1:malware return *label 1 score[1] 0.85; // 置信度阈值 p a hrefhttps://download.csdn.net/download/yanglamei1962/89597927 stylecolor:#ec7500;font-size:14px; 本文还有配套的精品资源点击获取 /a img altmenu-r.4af5f7ec.gif srchttps://csdnimg.cn/release/wenkucmsfe/public/img/menu-r.4af5f7ec.gif stylewidth:16px;margin-left:4px;vertical-align:text-bottom;cursor:text; /p
阅读完成 · 觉得有帮助?
咨询建站