【图像生成2】本地部署 LDM 并训练自己的数据集

西瓜和拾月

已于 2025-05-25 12:29:44 修改

阅读量1k

点赞数 20

CC 4.0 BY-SA版权

分类专栏：人工智能文章标签：深度学习 stable diffusion 计算机视觉

于 2025-05-25 12:29:13 首次发布

本文链接：https://blog.youkuaiyun.com/csdnn345/article/details/147728471

人工智能专栏收录该内容

2 篇文章

订阅专栏

一、背景

当需要检测的目标出现的概率很低但又必须检测的时候，直接问题就是目标数据获取困难，正负样本占比差异大，很难在短时间内收集足够数据，也没有公开数据集，这就大大限制了模型的学习能力。目前能想到的方法（1）小样本学习，使用针对小规模数据设计的网络模型；（2）使用某些多尺寸模型中较小的那个，保证模型规模和数据规模相差不会太悬殊，调整参数不断优化；（3）使用数据增强扩充数据集。

常用的数据增强方式比如旋转，缩放，裁剪，拼接，高斯模糊等等，都是对现有数据做一些变换，目标特征本质上没有发生变化，没有产生新的同类目标，也没有场景变化，对于模型的泛化性提升有限。

而近几年大火的扩散模型可谓是生成图像的一把利器，与 GAN 相比其训练过程更加稳定，生成的图像质量更好，且需要的训练数据相对更少，于是尝试一下是否能用扩散模型对自己的数据进行增强。结合数据情况和硬件资源，选择使用 Latent Diffusion Models (LDMs) 。

本文记录 LDMs 的本地部署训练，和无条件图像生成。

二、本地部署

代码仓库：https://github.com/CompVis/latent-diffusion

1. 环境搭建

官方的 environment.yaml 文件提供了一键环境配置，运行以下命令即可完成环境创建和库安装。

conda env create -f environment.yaml

参考网上资料[1]来看，直接使用上述命令安装，由于库之间的版本问题大概率会创建失败，所以最好是手动安装 python 和 pytorch 库，然后创建一个 requirements.txt，把其它需要安装的库加进去，具体版本可能需要根据自己的 python 和 torch 版本做调整。

albumentations==0.4.3
opencv-python==4.1.2.30
pudb==2019.2
imageio==2.9.0
imageio-ffmpeg==0.4.2
pytorch-lightning==1.4.2
omegaconf==2.1.1
test-tube>=0.7.5
streamlit>=0.73.1
einops==0.3.0
torch-fidelity==0.3.0
transformers==4.6.0
torchmetrics==0.6
kornia==0.5.10
six
-e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
-e git+https://github.com/openai/CLIP.git@main#egg=clip
-e .

安装之前手动下载 taming-transformer 和 clip，放到 src 文件夹（没有就创建一个）里，在安装过程中有如下提示，选忽略即可。

安装过程中会自动安装高版本 torch（我这里安装的是 2.4.1），这可能会导致程序运行失败[1]，所以要卸载掉重新安装跟自己 cuda 配套的 torch，比如我安装的是 torch1.7.0+cu110。

2. 测试

首先需要下载预训练权重（model-zoo），不同任务对应不同模型，下载后放到对应位置。

官方提供了很多测试用例，有条件生成（文生图、图像修复、类别条件生成）和无条件生成，如果要测试文生图，需要用到 Bert Tokenizer，主要用于将文本分割成 token，然后将这些token 转换为模型可以处理的数字 ID。可以按照[1]中的方法配置 Bert 后测试。

本文主要测试的无条件生成任务，首先下载预训练权重（官方提供了 4 种预训练 LDMs，分别对应了 4 种训练数据集，以 CelebA 为例），将下载的 model.ckpt 放到 latent-diffusion/models/ldm/celeba256 中，然后运行以下命令生成图像

# -n: 生成数量
# -c: ddim steps，DDIM加速采样步数，值越小速度越快
# -e: eta，DDIM差值因子，一般设置为0
CUDA_VISIBLE_DEVICES=0 python scripts/sample_diffusion.py -r models/ldm/celeba256/model.ckpt -n 5 --batch_size 1 -c 20 -e 0

生成图像如下：

3. 训练自己的LDM

LDM 的训练分两步，首先需要训练一个 AutoEncoder，再配合 AutoEncoder 训练 LDM。

（1）准备数据集

官方提供了两种数据类，LSUN 和 ImageNet，如果使用的是这两种数据集，那按照官方说明下载数据集并放置到正确位置即可。本文使用的是其他数据集（训练集 233 张，验证集 25 张），所以要仿照官方代码修改自己的数据类。

在 latent-diffusion/ldm/data/lsun.py 中添加数据类

class LSUNBottleTrain(LSUNBase):
    def __init__(self, **kwargs):
        super().__init__(txt_file="data/train/bottle_train.txt", data_root="data/train", **kwargs)


class LSUNBottleValidation(LSUNBase):
    def __init__(self, flip_p=0., **kwargs):
        super().__init__(txt_file="data/validation/bottle_val.txt", data_root="data/validation", flip_p=flip_p, **kwargs)

我们的数据类继承自 LSUNBase，所以要求数据集的存放格式要与 LSUN 数据集一致，如下

data
 |--train
     |--bottle_train.txt
     |--img001.jpg
     |--img002.jpg
     ...
 |--validation
     |--bottle_val.txt
     |--img001.jpg
     |--img002.jpg
     ...

（2）训练 AutoEncoder

官方提供了多种 AE，本文使用的是 kl-f8。修改configs/autoencoder/autoencoder_kl_32x32x64.yaml，主要修改 data 部分。

data:
  target: main.DataModuleFromConfig
  params:
    batch_size: 2
    wrap: true
    num_workers: 0
    train:
      target: ldm.data.lsun.LSUNBottleTrain # 改为自己的数据类
      params:
        size: 256
        #degradation: pil_nearest #注释掉该参数，否则报错
    validation:
      target: ldm.data.lsun.LSUNBottleVal # 改为自己的数据类
      params:
        size: 256
        #degradation: pil_nearest #注释掉该参数，否则报错

kl_32x32x64 对应，，batch_size 和 size 参数根据自己的显存设置，上述设置是在自己的 GPU（RTX3060，12G 占用 11G）上的极限值，再多一点就 OOM 了，同时需要将 num_workers 设置为 0，否则训练过程中也会报错。

开始训练。如果是单 GPU 训练，不要漏掉最后的逗号，CUDA_VISIBLE_DEVICES 和 --gpus 的 GPU id 要一致。

CUDA_VISIBLE_DEVICES=0 python main.py --base configs/autoencoder/autoencoder_kl_32x32x64.yaml -t --gpus 0,

如果没有指定 --max_epochs，默认训练 1000 个 epoch（大概跑不到 2 天），日志文件和checkpoints 保存在 logs 文件夹中，默认保留训练损失前三低的 checkpoints。

如果使用预训练权重，下载对应模型放到 models/first_stage_models/kl-f8 下，使用 --resume 指定。

CUDA_VISIBLE_DEVICES=0 python main.py --base configs/autoencoder/autoencoder_kl_32x32x64.yaml --resume models/first_stage_models/kl-f8/model.ckpt -t --gpus 0,

使用在大数据集上预训练的权重可以加速模型收敛。

（3）训练 LDM

仿照 configs/latent-diffusion/lsun_churches-ldm-kl-8.yaml 创建配置文件 xxx-ldm-kl-8.yaml，修改 model--params--first_stage_config--ckpt_path 参数为自己训练好的 autoencoder 路径

修改 data 参数，与 autoencoder_kl_32x32x64.yaml 中的一致即可。

开始训练

CUDA_VISIBLE_DEVICES=0 python main.py --base configs/latent-diffusion/bottle-ldm-kl-8.yaml -t --gpus 0,

（4）测试（无条件生成）

在 models/ldm 下创建文件夹 bottle256，将自己训练好的 LDM ckpt 放到该目录下，开始测试

CUDA_VISIBLE_DEVICES=0 python scripts/sample_diffusion.py -r models/ldm/<model_spec>/xxx.ckpt -n 20 --batch_size 2 -c 20 -e 0

生成的图像在 bottle256/samples/xxx/<time>/img 文件夹下，分辨率 256256。

部分生成效果还是可以的，当然也有一些效果不好的，如果数据量再多一些，再调调训练参数，效果可能会更好。

三、总结

总体来看 LDM 的使用不算复杂，而且生成效果还不错。在使用 LDM 之前试过 DDPM，训练时间太长，生成效果不如 LDM。因为目前只需要无条件生成，对于有条件生成没有测试过，后面如果用到了再具体研究吧。

参考资料

[1] latent-diffusion model环境配置，这可能是你能够找到的最细的博客了_latent diffusion model 训练 autoencoder-优快云博客

[2] Latent Diffusion部署&生成Unconditional Model - 哔哩哔哩

[3] BertTokenizer的使用方法(超详细)_from transformers import berttokenizer-优快云博客

[4] DDIM 简明讲解与 PyTorch 实现：加速扩散模型采样的通用方法 -知乎

[5] Stable Diffusion 图像生成攻略二 -知乎

[6] 使用DDPM(扩散模型)训练自己的数据集实现数据集扩容pytorch_ddpm数据处理-优快云博客