Distil-DIRE

Distil-DIRE 是 DIRE 的轻量级版本，可用于实时应用。它不直接计算 DIRE 图像，而是旨在通过带有单步噪声的 DDIM 反演，重建由 image-net 预训练分类器前向传播的对应 DIRE 图像的特征。（论文链接)

预训练 DistilDIRE 检查点

数据集	模型	链接
CelebA-HQ	224x224	链接
ImageNet	224x224	链接

预训练 ADM 扩散模型

我们使用 image-net 预训练的非条件 ADM 扩散模型进行特征重建。您可以从以下链接下载预训练模型： https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt

或者您可以使用以下脚本下载模型：

wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt -O models/256x256-adm.pt

数据准备

在您的自定义数据集上训练模型之前，您需要按照以下格式准备数据集：

mydataset/train|val|test
└── images
    ├── fakes
    │   └──img1.png...
    ├── reals
        └──rimg1.png...

准备数据集后，您可以使用以下脚本计算数据集的 epsilons 和 dire 图像：

bash compute_dire_eps.sh

运行脚本后，您将拥有以下目录结构：

mydataset/train|val|test
└── images
    ├── fakes
    │   └──img1.png...
    ├── reals
        └──rimg1.png...
└── eps
    ├── fakes
    │   └──img1.pt...
    ├── reals
        └──rimg1.pt...
└── dire
    ├── fakes
    │   └──img1.png...
    ├── reals
        └──rimg1.png...

对于 eps 和 dire 的计算，我们将 DDIM 步数设置为 20。推理时也应保持一致。

训练

为了训练 Distil-DIRE，请确保项目根目录下存在 datasets 目录，并且您的数据集位于 datasets 目录内。

torchrun --standalone --nproc_per_node 8 -m train --batch 128 --exp_name truemedia-distil-dire-mydataset --datasets mydataset --datasets_test mytestset --epoch 100 --lr 1e-4

微调

你也可以在自己的数据集上对模型进行微调。对于微调，你需要提供预训练模型的路径。

torchrun --standalone --nproc_per_node 8 -m train --batch 128 --exp_name truemedia-distil-dire-mydataset --datasets mydataset --datasets_test mytestset --epoch 100 --lr 1e-4 --pretrained_weights YOUR_PRETRAINED_MODEL_PATH

测试

要测试模型，你可以使用以下脚本：

python3 -m test --test True --datasets mydataset --pretrained_weights YOUR_PRETRAINED_MODEL_PATH

使用 Docker

export DOCKER_REGISTRY="YOUR_NAME" # Put your Docker Hub username here  
export DATE=`date +%Y%m%d` # Get the current date

# Build the Docker image for development
docker build -t "$DOCKER_REGISTRY/distil-dire:dev-$DATE" -f Dockerfile .


# Push your docker image to docker hub
docker login
docker push "$DOCKER_REGISTRY/distil-dire:dev-$DATE"

开发环境

export WORKDIR="YOUR_WORKDIR" # Put your working directory here
docker run --gpus=all --name=truemedia_gpu_all_distildire -v "$WORKDIR:/workspace/" -ti -e  "$DOCKER_REGISTRY/distil-dire:dev-$DATE"

# work inside the container (/workspace)

注意

• 本仓库运行于在 ImageNet 1k 数据集上训练的 ADM 扩散模型（256x256，无条件）以及在 ImageNet 1k 数据集上训练的 ResNet-50 分类器。
• 最低要求：1 个 GPU，10GB 显存

许可证

本作品采用 Creative Commons Attribution-NonCommercial 4.0 许可证授权。 只要你保留原作者署名和非商业许可，你可以自由阅读、分享和修改此代码。 请参阅 此网站 以获取详细的法律条款。

致谢

我们的代码基于 DIRE、guided-diffusion 和 CNNDetection 开发。感谢他们分享代码和模型。

引用

如果你发现本工作对你的研究有用，请引用我们的论文：

@misc{lim2024distildire,
      title={DistilDIRE: A Small, Fast, Cheap and Lightweight Diffusion Synthesized Deepfake Detection}, 
      author={Yewon Lim and Changyeon Lee and Aerin Kim and Oren Etzioni},
      year={2024},
      eprint={2406.00856},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}