TrafficFormer模型在昇腾平台上迁移适配实践

flow blue

119人浏览 · 2026-07-06 11:14:17

flow blue · 2026-07-06 11:14:17 发布

作者：昇腾实战派
知识地图：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）；
核心作用：捕捉三大流量特征：
1. 局部特征：单 Burst 内部报文、字段的关联；
2. 全局时序特征：同一 Flow 内多个 Burst 的长距离依赖；
3. 流间隐式关联：不同流量片段的语义共性；
优化点：针对流量长序列特点优化注意力计算效率，保证模型推理速度。

4. 预训练双任务头（核心创新，自监督阶段）

TrafficFormer 设计两个互补的自监督预训练任务，联合训练提升表征能力，是模型核心亮点：

任务 1：掩码 Burst 建模（MBM，Masked Burst Modeling）

对标 NLP 领域的掩码语言建模（MLM）：

随机遮蔽输入序列中的部分 Burst Token；
模型根据上下文语义预测被遮蔽的原始 Burst 内容；
作用：学习流量局部字段语义、报文组合规则，理解基础流量语法。

任务 2：同源-同向-同流多分类（SODF）

针对网络流量流属性定制的细粒度分类任务，弥补传统同源任务的缺陷：

对切分后的 Burst 片段做重组、颠倒、跨流拼接，构造 5 类不同样本；
模型判别样本类别（判断是否同源、流向是否一致、是否属于同一条流）；
作用：强制模型学习流量流向、流归属、报文顺序等高层行为特征，解决传统模型仅依赖五元组判断同源的浅层问题。

预训练损失：MBM 损失 + SODF 分类损失联合加权优化。

版本信息

软件	版本
cann	8.3.RC1
hdk	25.3.rc1
python	3.11
torch	2.1.0

CANN 镜像获取

docker pull 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 命令检查驱动是否正常。

创建容器

docker run --privileged -it -u root --ipc=host --network=host \
    --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 \
    --name=traffic swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-openeuler24.03-py3.11 /bin/bash

环境安装

git clone https://github.com/IDP-code/TrafficFormer.git  
cd TrafficFormer  
git apply ../trafficformer.patch
pip install -r requirements.txt

trafficformer.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_class=argparse.ArgumentDefaultsHelpFormatter)
diff --git a/requirements.txt b/requirements.txt
index 427c028..e6c9747 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,18 +1,20 @@
-cxxfilt==0.3.0
-docutils==0.20.1
-flowcontainer==7.2
-jieba==0.42.1
-matplotlib==3.8.0
-numpy==2.1.2
-packaging==24.1
-pandas==2.2.3
-scapy==2.5.0
-scikit_learn==1.3.1
-sentencepiece==0.1.99
-setuptools==68.0.0
-six==1.16.0
-Sphinx==8.1.3
-sphinx_rtd_theme==3.0.1
-torch==2.0.1
-torchvision==0.15.1
-tqdm==4.65.0
+flowcontainer
+jieba
+matplotlib
+packaging
+pandas
+scapy
+scikit_learn
+sentencepiece
+six
+Sphinx
+sphinx_rtd_theme
+tqdm
+torch==2.1.0
+torchvision==0.16.0
+pyyaml
+decorator
+scipy
+attrs
+psutil
+numpy==1.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
torch==2.1.0
torchvision==0.16.0
pyyaml
decorator
scipy
attrs
psutil
numpy==1.26.4

配置启动脚本

针对 pre_train.sh 文件：

export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

dataset=/data/12_13_masked.pt
python3 pre-training/pretrain.py --dataset_path ${dataset} \
    --vocab_path models/encryptd_vocab.txt \
    --output_model_path model.bin \
    --world_size 8 \
    --gpu_ranks 0 1 2 3 4 5 6 7 \
    --master_ip tcp://localhost:12345 \
    --total_steps 90000 \
    --save_checkpoint_steps 10000 \
    --batch_size 32 \
    --embedding word_pos_seg \
    --encoder transformer \
    --mask fully_visible \
    --target mlm