vishgit19/Duet-MalCross

# Duet-MalCross 静态-动态多模态恶意软件分类。 这是一个用于 8 类 Windows 恶意软件分类的 PyTorch 代码库， 它结合了**静态**（EMBER 2024）模态和**动态**（CAPE 行为图）模态，在**两阶段** pipeline 中进行训练，并使用了 **cross-attention fusion** 模块和**学习型 MLP logit stacker**。 ## 架构 **阶段 2（训练）** — 仅训练 fusion 模块和 concat 分类器； 冻结的 encoder 生成 token 序列，而不是池化后的 embedding： ``` static input (EMBER 2628-d) graph input (PyG Data) │ │ ▼ ▼ static_encoder.forward_tokens graph_encoder.forward_tokens │ │ └──────────┐ ┌──────────────┘ ▼ ▼ CrossAttentionFusion (two cross-attn encoders, M1 & M2) │ m1_cls m2_cls │ │ └───── cat ────┘ │ ▼ stage-2 concat classifier ──► concat_logits (= final during training) ``` **评估（激活 MLP stacker）** — 也会运行冻结的单模态分类器， 以便 stacker 重新组合三个分支的 logit： ``` static input ─► static_encoder ─► static_classifier ─► static_logits ─┐ │ graph input ─► graph_encoder ─► dynamic_classifier ─► dynamic_logits ┤ │ [fusion path above] ─► concat_logits ──────────────────────────────────┤ ▼ MLPStacker([static, dynamic, concat]) ──► final_logits ``` 静态 encoder 是一个带有 block-pair-bias attention 和 drop-path 的 Transformer；图 encoder 使用了 `TransformerConv` 消息传递机制以及 type-wise attention pooling。 ## 两阶段训练 **阶段 1 — 独立的 encoder 预训练。** 每个 encoder 都 独立训练，并配有各自的单模态分类器，在验证集 macro-F1 上进行 early stopping。每种模态的最佳 checkpoint 会 被单独保存。 ``` static_encoder + static_classifier -> static_logits (CE loss) graph_encoder + dynamic_classifier -> dynamic_logits (CE loss) ``` **阶段 2 — fusion + concat 分类器训练。** 两个 encoder 都 从它们阶段 1 的 checkpoint 中加载并冻结。仅训练 `CrossAttentionFusion` 模块（两个 cross-attention encoder M1/M2）以及 阶段 2 的 concat 分类器。然后，在验证集上对三个分支的 logit（静态、 动态、concat）事后拟合 MLP stacker — 不存在训练集数据泄露。 ## 仓库结构 ``` config/ config.py # single config dataclass — all hyperparameters duet_dataset/ csv/ # label_map.json + split manifests (CSV) download.py # MalwareBazaar safe-ZIP downloader static_process.py # EMBER 2024 feature extraction from PE files dynamic_process.py # CAPE report.json -> behavior graph model/ static_encoder.py # EMBER Transformer (block-pair-bias attention) graph_encoder.py # relational graph encoder (TransformerConv) fusion.py # CrossAttentionFusion + MLPStacker multimodal_model.py # top-level model wiring classification_heads.py pooling.py trainer/ trainer.py # two-stage trainer + MLP stacker fit pipeline.py # data pipeline assembly utils/ ember_layout.py # EMBER TOKEN_SPECS (token contract) ember_preprocess.py # semantic_preprocess_pe_v2 (2568 -> 2628) transforms.py # EmberSemanticScaler graph_builder.py # parsed dict -> PyG Data vocab.py # VocabSet, fit_vocabs dataset.py # MultimodalDataset collate.py # batch collation static_loader.py # mmap'd .npy shard loader dynamic_loader.py # behavior-graph JSON loader level_features.py # graph/walk-level scalar features manifest.py # CSV manifest + label_map loading metrics.py # macro-F1, per-branch evaluation seed.py, io.py, hashing.py train.py eval.py requirements.txt ``` ## 环境配置 ``` pip install -r requirements.txt ``` 这将安装 PyTorch、PyTorch-Geometric、scikit-learn、requests 和 [EMBER](https://github.com/elastic/ember) 库。 ## 数据准备 1. **下载样本** 从 MalwareBazaar： export MALWAREBAZAAR_AUTH_KEY="your-key" # 在 duet_dataset/csv/hashes.txt 中每行一个 SHA256 python duet_dataset/download.py ZIP 文件将保存到 `data/samples_zip/`（密码：`infected`）。 2. **提取静态特征**（EMBER 2024，2568 维）从 PE 文件中： python duet_dataset/static_process.py \ --input data/samples/extracted \ --output data \ --shard-size 4096 将分片的 `.npy` 文件写入 `data/static/`，并将清单写入 `data/csv/static_manifest.csv`。 3. **构建行为图** 从 CAPE 沙箱报告中： from duet_dataset.dynamic_process import build_from_report_path graph = build_from_report_path("path/to/report.json") # graph 包含: nodes, edges, temporal_walk, graph_features, # walk_features, summary_features 将图字典保存为 `.json` 或 `.json.gz`（每个样本一个文件）。 4. **组装划分的 CSV**（`train.csv`, `val.csv`, `test.csv`）放在 `duet_dataset/csv/` 中。有关 所需的列，请参见 `duet_dataset/csv/manifest_schema.md`。每一行将样本的静态分片路径、 动态图路径和类标签结合在一起。 ## 训练 编辑 `config/config.py`（路径、超参数、阶段开关），然后： ``` python train.py ``` 进行快速冒烟测试（微小的子集，2 个 epoch）： ``` # 在 config.py 中 smoke_test = True ``` 跳过阶段 1，仅使用现有的 encoder checkpoint 重新运行阶段 2： ``` run_stage1_static = False run_stage1_dynamic = False run_stage2 = True # 确保 static_checkpoint / dynamic_checkpoint 指向有效的文件 ``` ## 评估 ``` python eval.py ``` 报告各个分支（静态、动态、concat、final）在验证集和测试集上的 macro-F1 和 准确率。`final` 分支是 MLP stacker 的输出。 ## 配置 所有超参数都位于 `config/config.py` 中的一个单独的 `Config` dataclass 中。没有 CLI，也没有 YAML。关键部分： - **路径** — `dataset_root`, `output_root`, checkpoint 路径 - **架构** — encoder 维度，fusion heads/layers，dropout - **训练阶段** — `run_stage1_static`, `run_stage1_dynamic`, `run_stage2`，每个阶段的 LR/epochs/patience - **MLP stacker** — `mlp_stacker_hidden`, `mlp_stacker_epochs`, `mlp_stacker_lr` - **优化** — AdamW，余弦调度，AMP，label smoothing， class-weighted loss ## 注意事项 - 静态 encoder 需要 EMBER 2024 特征向量（原始为 2568 维， 通过 `semantic_preprocess_pe_v2` 扩展到 2628 维）。 `EmberSemanticScaler` 负责处理这种扩展；仅在训练集上拟合它。 - 动态 encoder 使用 carrier graph（而不是 PyG `HeteroData`）： 节点/边类型和特征被编码为 embedding，可以轻松地进行批次化， 并且对空图具有鲁棒性。 - MLP stacker 仅在验证集上训练（没有训练集数据泄露）。其 权重单独保存为 `mlp_stacker.pt`。

标签：Apex, PyTorch, 凭据扫描, 多模态融合, 机器学习, 网络攻击, 逆向工具