Stable Diffusion火影数据集训练:SwanLab可视化训练

本文涉及的产品
实时数仓Hologres,5000CU*H 100GB 3个月
智能开放搜索 OpenSearch行业算法版,1GB 20LCU 1个月
大数据开发治理平台 DataWorks,不限时长
简介: **使用Stable Diffusion 1.5模型训练火影忍者风格的文生图模型。在22GB显存的GPU上,通过Huggingface的`lambdalabs/naruto-blip-captions`数据集进行训练,利用SwanLab进行监控。所需库包括`swanlab`, `diffusers`, `datasets`, `accelerate`, `torchvision`, `transformers`。代码、日志和更多资源可在GitHub和SwanLab找到。训练涉及数据下载、模型配置、训练过程可视化及结果评估。**

Stable Diffusion 1.5(SD1.5)是由Stability AI在2022年8月22日开源的文生图模型,是SD最经典也是社区最活跃的模型之一。

以SD1.5作为预训练模型,在火影忍者数据集上微调一个火影风格的文生图模型(非Lora方式),是学习SD训练的入门任务。

sd-1.png

显存要求 22GB左右

在本文中,我们会使用SD-1.5模型在火影忍者数据集上做训练,同时使用SwanLab监控训练过程、评估模型效果。

1.环境安装

本案例基于Python>=3.8,请在您的计算机上安装好Python;

另外,您的计算机上至少要有一张英伟达显卡(显存大约要求22GB左右)。

我们需要安装以下这几个Python库,在这之前,请确保你的环境内已安装了pytorch以及CUDA:

swanlab
diffusers
datasets
accelerate
torchvision
transformers

一键安装命令:

pip install swanlab diffusers datasets accelerate torchvision transformers
本文的代码测试于diffusers==0.29.0、accelerate==0.30.1、datasets==2.18.0、transformers==4.41.2、swanlab==0.3.11,更多库版本可查看 SwanLab记录的Python环境

2.准备数据集

本案例是用的是火影忍者数据集,该数据集主要被用于训练文生图模型。

该数据集由1200条(图像、描述)对组成,左边是火影人物的图像,右边是对它的描述:

sd-2.png

我们的训练任务,便是希望训练后的SD模型能够输入提示词,生成火影风格的图像:

sd-3.png


数据集的大小大约700MB左右;数据集的下载方式有两种:

  1. 如果你的网络与HuggingFace连接是通畅的,那么直接运行我下面提供的代码即可,它会直接通过HF的datasets库进行下载。
  2. 如果网络存在问题,我也把它放到百度网盘(提取码: gtk8),下载naruto-blip-captions.zip到本地解压后,运行到与训练脚本同一目录下。

3.准备模型

这里我们使用HuggingFace上Runway发布的stable-diffusion-v1-5模型。

sd-4.png

模型的下载方式同样有两种:

  1. 如果你的网络与HuggingFace连接是通畅的,那么直接运行我下面提供的代码即可,它会直接通过HF的transformers库进行下载。
  2. 如果网络存在问题,我也把它放到百度网盘(提取码: gtk8),下载stable-diffusion-v1-5.zip到本地解压后,运行到与训练脚本同一目录下。

4. 配置训练可视化工具

我们使用SwanLab来监控整个训练过程,并评估最终的模型效果。

如果你是第一次使用SwanLab,那么还需要注册一个账号,在用户设置页面复制你的API Key,然后在训练开始时粘贴进去即可:

sd-swanlab-login.png

5.开始训练

由于训练的代码比较长,所以我把它放到了Github里,请Clone里面的代码:

git clone https://github.com/Zeyi-Lin/Stable-Diffusion-Example.git

如果你与HuggingFace的网络连接通畅,那么直接运行训练:

python train_sd1-5_naruto.py \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --seed=42 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model"

如果你的模型或数据集用的是上面的网盘下载,那么你需要做下面的两件事:

第一步:将数据集和模型文件夹放到训练脚本同一目录下,文件结构如下:

|--- sd_config.py
|--- train_sd1-5_naruto.py
|--- stable-diffusion-v1-5
|--- naruto-blip-captions

stable-diffusion-v1-5是下载好的模型文件夹,naruto-blip-captions是下载好的数据集文件夹。

第二步:修改sd_config.py的代码,将pretrained_model_name_or_pathdataset_name的default值分别改为下面这样:

    parser.add_argument(
        "--pretrained_model_name_or_path",
        type=str,
        default="./stable-diffusion-v1-5",
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="./naruto-blip-captions",
    )

然后运行启动命令即可。


看到下面的进度条即代表训练开始:

sd-10.png

6. 训练结果演示

我们在SwanLab上查看最终的训练结果:

sd-5.png

可以看到SD训练的特点是loss一直在震荡,随着epoch的增加,loss在最初下降后,后续的变化其实并不大:

sd-6.png

我们来看看主观生成的图像,第一个epoch的图像长这样:

sd-7.png

可以看到詹姆斯还是非常的“原生态”,迈克尔杰克逊生成的也怪怪的。。。

再看一下中间的状态:

sd-8.png

sd-9.png

经过比较长时间的训练后,效果就好了不少。

比较有意思的是,比尔盖茨生成出来的形象总是感觉非常邪恶。。。

详细训练过程看这里:SD-Naruto - SwanLab

至此,你已经完成了SD模型在火影忍者数据集上的训练。

相关文章
|
2月前
|
机器学习/深度学习 自然语言处理
文生图模型-Stable Diffusion | AIGC
所谓的生成式模型就是通过文本或者随机采样的方式来得到一张图或者一段话的模型,比如文生图,顾名思义通过文本描述来生成图像的过程。当前流行的文生图模型,如DALE-2, midjourney以及今天要介绍的Stable Diffusion,这3种都是基于Diffusion扩散模型【1月更文挑战第6天】
550 0
|
9月前
|
PyTorch 测试技术 API
Pytorch模型训练与在线部署
本文以CIFAR10数据集为例,通过自定义神经元网络,完成模型的训练,并通过Flask完成模型的在线部署与调用,考略到实际生产模型高并发调用的述求,使用service_streamer提升模型在线并发能力。
608 1
|
存储 机器学习/深度学习 算法
MMDetection3d对KITT数据集的训练与评估介绍
MMDetection3d对KITT数据集的训练与评估介绍
1618 0
MMDetection3d对KITT数据集的训练与评估介绍
|
1月前
|
计算机视觉
【YOLOv10训练教程】如何使用YOLOv10训练自己的数据集并且推理使用
【YOLOv10训练教程】如何使用YOLOv10训练自己的数据集并且推理使用
|
2月前
|
机器学习/深度学习 人工智能 边缘计算
为何人们喜欢推理胜于训练大模型?
在AI和机器学习领域,越来越多的人转向重视推理而非大规模模型训练。推理的即时性和高效性使其在需要快速响应的场景中占优,如自然语言处理和图像识别。推理过程的可视化能帮助用户理解模型决策,便于调试和提升性能。此外,推理在边缘计算和移动设备上的应用降低了延迟和带宽成本,同时保护了用户隐私。相比于训练大模型的高资源消耗,推理更为节能且成本效益高,尤其在数据挖掘和新知识探索方面展现出创新潜力。推理在实际应用中与训练模型相结合,提供了性能与成本的有效平衡。随着技术进步,推理将在推动人工智能领域发展中发挥更大作用。
|
2月前
|
机器人 Linux Docker
【LLM】基于Stable-Diffusion模型构建可以生成图像的聊天机器人
【4月更文挑战第13天】基于Stable-Diffusion模型构建可以生成图像的聊天机器人
37 0
|
2月前
|
机器学习/深度学习 自然语言处理 测试技术
Stable Diffusion——外挂VAE模型
Stable Diffusion——外挂VAE模型
375 0
|
7月前
|
数据采集 机器学习/深度学习 自然语言处理
本地训练,开箱可用,Bert-VITS2 V2.0.2版本本地基于现有数据集训练(原神刻晴)
按照固有思维方式,深度学习的训练环节应该在云端,毕竟本地硬件条件有限。但事实上,在语音识别和自然语言处理层面,即使相对较少的数据量也可以训练出高性能的模型,对于预算有限的同学们来说,也没必要花冤枉钱上“云端”了,本次我们来演示如何在本地训练Bert-VITS2 V2.0.2模型。
本地训练,开箱可用,Bert-VITS2 V2.0.2版本本地基于现有数据集训练(原神刻晴)
|
2月前
|
机器学习/深度学习 缓存 PyTorch
Yolov5如何训练自定义的数据集,以及使用GPU训练,涵盖报错解决
Yolov5如何训练自定义的数据集,以及使用GPU训练,涵盖报错解决
683 0
|
2月前
|
机器学习/深度学习 数据采集 算法
Midjourney如何训练模型
Midjourney如何训练模型
207 0