ET-BERT实战：如何用你自己的PCAP数据训练一个加密流量分类模型？

ET-BERT实战：如何用你自己的PCAP数据训练一个加密流量分类模型？

当企业安全团队需要监控内部网络中的异常加密流量时，通用模型往往难以识别特定业务场景下的威胁。这时，使用私有PCAP数据训练定制化的ET-BERT模型就成为刚需。本文将手把手带您完成从原始流量到分类模型的完整Pipeline。

1. 环境准备与数据规范检查

在开始前，请确保您的开发环境满足以下基础要求：

• Python 3.8+ 与 PyTorch 1.10+
• 安装Wireshark的tshark工具（用于PCAP解析）
• 至少16GB内存（处理大型流量文件时需要）

数据格式验证是首要步骤。打开您的PCAP文件样本，检查是否包含：

tshark -r your_traffic.pcap -T fields -e tcp.payload | head -n 5

理想输出应显示十六进制负载数据。如果您的数据是pcapng格式，需要先转换：

from pcapng import FileScanner with open('input.pcapng', 'rb') as fp: scanner = FileScanner(fp) packets = list(scanner)

关键目录结构需要与ET-BERT兼容：

├── datasets │ └── your_dataset │ ├── raw_pcaps # 原始PCAP存放处 │ ├── splitcap # 自动生成的会话分割文件 │ └── processed # 最终训练数据

2. 深度解析数据预处理流程

2.1 会话分割与特征提取

修改data_process/main.py中的核心参数：

# 关键参数说明 pcap_path = "/path/to/your/pcaps" # 原始数据路径 dataset_save_path = "/output/path" # 处理结果路径 samples = 10000 # 每类样本数上限 dataset_level = "packet" # 可选"flow"或"packet"

执行会话分割时，特别注意：

当splitcap=True时，程序会执行耗时较长的流量重组操作。建议对大型数据集分批次处理。

典型问题排查表：

2.2 标签映射策略

对于自定义标签体系，需要修改open_dataset_deal.py中的类别映射逻辑。例如企业内部分类：

label_mapping = { "normal_web": 0, "vpn_zoom": 1, "malware_c2": 2, # 添加您的特有类别... }

3. 模型训练实战技巧

3.1 预训练模型适配

下载基础预训练模型后，需调整词汇表：

python vocab_process/main.py \ --corpus_path corpora/your_traffic.txt \ --vocab_path models/custom_vocab.txt \ --min_count 3 # 低频词过滤阈值

关键参数调整建议：

• seq_length: 根据平均包长设置（可用tshark -z io,phs统计）
• mask: 企业流量建议用fully_visible而非causal

3.2 微调阶段优化

执行微调时推荐使用的参数组合：

python fine-tuning/run_classifier.py \ --pretrained_model_path models/pre-trained_model.bin \ --vocab_path models/custom_vocab.txt \ --train_path datasets/your_dataset/train.tsv \ --batch_size 64 \ # 根据GPU显存调整 --learning_rate 1e-5 \ # 小数据集需更低学习率 --epochs_num 15 \ # 早停机制建议设为8-20 --focal_loss 1 # 处理类别不平衡

性能提升技巧：

• 在models/目录下保存多个checkpoint
• 使用--weight_decay 0.01防止过拟合
• 对少量数据启用--freeze_bert冻结底层参数

4. 生产环境部署方案

4.1 模型轻量化处理

使用ONNX转换提升推理速度：

torch.onnx.export( model, dummy_input, "model.onnx", opset_version=11, input_names=["input_ids", "attention_mask"], dynamic_axes={ "input_ids": {0: "batch"}, "attention_mask": {0: "batch"} } )

4.2 实时分类实现

构建Flask API服务：

@app.route('/classify', methods=['POST']) def classify(): pcap = request.files['pcap'] tmp_path = f"/tmp/{uuid.uuid4()}.pcap" pcap.save(tmp_path) # 执行预处理 preprocess_pipeline(tmp_path) # 加载模型预测 result = model.predict(preprocessed_data) return jsonify(result)

性能指标参考值（基于i7-11800H + RTX 3070）：

对于需要更高性能的场景，建议：

• 使用Triton Inference Server
• 实现批量预测（batch_size=32时效率提升3倍）
• 对TCP重组等耗时操作用C++重写