TabDDPM：用扩散模型建模表格数据

这是我们的论文《TabDDPM：用扩散模型建模表格数据》的官方代码（论文）。

环境搭建

1. 安装 conda（仅用于管理环境）。
2. 运行以下命令：

   export REPO_DIR=/path/to/the/code
   cd $REPO_DIR
   
   conda create -n tddpm python=3.9.7
   conda activate tddpm
   
   pip install torch==1.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html
   pip install -r requirements.txt
   
   # 如果以下命令执行失败，请更新 conda
   conda env config vars set PYTHONPATH=${PYTHONPATH}:${REPO_DIR}
   conda env config vars set PROJECT_DIR=${REPO_DIR}
   
   conda deactivate
   conda activate tddpm
   

实验运行

此处介绍重现实验结果所需的信息。
使用 agg_results.ipynb 可以打印所有数据集和所有方法的结果。

数据集

我们上传了论文中使用的数据集及其训练/验证/测试划分（链接见下文）。我们不对原始数据集的许可施加额外限制，数据来源已在论文附录中列出。

您可以通过以下命令加载数据集：

conda activate tddpm
cd $PROJECT_DIR
wget "https://www.dropbox.com/s/rpckvcs3vx7j605/data.tar?dl=0" -O data.tar
tar -xvf data.tar

文件结构

tab-ddpm/ —— 所提出方法的实现
tuned_models/ —— 评估模型（CatBoost 或 MLP）的调优超参数

所有主要脚本均位于 scripts/ 文件夹中：

• scripts/pipeline.py 用于根据给定配置训练、采样并评估 TabDDPM
• scripts/tune_ddpm.py —— 调优 TabDDPM 的超参数
• scripts/eval_[catboost|mlp|simple].py —— 使用调优后的评估模型或简单模型评估合成数据
• scripts/eval_seeds.py —— 使用多次采样和多个评估种子进行评估
• scripts/eval_seeds_simple.py —— 使用多次采样和多个评估种子进行评估（针对简单模型）
• scripts/tune_evaluation_model.py —— 调优评估模型（CatBoost 或 MLP）的超参数
• scripts/resample_privacy.py —— 隐私计算

实验文件夹 (exp/)：

• 所有结果和合成数据均存储在 exp/[ds_name]/[exp_name]/ 文件夹中
• exp/[ds_name]/config.toml 是用于调优 TabDDPM 的基础配置
• exp/[ds_name]/eval_[catboost|mlp].json 存储评估结果（scripts/eval_seeds.py）

要了解 config.toml 文件的结构，请阅读 CONFIG_DESCRIPTION.md。

基线方法：

• smote/
• CTGAN/ —— TVAE 官方仓库
• CTAB-GAN/ —— 官方仓库
• CTAB-GAN-Plus/ —— 官方仓库

示例

运行 TabDDPM 调优。

模板与示例（--eval_seeds 为可选）：

python scripts/tune_ddpm.py [ds_name] [train_size] synthetic [catboost|mlp] [exp_name] --eval_seeds
python scripts/tune_ddpm.py churn2 6500 synthetic catboost ddpm_tune --eval_seeds

运行 TabDDPM 流水线。

模板与示例（--train、--sample、--eval 为可选）：

python scripts/pipeline.py --config [path_to_your_config] --train --sample --eval
python scripts/pipeline.py --config exp/churn2/ddpm_cb_best/config.toml --train --sample

上述脚本在 NVIDIA GeForce RTX 2080 Ti 上大约需要 7 分钟运行。

运行基于种子的评估
在进行评估之前，您必须使用给定的超参数训练模型（如上例所示）。

模板与示例：

python scripts/eval_seeds.py --config [path_to_your_config] [n_eval_seeds] [ddpm|smote|ctabgan|ctabgan-plus|tvae] synthetic [catboost|mlp] [n_sample_seeds]
python scripts/eval_seeds.py --config exp/churn2/ddpm_cb_best/config.toml 10 ddpm synthetic catboost 5

TabDDPM 快速上手指南

环境准备

系统要求

• 操作系统：支持 Linux 或 macOS
• Python 版本：3.9.7（推荐使用 Conda 管理环境）
• CUDA 支持：建议使用 NVIDIA GPU 进行训练和采样，以加速计算过程

前置依赖

• Conda：用于管理虚拟环境
• PyTorch 1.10.1 + CUDA 11.1
• 其他依赖项见 requirements.txt

---

安装步骤

1. 设置项目目录路径：

   export REPO_DIR=/path/to/the/code
   cd $REPO_DIR
   

2. 创建并激活 Conda 虚拟环境：

   conda create -n tddpm python=3.9.7
   conda activate tddpm
   

3. 安装 PyTorch（国内用户可使用清华源加速）：

   pip install torch==1.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html
   

4. 安装其他依赖：

   pip install -r requirements.txt
   

5. 设置环境变量（可选）：

   conda env config vars set PYTHONPATH=${PYTHONPATH}:${REPO_DIR}
   conda env config vars set PROJECT_DIR=${REPO_DIR}
   

6. 重新激活环境使配置生效：

   conda deactivate
   conda activate tddpm
   

---

基本使用

下载数据集

运行以下命令下载论文中使用的数据集：

conda activate tddpm
cd $PROJECT_DIR
wget "https://www.dropbox.com/s/rpckvcs3vx7j605/data.tar?dl=0" -O data.tar
tar -xvf data.tar

示例：运行 TabDDPM 调参流程

python scripts/tune_ddpm.py churn2 6500 synthetic catboost ddpm_tune --eval_seeds

示例：运行 TabDDPM 训练、采样与评估流程

python scripts/pipeline.py --config exp/churn2/ddpm_cb_best/config.toml --train --sample

注意：此脚本在 NVIDIA GeForce RTX 2080 Ti 上大约需要 7 分钟完成。

示例：使用多个种子进行评估

python scripts/eval_seeds.py --config exp/churn2/ddpm_cb_best/config.toml 10 ddpm synthetic catboost 5

在执行此命令前，请确保已使用对应超参数训练好模型。