​作者​昇腾实战派​知识地图​https://blog.csdn.net/Lumos_Lovegood/article/details/161601003背景概述网络流量分析是网络安全与运维的关键技术但传统方法依赖大量人工标注数据泛化能力差。自监督预训练模型的出现为解决标注数据稀缺问题提供了新思路。本文基于昇腾 AI 平台介绍 TrafficFormer 的模型架构、环境部署与预训练实践。模型介绍TrafficFormer 是面向网络流量Network Traffic的高效 Transformer 预训练模型由清华大学等团队提出核心用于解决网络流量标注数据稀缺、流量语义挖掘困难的行业痛点主打自监督预训练 下游任务微调范式可适配流量分类、恶意流量识别、网站指纹识别、协议交互理解等网络安全与流量分析任务。网络流量数据具备时序性、结构化、字段密集的特点传统模型依赖人工标注泛化能力差。TrafficFormer 借鉴 NLP 预训练思路针对流量独有特征定制双自监督预训练任务搭配专属流量数据增强策略大幅提升无标注数据利用率与下游任务精度对比同类流量预训练模型性能更优F1 分数最高提升 10%。模型架构如下图所示1. 数据预处理模块输入层这是流量模型的专属前置模块区别于通用 NLP 预处理流量单元切分将完整网络流Flow划分为多个连续报文簇Burst再将单个 Burst 切分为两段基础单元Token 化把流量五元组、协议字段、报文长度、时间间隔等结构化字段转为模型可识别的 Token序列规整截断/补全序列长度统一输入维度添加分隔符[SEP]、特殊标识位区分 Burst 分段。2. 嵌入层Embedding Layer融合三类嵌入适配流量时序与结构特征字段嵌入对网络协议、端口、报文内容等离散字段做向量映射位置嵌入建模报文、Burst 的时序先后关系捕捉流量时间依赖分段嵌入区分同一个 Burst 切分后的两段单元辅助流向、顺序特征学习。三类嵌入向量逐元素相加作为 Transformer 编码器的初始输入。3. 主干网络Transformer 编码器基础配置12 层标准 Transformer Encoder搭载多头自注意力机制与层归一化Layer Norm核心作用捕捉三大流量特征局部特征单 Burst 内部报文、字段的关联全局时序特征同一 Flow 内多个 Burst 的长距离依赖流间隐式关联不同流量片段的语义共性优化点针对流量长序列特点优化注意力计算效率保证模型推理速度。4. 预训练双任务头核心创新自监督阶段TrafficFormer 设计两个互补的自监督预训练任务联合训练提升表征能力是模型核心亮点任务 1掩码 Burst 建模MBMMasked Burst Modeling对标 NLP 领域的掩码语言建模MLM随机遮蔽输入序列中的部分 Burst Token模型根据上下文语义预测被遮蔽的原始 Burst 内容作用学习流量局部字段语义、报文组合规则理解基础流量语法。任务 2同源-同向-同流多分类SODF针对网络流量流属性定制的细粒度分类任务弥补传统同源任务的缺陷对切分后的 Burst 片段做重组、颠倒、跨流拼接构造 5 类不同样本模型判别样本类别判断是否同源、流向是否一致、是否属于同一条流作用强制模型学习流量流向、流归属、报文顺序等高层行为特征解决传统模型仅依赖五元组判断同源的浅层问题。预训练损失MBM 损失 SODF 分类损失 联合加权优化。版本信息软件版本cann8.3.RC1hdk25.3.rc1python3.11torch2.1.0CANN 镜像获取dockerpull swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-openeuler24.03-py3.11更多版本下载请参考 CANN 镜像仓库注意Atlas 800I A3 机器请下载带 a3 关键字的版本Atlas 800I A2 机器请下载 910b 关键字版本。可使用npu-smi info命令检查驱动是否正常。创建容器dockerrun--privileged-it-uroot--ipchost--networkhost\--device/dev/davinci_manager\--device/dev/devmm_svm\--device/dev/hisi_hdc\-v/etc/localtime:/etc/localtime\-v/usr/local/Ascend/driver:/usr/local/Ascend/driver\-v/usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi\-v/home:/home\-v/data:/data\--nametraffic swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-openeuler24.03-py3.11 /bin/bash环境安装gitclone https://github.com/IDP-code/TrafficFormer.gitcdTrafficFormergitapply../trafficformer.patch pipinstall-rrequirements.txttrafficformer.patch内容如下diff --git a/pre-training/pretrain.py b/pre-training/pretrain.py index 0fa0cac..a35ce18 100644 --- a/pre-training/pretrain.py b/pre-training/pretrain.py -6,7 6,15 import torch import uer.trainer as trainer from uer.utils.config import load_hyperparam from uer.opts import * - import torch_npu from torch_npu.contrib import transfer_to_npu import numpy as np import random seed 7 torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) def main(): parser argparse.ArgumentParser(formatter_classargparse.ArgumentDefaultsHelpFormatter) diff --git a/requirements.txt b/requirements.txt index 427c028..e6c9747 100644 --- a/requirements.txt b/requirements.txt -1,18 1,20 -cxxfilt0.3.0 -docutils0.20.1 -flowcontainer7.2 -jieba0.42.1 -matplotlib3.8.0 -numpy2.1.2 -packaging24.1 -pandas2.2.3 -scapy2.5.0 -scikit_learn1.3.1 -sentencepiece0.1.99 -setuptools68.0.0 -six1.16.0 -Sphinx8.1.3 -sphinx_rtd_theme3.0.1 -torch2.0.1 -torchvision0.15.1 -tqdm4.65.0 flowcontainer jieba matplotlib packaging pandas scapy scikit_learn sentencepiece six Sphinx sphinx_rtd_theme tqdm torch2.1.0 torchvision0.16.0 pyyaml decorator scipy attrs psutil numpy1.26.4 diff --git a/uer/trainer.py b/uer/trainer.py index 502c822..c1a70b8 100644 --- a/uer/trainer.py b/uer/trainer.py -154,13 154,16 class MlmTrainer(Trainer): done_tokens self.batch_size * self.seq_length * self.report_steps if self.dist_train: done_tokens * self.world_size cost_time time.time() - self.start_time print(| {:8d}/{:8d} steps | {:3.3f} s | {:8.2f} tokens/s | loss {:7.2f} | acc: {:3.3f}.format( self.current_step, self.total_steps, - done_tokens / (time.time() - self.start_time), cost_time, done_tokens / cost_time, self.total_loss / self.report_steps, self.total_correct / self.total_denominator))requirements.txt内容如下flowcontainer jieba matplotlib packaging pandas scapy scikit_learn sentencepiece six Sphinx sphinx_rtd_theme tqdm torch2.1.0 torchvision0.16.0 pyyaml decorator scipy attrs psutil numpy1.26.4配置启动脚本针对pre_train.sh文件exportASCEND_RT_VISIBLE_DEVICES0,1,2,3,4,5,6,7dataset/data/12_13_masked.pt python3 pre-training/pretrain.py--dataset_path${dataset}\--vocab_pathmodels/encryptd_vocab.txt\--output_model_pathmodel.bin\--world_size8\--gpu_ranks01234567\--master_iptcp://localhost:12345\--total_steps90000\--save_checkpoint_steps10000\--batch_size32\--embeddingword_pos_seg\--encodertransformer\--maskfully_visible\--targetmlm注意数据集需要提供。启动预训练bashpre_train.sh测试结果预训练过程中观测loss无精度问题。训练9W步loss趋势稳定下降并且性能达到预期。