# TrafficModel

**Repository Path**: wangpengyang123/traffic-model

## Basic Information

- **Project Name**: TrafficModel
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Apache-2.0
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 13
- **Created**: 2026-06-04
- **Last Updated**: 2026-06-04

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 基于多模态深度学习的零日恶意流量分类实现

## 项目简介

本项目基于Reformer架构实现了一个用于网络流量分析的深度学习模型。Reformer是一种高效的Transformer变体，使用局部敏感哈希(LSH)注意力机制，能够处理长序列数据的同时保持较低的计算复杂度。

该模型主要用于网络流量的多目标分类和深度度量学习，适用于网络安全领域的流量分析任务。

## 目录结构

```
├── reformer_pytorch/   # Reformer模型核心实现
│   ├── __init__.py     # 模块导出
│   ├── attention.py    # 注意力机制
│   ├── autopadder.py   # 自动填充工具
│   ├── DML.py          # 深度度量学习模型
│   ├── generative_tools.py  # 生成工具
│   ├── overall.py      # 整体模型架构
│   ├── recorder.py     # 记录器
│   ├── reformer_enc_dec.py  # 编码器-解码器结构
│   ├── reformer_pytorch.py  # 主Reformer实现
│   └── reversible.py   # 可逆层实现
├── train_n_ford/       # 训练脚本
│   ├── task_zero_day.py    # 零日检测任务
│   ├── train_five_cls.py   # 五类分类任务
│   └── trian_five_dml.py   # 五类DML训练
├── results/            # 训练结果
├── config.py           # 配置文件
└── utils.py            # 工具函数
```

## 核心功能

### 1. Reformer模型

- 基于局部敏感哈希(LSH)的注意力机制
- 可逆层设计，减少内存占用
- 支持长序列处理
- 多种位置编码选项（轴向、绝对、固定）
- 支持局部注意力和全局注意力混合使用

### 2. 流量分析模型

- **HAE (Header Autoencoder)**: 用于处理网络数据包头部特征
- **PAE (Payload Autoencoder)**: 用于处理网络数据包负载特征
- 多目标分类器：支持同时进行多个分类任务
- 深度度量学习：用于学习流量特征的相似度表示

### 3. 训练框架

- 支持五折交叉验证
- 多种训练任务：零日检测、五类分类、深度度量学习
- 完整的结果记录和分析

## 模型架构

### 整体架构

```
输入 (流量数据)
├── 头部特征 ──> HAE (Header Autoencoder) ──┐
│                                           │
└── 负载特征 ──> PAE (Payload Autoencoder) ──┼──> 特征融合 ──> Reformer注意力层 ──> 分类器/DML头
                                            │
                                            └──> 深度度量学习 ──> 特征映射
```

### HAE (Header Autoencoder)

- 卷积编码器 + 残差特征提取 + 解码器
- 用于提取数据包头部的高级特征
- 支持可配置的深度和特征维度

### PAE (Payload Autoencoder)

- 膨胀卷积编码器 + 多层残差特征提取
- 用于提取数据包负载的结构化特征
- 支持可配置的深度和特征维度

### Reformer注意力层

- 基于LSH的自注意力机制
- 支持多头注意力
- 可配置的桶大小和哈希轮数
- 支持局部注意力和全局注意力混合

## 安装依赖

```bash
pip install torch torchvision einops local_attention axial_positional_embedding product_key_memory
```

## 使用方法

### 1. 数据准备

确保输入数据格式符合模型要求：
- 头部特征：形状为 [batch_size, seq_len, h_infeature]
- 负载特征：形状为 [batch_size, seq_len, p_infeature]

### 2. 模型初始化

```python
from reformer_pytorch import Model

# 初始化模型
target_name = [['class1', 'class2'], ['class3', 'class4']]  # 多目标分类
h_infeature = 128  # 头部输入特征维度
h_depth = 4        # HAE深度
p_infeature = 1    # 负载输入特征维度
p_depth = 4        # PAE深度
out_feature = 256  # 输出特征维度
num_head = 8       # 注意力头数
drop = 0.3         # dropout率
att_mask = None    # 注意力掩码
h_att_mask = None  # 头部注意力掩码
p_att_mask = None  # 负载注意力掩码

model = Model(target_name, h_infeature, h_depth, p_infeature, p_depth, out_feature, num_head, drop, att_mask, h_att_mask, p_att_mask)
```

### 3. 模型训练

使用train_n_ford目录下的训练脚本进行训练：

```bash
python train_n_ford/train_five_cls.py  # 五类分类任务
python train_n_ford/trian_five_dml.py  # 五类DML训练
python train_n_ford/task_zero_day.py   # 零日检测任务
```

### 4. 模型推理

```python
# 输入数据
hdr = torch.randn(batch_size, seq_len, h_infeature)
pay = torch.randn(batch_size, seq_len, p_infeature)

# 模型推理
targetout, atten = model(hdr, pay)
```

## 配置说明

配置文件 `config.py` 包含模型和训练的主要参数：

- `BATCH_SIZE`: 批次大小
- `LEARNING_RATE`: 学习率
- `EPOCHS`: 训练轮数
- `NUM_HEADS`: 注意力头数
- `DROP_RATE`: Dropout率
- 其他模型特定参数


## 数据集

本模型在以下公开数据集上进行训练和验证：

### ISCX-VPN-2016 Dataset

- **发布机构**: 加拿大新不伦瑞克大学网络安全研究所 (Canadian Institute for Cybersecurity, UNB)
- **描述**: 用于区分VPN流量和非VPN流量的数据集，包含多种VPN协议的流量样本
- **下载链接**: https://www.unb.ca/cic/datasets/index.html
- **应用**: VPN流量识别、加密流量分类

### ISCX-Tor-2016 Dataset

- **发布机构**: 加拿大新不伦瑞克大学网络安全研究所 (Canadian Institute for Cybersecurity, UNB)
- **描述**: 用于区分Tor流量和非Tor流量的数据集，包含匿名网络流量样本
- **下载链接**: https://www.unb.ca/cic/datasets/index.html
- **应用**: Tor流量识别、匿名网络检测

### USTC-TFC-2016 Dataset

- **发布机构**: 中国科学技术大学 (University of Science and Technology of China)
- **描述**: 涵盖恶意和良性应用流量的数据集，包含多种应用类型的网络流量
- **下载链接**: http://www.nlpir.org/wordpress/USTC-TFC2016%E8%AF%AD%E6%96%87%E5%88%86%E7%B1%BB%E6%8C%91%E6%88%98%E8%B5%9B/
- **应用**: 恶意流量检测、应用分类、零日攻击检测

## 作者

ZhangZheng

## 许可证

MIT License