Sd_Met: Stable Diffusion Enhanced Mamba Architecture

Sd_Met是一个将NetMamba的Mamba架构与Stable Diffusion扩散模型相结合的深度学习框架，专门用于网络流量分类任务。

架构特点

• Mamba编码器: 基于Mamba架构的高效流量编码器
• Stable Diffusion: 集成扩散模型的预训练能力
• LoRA微调: 参数高效的微调机制
• 多层次特征融合: 结合编码器和扩散模型特征
• NetMamba兼容: 完全兼容NetMamba数据集格式

目录结构

Sd_Met/
├── src/
│   ├── models/                 # 核心模型架构
│   │   ├── sd_mamba.py         # Sd_Met主模型
│   │   ├── mamba_components.py # Mamba组件实现
│   │   └── net_mamba_components.py # NetMamba兼容组件
│   ├── engine/                 # 训练引擎
│   │   ├── pretrain_engine.py  # 预训练引擎
│   │   └── finetune_engine.py  # 微调引擎
│   ├── datasets/               # 数据集处理
│   │   └── sd_mamba_dataset.py # 数据集类
│   └── utils/                  # 工具函数
├── configs/
│   └── model_config/
│       └── sd_mamba_config.py  # 模型配置
├── scripts/
│   ├── pretrain/
│   │   └── pretrain_sd_mamba.py
│   ├── finetune/
│   │   └── finetune_sd_mamba.py
│   └── eval/
│       └── eval_sd_mamba.py
├── pretrained/                 # 预训练模型存储
└── examples/                   # 使用示例

安装

环境要求

• Python 3.8+
• PyTorch 1.12+
• CUDA 11.0+

安装步骤

1. 克隆项目:

git clone <repository-url>
cd Sd_Met

2. 创建虚拟环境:

conda create -n sd_mamba python=3.10
conda activate sd_mamba

3. 安装依赖:

pip install -r requirements.txt

4. 安装Mamba:

pip install mamba-ssm==1.1.1

快速开始

预训练

python scripts/pretrain/pretrain_sd_mamba.py \
    --data_path /path/to/dataset \
    --output_dir ./output/pretrain \
    --batch_size 128 \
    --steps 150000 \
    --mask_ratio 0.9

微调

python scripts/finetune/finetune_sd_mamba.py \
    --data_path /path/to/dataset \
    --pretrained ./output/pretrain/best_model.pth \
    --output_dir ./output/finetune \
    --epochs 120 \
    --nb_classes 2

评估

python scripts/eval/eval_sd_mamba.py \
    --data_path /path/to/dataset \
    --checkpoint ./output/finetune/best_model.pth \
    --output_dir ./output/eval

数据准备

数据集格式

Sd_Met兼容NetMamba数据集格式:

dataset/
├── train/
│   ├── class1/
│   │   ├── sample1.png
│   │   └── sample2.png
│   └── class2/
│       ├── sample3.png
│       └── sample4.png
├── val/
│   ├── class1/
│   └── class2/
└── test/
        ├── class1/
        └── class2/

自定义数据集

可以使用numpy数组创建数据集:

from src.datasets.sd_mamba_dataset import create_dataloader_from_arrays

# images: [N, H, W] numpy数组
# labels: [N] numpy数组
train_loader = create_dataloader_from_arrays(
    train_images, train_labels, 
    batch_size=32, shuffle=True, img_size=40
)

配置说明

模型配置

在 configs/model_config/sd_mamba_config.py 中可以配置:

• 数据配置: img_size, stride_size, in_chans
• Mamba配置: embed_dim, depth, drop_rate
• 扩散配置: num_timesteps, beta_min, beta_max
• LoRA配置: use_lora, lora_rank, lora_alpha

训练配置

预训练配置:

config = SdMambaPretrainConfig(
    mask_ratio=0.9,
    steps=150000,
    batch_size=128
)

微调配置:

config = SdMambaFinetuneConfig(
    num_classes=2,
    epochs=120,
    freeze_encoder=False
)

使用示例

Python API使用

import torch
from src.models.sd_mamba import sd_mamba_pretrain, sd_mamba_classifier
from configs.model_config.sd_mamba_config import SdMambaPretrainConfig, SdMambaFinetuneConfig

# 预训练
pretrain_config = SdMambaPretrainConfig()
model = sd_mamba_pretrain(**pretrain_config.__dict__)

# 微调
finetune_config = SdMambaFinetuneConfig(num_classes=2)
model = sd_mamba_classifier(**finetune_config.__dict__)

# 推理
model.eval()
with torch.no_grad():
    outputs = model(images)
    logits = outputs['logits']
    predictions = torch.argmax(logits, dim=-1)

LoRA微调

from src.models.sd_mamba import sd_mamba_classifier

# 创建模型
model = sd_mamba_classifier(num_classes=2, use_lora=True, lora_rank=16)

# 只训练LoRA参数
for name, param in model.named_parameters():
    if 'lora' in name:
        param.requires_grad = True
    else:
        param.requires_grad = False

性能比较

方法	准确率	参数量	训练时间
NetMamba	95.2%	2.5M	2h
Sd_Met	96.8%	3.1M	3h
Sd_Met + LoRA	96.5%	0.5M*	1h

*LoRA微调时实际训练参数量

可视化工具

Sd_Met提供了可视化工具来分析模型性能:

• 混淆矩阵
• ROC曲线（二分类）
• 训练曲线
• 特征可视化

故障排除

常见问题

1. CUDA内存不足: 减小batch_size或使用梯度累积
2. 训练速度慢: 使用pin_memory和num_workers优化数据加载
3. 模型收敛差: 调整学习率和warmup步骤

调试建议

1. 使用较小的数据集测试流程
2. 检查数据预处理是否正确
3. 监控训练损失和验证准确率

许可证

本项目采用MIT许可证，详见LICENSE文件。

引用

如果您使用Sd_Met，请引用：

@misc{sd_mamba_2024,
  title={Sd_Met: Stable Diffusion Enhanced Mamba Architecture for Network Traffic Classification},
  author={Sd_Met Team},
  year={2024},
  url={https://github.com/your-repo/sd_met}
}