-
Hardware Requirements and Software Requirements
- GPU: 1x high-end NVIDIA GPU with at least 32GB memory
- Python: 3.8
- CUDA: 11.7
- cuDNN: 8.9.7
To install other requirements:
pip install -r requirements.txt -
Datasets and Pretrain-weights
Please download the dataset GenImage and the data root directory should be organized as follows.
├── ADM
│ ├── imagenet_ai
│ │ ├── train
│ │ │ ├── ai
│ │ │ ├── nature
│ │ ├── val
│ │ │ ├── ai
│ │ │ ├── nature
├── BigGAN
│ ├── imagenet_ai
│ │ ├── train
│ │ │ ├── ai
│ │ │ ├── nature
│ │ ├── val
│ │ │ ├── ai
│ │ │ ├── nature
├── Glide
│ ├── imagenet_ai
│ │ ├── train
│ │ │ ├── ai
│ │ │ ├── nature
│ │ ├── val
│ │ │ ├── ai
│ │ │ ├── nature
├── Midjourney
│ ├── ...
├── Stable Diffusion V1.4
│ ├── ...
├── Stable Diffusion V1.5
│ ├── ...
├── VQDM
│ ├── ...
├── Wukong
│ ├── ...
The pre-trained weights are located at checkpoints, with the extraction code z8se. Place them in the checkpoints folder.
To train a detector with full GenImage training set, run this command:
python train.py --exam_dir <log path> --name <model name> --blur_prob 0.5 --blur_sig 0.0,3.0 --jpg_prob 0.5 --jpg_method cv2,pil --jpg_qual 30,100 --dataroot <dataset path> --classes "ADM/imagenet_ai,BigGAN/imagenet_ai,glide/imagenet_ai,Midjourney/imagenet_ai,stable_diffusion_v_1_4/imagenet_ai,stable_diffusion_v_1_5/imagenet_ai,VQDM/imagenet_ai,wukong/imagenet_ai" --gpu_ids <gpu id> --loadSize <load size> --cropSize <crop size> --batch_size <batch size> --class_bal --arch <model type>
You can use the above commands to train Efficientnet-B0 (E), ResNet50 (R), DeiT (D), and Swin-T (S) models. After training, you can test the performance of the model in demo_dir.py. Replace "efficientnet-b0," "resnet50," "deit," and "swin-t" in the dic with the corresponding training weight paths. run this command:
python demo_dir.py -d <test dataset> -m <model path>
To train Control-VAE with full GenImage training set, run this command:
python train_vae.py --iter <training iterations> --save_dir <save path> --dir "/path/to/genimage"
First, you should replace E, R, D, and S in other_attacks.py's dic with the corresponding training weight paths.
To generate adversarial examples, run this command:
python main.py --model_name <surrogate model> --save_dir <save path> --images_root </path/to/GenImage> --encoder_weights </path/to/Controlvae.pt>
Some of the codes are built upon DiffAttack and CNNDetection. The experimental dataset comes from GenImage. Thanks them for their great works!