AMPLE实验复现

  • 2
  • 789 次浏览
  • A+
所属分类：笔记

记录实验复现流程，包括预处理步骤，遇到的错误及解决方案。

原文链接：Vulnerability Detection with Graph Simplification and Enhanced Graph Representation Learning | IEEE Conference Publication | IEEE Xplore

项目地址：https://github.com/AMPLE001/AMPLE

1 预处理

1.1 joern工具下载

首先从作者提供的网址：https://zenodo.org/records/7323504/files/joern.zip?download=1 下载编译好的joern 工具。（不同版本的 joern 生成的 cpg 图可能存在一定的差异，为尽可能复现论文中报告的实验结果，建议使用作者提供的 joern 版本）。

解压至 data_processing 文件夹中，此时的文件夹结构应为：

data_processing
├─ cpg_GS.py
├─ cpg_original.py
├─ cpg_TGS.py
├─ cpg_VGS.py
├─ joern.md
├─ process.py
├─ utils.py
├─ word2vec.py
└─ joern
   ├─ process.py
   ├─ readme.txt
   ├─ slicer.sh
   └─ joern
      └─ joern
         ├─ .gitignore
         ├─ AUTHORS
         ├─ build.gradle
         ├─ build.sh
         ├─ joern-parse
         ├─ ....
         ├─ ....

1.2 数据集预处理

下载 Devign 数据集：function.json ，（需科学上网，文中采用了 Devign, Reveal 和 BigVul 三种数据集，本文以 Devign 数据集为例进行演示），并放在AMPLE 根目录中。

新建 dataset_preprocess.py 文件，用于提取 function.json 中的每一个函数样本，并按照 0.c, 1.c, 2.c, ... 的命名方式生成 C 文件：

import json
import os

# 读取 function.json 文件
with open('function.json', 'r', encoding='utf-8') as file:
    data = json.load(file)

# 提取所有键为 "func" 的值
func_values = []

# 递归提取所有键为 "func" 的值
def extract_func_values(data):
    if isinstance(data, dict):  # 如果是字典
        for key, value in data.items():
            if key == "func":
                func_values.append(value)
            else:
                extract_func_values(value)  # 递归处理嵌套的字典
    elif isinstance(data, list):  # 如果是列表
        for item in data:
            extract_func_values(item)  # 递归处理列表中的每个元素

# 调用提取函数
extract_func_values(data)

# 创建 devign_dataset 文件夹，如果不存在的话
output_folder = 'devign_dataset/devign_raw_code'
os.makedirs(output_folder, exist_ok=True)

# 为每个提取的值生成一个 .c 文件并保存到 devign_dataset 文件夹
for i, func_value in enumerate(func_values):
    filename = os.path.join(output_folder, f"{i}.c")
    
    # 将 func_value 写入 .c 文件
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(func_value)

    print(f"已生成文件: {filename}")

1.3 代码路径修改

原代码中存在一些路径问题，具体修改为：

# AMPLE/data_processing/process.py 文件:
# 第 24 行
"bash ./slicer.sh "  修改为  "bash ./joern/slicer.sh " # 注意后面的空格
# 第 38 行
"../devign_dataset"  修改为  "../devign_dataset/devign_raw_code"
# AMPLE/data_processing/joern/slicer.sh 文件，第 19 行
./joern/joern-parse  修改为  ./joern/joern/joern/joern-parse

执行下列命令，开始调用 joern 解析源代码（解析过程相当漫长，大约4-5小时，慢慢等待即可...）：

cd data_processing
python process.py --start 0 --end 27317

如果执行上述命令遇到 ./joern/slicer.sh: line 19: ./joern/joern/joern/joern-parse: Permission denied 等错误，通过下列命令修改执行权限后再次运行即可：

chmod +x ./joern/joern/joern/joern-parse

2 图简化

安装作者提供的库版本，避免因 torchdata 和 dgl 库版本不符导致的报错：

conda install torch==1.9.0
conda install pytorch::torchdata
conda install pyyaml
conda install pydantic

下载作者训练好的 word2vec 模型：https://zenodo.org/records/7333062#.Y3c5SHZByUk，放入文件夹`devign_dataset/devign_wv_models`中。

新建`devign_dataset/devign_parsed`文件夹，并将`data_processing`文件夹下生成的每个解析结果移动到该文件夹中。

将devign数据集按照8:1:1的比例划分为train、valid、test三个部分，并放入文件夹`devign_dataset/devign_data_split`中。

修改 /AMPLE_code/data_loader/batch_graph.py 以及 /AMPLE_code/data_loader/dataset.py 中的DGLGraph() 为 graph()，以消除下列报警问题：

/home/lxh/anaconda3/envs/ample/lib/python3.9/site-packages/dgl/heterograph.py:72: DGLWarning: Recommend creating graphs by dgl.graph(data) instead of dgl.DGLGraph(data).dgl_warning('Recommend creating graphs by dgl.graph(data)'
# 修改 661、664、668 为按行读取json
# 修改第 649-653 行的超参数路径
# 修改 702 行为 file_name = str(i) + '.c'
# 修改 705 行为 label = int(entry['target'])
python cpg_GS.py

执行完毕后会在`devign_dataset/decign_cpg_c2_2`文件夹中生成9个文件：

bad_file.json
devign-test-num-v2.json
devign-test-v2.json
devign-train-file.json
devign-train-num-v2.json
devign-train-v2.json
devign-valid-file.json
devign-valid-num-v2.json
devign-valid-v2.json

3 模型训练

安装torch（建议安装此版本，避免dgl冲突问题）：

pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/cu111/torch_stable.html

安装dgl-gpu版本（注意torch和cuda版本对应）：

pip install dgl-cu111 -f https://data.dgl.ai/wheels/repo.html

开始训练模型：

cd AMPLE_code
python main.py

发表评论

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

目前评论：2   其中：访客  0   博主  0

    • dleo 0
      回复 2025年6月25日 16:39  沙发

      博主你好，我按照您的方法步骤复现时，在执行cpg_GS.py文件生成decign_cpg_c2_2文件夹内的文件时，生成的文件都是空的，如果方便的话，希望您可以回复更详细的步骤或者复现时的代码，万分感谢

      • jack 0
        回复 2025年9月18日 16:58  板凳

        哥，能复现一下MAGNET那篇论文吗，这是他的文章《Meta-path based Attentional Graph Learning Model for Vulnerability Detection》，他也没给步骤，属实不知道该怎么复现了