IDDM：集成设计扩散模型

English Document

关于模型

该扩散模型为经典的ddpm、ddim和plms，来源于论文《Denoising Diffusion Probabilistic Models》、《Denoising Diffusion Implicit Models》和《Pseudo Numerical Methods for Diffusion Models on Manifolds》。

我们将此项目命名为IDDM: Integrated Design Diffusion Model，中文名为集成设计扩散模型。在此项目中进行模型复现、训练器和生成器编写、部分算法和网络结构的改进与优化，该仓库持续维护。

如果有任何问题，请先到此issue进行问题查询，若无法解决可以加入我们的QQ群：949120343、开启新issue提问或联系我的邮箱：chenyu1998424@gmail.com/chairc1998@163.com

本仓库整体结构

Integrated Design Diffusion Model
├── config
│   ├── choices.py
│   └── version.py
├── datasets
│   └── dataset_demo
│       ├── class_1
│       ├── class_2
│       └── class_3
├── model
│   ├── modules
│   │   ├── activation.py
│   │   ├── attention.py
│   │   ├── block.py
│   │   ├── conv.py
│   │   ├── ema.py
│   │   └── module.py
│   ├── networks
│   │   ├── sr
│   │   │   └── srv1.py
│   │   ├── base.py
│   │   ├── cspdarkunet.py
│   │   └── unet.py
│   └── samples
│       ├── base.py
│       ├── ddim.py
│       ├── ddpm.py
│       └── plms.py
├── results
├── sr
│   ├── dataset.py
│   ├── demo.py
│   ├── interface.py
│   └── train.py
├── test
│   ├── noising_test
│   │   ├── landscape
│   │   └── noise
│   └── test_module.py
├── tools
│   ├── deploy.py
│   ├── FID_calculator.py
│   ├── FID_calculator_plus.py
│   ├── generate.py
│   └── train.py
├── utils
│   ├── check.py
│   ├── checkpoint.py
│   ├── dataset.py
│   ├── initializer.py
│   ├── logger.py
│   ├── lr_scheduler.py
│   └── utils.py
├── webui
│   └──web.py
└── weight

接下来要做

• 1. 新增cosine学习率优化（2023-07-31）
• 2. 使用效果更优的U-Net网络模型（2023-11-09）
• 3. 更大尺寸的生成图像（2023-11-09）
• 4. 多卡分布式训练（2023-07-15）
• 5. 云服务器快速部署和接口（2023-08-28）
• 6. 增加DDIM采样方法（2023-08-03）
• 7. 支持其它图像生成（2023-09-16）
• 8. 低分辨率生成图像进行超分辨率增强[超分模型效果待定]（2024-02-18）
• 9. 使用Latent方式降低显存消耗
• 10. 重构model整体结构（2023-12-06）
• 11. 编写可视化webui界面（2024-01-23）
• 12. 增加PLMS采样方法（2024-03-12）
• 13. 增加FID方法验证图像质量（2024-05-06）

训练

注意

本自README的训练GPU环境如下：使用具有6GB显存的NVIDIA RTX 3060显卡、具有11GB显存的NVIDIA RTX 2080Ti显卡和具有24GB（总计48GB，分布式训练）显存的NVIDIA RTX 6000（×2）显卡对模型进行训练和测试。上述GPU均可正常训练。

开始你的第一个训练（以cifar10为例，模式单卡）

1. 导入数据集

   首先，将数据集上传至目标文件夹datasets中[issue]。上传后文件夹格式（例如：cifar10文件夹下存放着所有类别；class0文件夹下存储着class0这个类别的所有图片）如下方列表所示：

    datasets
    └── cifar10
        ├── class0
        ├── class1
        ├── class2
        ├── class3
        ├── class4
        ├── class5
        ├── class6
        ├── class7
        ├── class8
        └── class9

   此时你的训练前准备已经完毕。

2. 设置训练参数

   打开train.py文件，修改if __name__ == "__main__":中的parser参数；

   设置--conditional参数为True，因为是多类别训练，所以需要开启，单类别可以不开启也可以开启；

   设置--run_name参数为你想创建的文件名称，例如cifar_exp1；

   设置--dataset_path参数为/你的/本地/或/远程服务器/文件/地址/datasets/cifar10；

   设置--result_path参数为/你的/本地/或/远程服务器/文件/地址/results；

   设置--num_classes参数为10，这是你的类别总数；

   设置更多参数（自定义），如果报CUDA out of memory错误，将--batch_size、--num_workers调小；

   在自定义参数中，你可以设置不同的--sample例如ddpm或ddim，设置不同的训练网络--network例如unet或cspdarkunet。当然激活函数--act，优化器--optim，混合精度训练--amp，学习率方法--lr_func等参数也都是可以自定义设置的。

   详细命令可参考训练参数。

3. 等待训练过程

   点击run运行后，项目会在results文件夹中生成cifar_exp1文件夹，该文件夹中会保存训练日志文件、模型训练文件、模型EMA文件、模型优化器文件、训练的所有最后一次保存的文件和评估后生成的图片。

4. 查看结果

   找到results/cifar_exp1文件夹即可查看训练结果。

↓↓↓↓↓↓↓↓↓↓下方为多种训练方式、训练详细参数讲解↓↓↓↓↓↓↓↓↓↓

普通训练

1. 以landscape数据集为例，将数据集文件放入datasets文件夹中，该数据集的总路径如下/your/path/datasets/landscape，图片存放在/your/path/datasets/landscape/images，数据集图片路径如下/your/path/datasets/landscape/images/*.jpg

2. 打开train.py文件，找到--dataset_path参数，将参数中的路径修改为数据集的总路径，例如/your/path/datasets/landscape

3. 设置必要参数，例如--sample，--conditional，--run_name，--epochs，--batch_size，--image_size，--result_path等参数，若不设置参数则使用默认设置。我们有两种参数设置方法，其一是直接对train.py文件if __name__ == "__main__":中的parser进行设置（我们推荐这种方式）；其二是在控制台在/your/path/Defect-Diffiusion-Model/tools路径下输入以下命令：
   有条件训练命令

   python train.py --sample ddpm --conditional --run_name df --epochs 300 --batch_size 16 --image_size 64 --num_classes 10 --dataset_path /your/dataset/path --result_path /your/save/path

   无条件训练命令

   python train.py --sample ddpm --run_name df --epochs 300 --batch_size 16 --image_size 64 --dataset_path /your/dataset/path --result_path /your/save/path

4. 等待训练即可

5. 若因异常原因中断训练[issue]，我们可以在train.py文件，首先设置--resume为True，其次设置异常中断的迭代编号，再写入该次训练的所在文件夹（run_name），最后运行文件即可。也可以使用如下命令进行恢复：
   有条件恢复训练命令

   # 此处为输入--start_epoch参数，使用当前编号权重
   python train.py --resume --start_epoch 10 --sample ddpm --conditional --run_name df --epochs 300 --batch_size 16 --image_size 64 --num_classes 10 --dataset_path /your/dataset/path --result_path /your/save/path

   # 此处为不输入--start_epoch参数，默认使用last权重
   python train.py --resume --sample ddpm --conditional --run_name df --epochs 300 --batch_size 16 --image_size 64 --num_classes 10 --dataset_path /your/dataset/path --result_path /your/save/path

   无条件恢复训练命令

   python train.py --resume --start_epoch 10 --sample ddpm --run_name df --epochs 300 --batch_size 16 --image_size 64 --dataset_path /your/dataset/path --result_path /your/save/path

   # 此处为不输入--start_epoch参数，默认使用last权重
   python train.py --resume --sample ddpm --run_name df --epochs 300 --batch_size 16 --image_size 64 --dataset_path /your/dataset/path --result_path /your/save/path

6. 预训练模型在每次大版本Release中发布，请留意。预训练模型使用方法如下[issue]，首先将对应network、image_size、act等相同参数的模型下到本地任意文件夹下。直接调整train.py中--pretrain和--pretrain_path即可。也可以使用如下命令进行预训练：
   使用有条件预训练模型训练命令

   python train.py --pretrain --pretrain_path /your/pretrain/path/model.pt --sample ddpm --conditional --run_name df --epochs 300 --batch_size 16 --image_size 64 --dataset_path /your/dataset/path --result_path /your/save/path

   使用无条件预训练模型训练命令

   python train.py --pretrain --pretrain_path /your/pretrain/path/model.pt --sample ddpm --run_name df --epochs 300 --batch_size 16 --image_size 64 --dataset_path /your/dataset/path --result_path /your/save/path

分布式训练

1. 基本配置与普通训练相似，值得注意的是开启分布式训练需要设置--distributed。为了防止随意设置分布式训练，我们为开启分布式训练设置了几个基本条件，例如args.distributed、torch.cuda.device_count() > 1和torch.cuda.is_available()。

2. 设置必要的参数，例如--main_gpu和--world_size。--main_gpu通常设置为主要GPU，例如做验证、做测试或保存权重，我们仅在单卡中运行即可。而world_size的值会与实际使用的GPU数量或分布式节点数量相对应。

3. 我们有两种参数设置方法，其一是直接对train.py文件if __name__ == "__main__":中的parser进行设置；其二是在控制台在/your/path/Defect-Diffiusion-Model/tools路径下输入以下命令：

   有条件训练命令

   python train.py --sample ddpm --conditional --run_name df --epochs 300 --batch_size 16 --image_size 64 --num_classes 10 --dataset_path /your/dataset/path --result_path /your/save/path --distributed --main_gpu 0 --world_size 2

   无条件训练命令

   python train.py --sample ddpm --run_name df --epochs 300 --batch_size 16 --image_size 64 --dataset_path /your/dataset/path --result_path /your/save/path --distributed --main_gpu 0 --world_size 2

4. 等待训练即可，中断恢复同基本训练一致。

训练参数

参数讲解

生成

1. 打开generate.py文件，找到--weight_path参数，将参数中的路径修改为模型权重路径，例如/your/path/weight/model.pt

2. 设置必要参数，例如--conditional，--generate_name，--num_images，--num_classes，--class_name，--image_size，--result_path等参数，若不设置参数则使用默认设置。我们有两种参数设置方法，其一是直接对generate.py文件if __name__ == "__main__":中的parser进行设置；其二是在控制台在/your/path/Defect-Diffiusion-Model/tools路径下输入以下命令：
   有条件生成命令（1.1.1版本以上）

   python generate.py --generate_name df --num_images 8 --class_name 0 --image_size 64 --weight_path /your/path/weight/model.pt --sample ddpm

   无条件生成命令（1.1.1版本以上）

   python generate.py --generate_name df --num_images 8 --image_size 64 --weight_path /your/path/weight/model.pt --sample ddpm

   有条件生成命令（1.1.1版本及以下）

   python generate.py --conditional --generate_name df --num_images 8 --num_classes 10 --class_name 0 --image_size 64 --weight_path /your/path/weight/model.pt --sample ddpm --network unet --act gelu 

   无条件生成命令（1.1.1版本及以下）

   python generate.py --generate_name df --num_images 8 --image_size 64 --weight_path /your/path/weight/model.pt --sample ddpm --network unet --act gelu 

3. 等待生成即可

生成参数

参数讲解

结果

我们在以下4个数据集做了训练，开启conditional，采样器为DDPM，图片尺寸均为64*64，激活函数为gelu，学习率为3e-4，采用线性学习方法，迭代次数为300，分别是cifar10，NEUDET，NRSD-MN和WOOD。结果如下图所示：

cifar10

NEUDET

NRSD

WOOD

Animate face（整活生成）

基于64×64模型生成160×160（任意大尺寸）图像

当然，我们根据64×64的基础模型，在generate.py文件中生成160×160的NEU-DET图片（单张输出，每张图片占用显存21GB）。请注意这个[issue]：如果是缺陷纹理那种图片，特征物不明显的直接生成大尺寸就不会有这些问题，例如NRSD、NEU数据集。如果是含有背景有特定明显特征的则需要超分或者resize提升尺寸，例如Cifar10、CelebA-HQ等。**如果实在需要大尺寸图像，在显存足够的情况下直接训练大像素图片。**详细图片如下：

部署

未完待续

关于引用

如果在学术论文中使用该项目进行实验，在可能的情况下，请适当引用我们的项目，为此我们表示感谢。具体引用格式可访问此网站。

@software{chen_2024_10866128,
  author       = {Chen Yu},
  title        = {IDDM: Integrated Design Diffusion Model},
  month        = mar,
  year         = 2024,
  publisher    = {Zenodo},
  doi          = {10.5281/zenodo.10866128},
  url          = {https://doi.org/10.5281/zenodo.10866128}
}

致谢

人员：

@dome272

组织：

@JetBrains

@Python

@Pytorch