# Devign-A3M

**Repository Path**: rzzrz/Devign-A3M

## Basic Information

- **Project Name**: Devign-A3M
- **Description**: No description available
- **Primary Language**: Python
- **License**: MulanPSL-2.0
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-11-22
- **Last Updated**: 2025-11-22

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Devign-A3M: Attention-Augmented Attention-Map Aware Devign Model

## 项目简介

Devign-A3M 是对原始 Devign 模型的改进，通过以下创新提升漏洞检测性能：

1. **GNN生成的注意力图**：使用图神经网络（GNN）生成节点注意力分数，并将其转换为注意力图用于增强CNN层
2. **双路自注意力**：在GNN和CNN两处都引入自注意力机制（GAT和SE-Block/Non-local），提高模型的表示能力
3. **注意力图融合**：将GNN生成的注意力图与CNN输入特征进行融合（乘性/加性/拼接）

## 项目结构

```
devign-a3m/
├── data/                          # 数据目录
│   ├── raw/                      # 原始数据
│   ├── processed/                # 预处理后的图数据
│   │   ├── train/               # 训练集
│   │   ├── val/                 # 验证集
│   │   └── test/                # 测试集
│   └── scripts/                 # 数据预处理脚本
├── joern_scripts/               # Joern图提取脚本
├── src/                         # 源代码
│   ├── datasets/
│   │   └── vul_dataset.py      # 数据集类
│   ├── models/
│   │   ├── gnn_layers.py       # GNN层（带自注意力）
│   │   ├── attention_gen.py    # 注意力图生成
│   │   ├── cnn_backbone.py     # CNN骨干网络
│   │   └── devign_a3m.py       # 完整模型
│   ├── utils/
│   │   ├── metrics.py          # 评估指标
│   │   ├── logger.py           # 日志工具
│   │   └── seed.py             # 随机种子设置
│   ├── configs/
│   │   └── devign_a3m.yaml     # 配置文件
│   ├── train.py                # 训练脚本
│   └── eval.py                 # 评估脚本
├── experiments/                 # 实验输出
│   ├── logs/                   # 日志文件
│   └── checkpoints/            # 模型检查点
├── notebooks/
│   └── visualize_attention.ipynb  # 注意力可视化（待实现）
├── requirements.txt            # Python依赖
└── README.md                  # 本文件
```

## 环境配置

### 1. Python环境

```bash
# 创建conda环境
conda create -n devign-a3m python=3.10 -y
conda activate devign-a3m

# 安装PyTorch（根据你的CUDA版本调整）
# CUDA 11.8:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

# 安装PyTorch Geometric
pip install torch-geometric

# 安装其他依赖
pip install -r requirements.txt
```

### 2. Joern安装（用于图提取）

参考 [Joern官方文档](https://github.com/joernio/joern) 安装Joern。

## 数据准备

### 1. 数据下载

推荐使用以下数据集：
- **DiverseVul**: 包含Devign样本的大型漏洞数据集
- **Big-Vul**: 微软的大规模漏洞数据集
- **MSR dataset**: 微软研究数据集

### 2. 数据预处理

使用Joern提取代码属性图（Code Property Graph）：

```bash
# 运行Joern提取脚本（需要根据实际情况编写）
cd joern_scripts
# 执行提取脚本
```

预处理后的数据应保存为`.pt`格式，每个文件包含：
- `node_features`: [N, D] 节点特征
- `edge_index`: [2, E] 边索引
- `label`: int 标签（0或1）
- `node_order`: [N] 节点顺序（可选）
- `meta`: dict 元数据

### 3. 创建测试数据（用于快速验证）

```python
from src.datasets.vul_dataset import create_dummy_sample
import os

# 创建一些dummy样本用于测试
data_dir = "data/processed/train"
os.makedirs(data_dir, exist_ok=True)

for i in range(10):
    label = i % 2  # 交替标签
    sample = create_dummy_sample(
        node_dim=100,
        num_nodes=50 + i * 5,
        num_edges=100 + i * 10,
        label=label,
        save_path=f"{data_dir}/sample_{i:04d}.pt"
    )
```

## 训练

### 1. 配置模型

编辑 `src/configs/devign_a3m.yaml` 调整模型和训练参数。

### 2. 开始训练

```bash
cd src
python train.py \
    --config configs/devign_a3m.yaml \
    --data_root ../data/processed \
    --exp_name devign_a3m_run1 \
    --seed 42
```

### 3. 恢复训练

```bash
python train.py \
    --config configs/devign_a3m.yaml \
    --data_root ../data/processed \
    --exp_name devign_a3m_run1 \
    --resume ../experiments/devign_a3m_run1/checkpoints/best_model.pth \
    --seed 42
```

## 评估

```bash
cd src
python eval.py \
    --ckpt ../experiments/devign_a3m_run1/checkpoints/best_model.pth \
    --data_root ../data/processed \
    --split test \
    --output_dir ../experiments/evaluation
```

## 模型架构

### 1. GNN模块
- 使用GAT（Graph Attention Network）进行图表示学习
- 输出节点embedding和节点注意力分数

### 2. 注意力图生成
- 将节点注意力分数转换为2D注意力图
- 支持多种归一化方法：softmax、minmax、z-score

### 3. CNN模块
- 接收图映射后的网格特征和注意力图
- 使用SE-Block进行通道注意力
- 使用Non-local Block进行空间自注意力

### 4. 分类器
- 融合GNN的图级特征和CNN的全局特征
- 输出二分类结果

## 实验结果

### 指标
- Accuracy（准确率）
- Precision（精确率）
- Recall（召回率）
- F1-Score
- ROC AUC
- PR AUC

### 可视化
- ROC曲线
- PR曲线
- 注意力热图（overlay到源代码）

## 消融实验建议

1. **Baseline (Devign原始)**：复现原始Devign作为基线
2. **Devign + GNN注意力图**：仅使用GNN生成的注意力图
3. **Devign + GNN自注意力**：仅在GNN中使用自注意力
4. **Devign + CNN自注意力**：仅在CNN中使用自注意力
5. **Devign-A3M (完整)**：使用所有改进

## 常见问题

### 1. Joern解析失败
- 确保Joern版本与源码兼容
- 先用小样本测试

### 2. 模型损失不降
- 检查学习率设置
- 先用小模型和小数据集测试
- 检查数据格式是否正确

### 3. 显存不足
- 减小batch_size
- 减小模型维度
- 使用梯度累积

### 4. 注意力图不直观
- 检查归一化方法
- 调整温度参数
- 可视化时overlay到源代码行

## 引用

如果您使用了本项目，请引用：

```bibtex
@article{devign-a3m,
  title={Attention-Augmented Attention-Map Aware Devign Model for Vulnerability Detection},
  author={Your Name},
  journal={Your Journal},
  year={2024}
}
```

## 许可证

MIT License

## 联系方式

如有问题，请提交Issue或联系作者。