共计 8251 个字符，预计需要花费 21 分钟才能阅读完成。
稳固扩散模型因其从文本形容生成高质量、多样化图像的能力而取得了极大的关注。然而这些预训练模型在生成高度定制或个性化主题的图像时可能会有所有余。
这时就须要咱们进行手动的微调。微调能够依据绝对较小的图像集向模型传授新的、独特的主题。咱们明天应用 DreamBooth 在不影响模型原始性能的状况下实现微调过程。
根底概念
1、生成模型和文本到图像的合成
生成模型是一类机器学习模型，旨在生成与给定数据集类似的新数据实例。他们捕获潜在的数据分布，产生新的样本。
文本到图像模型是生成模型的一个子集，因为它们以极高的准确性和保真度将文本形容转换为相应的视觉示意而特地乏味。
Stable Diffusion 是一种文本到图像的模型，它利用 Transformer 架构的一种变体来依据文本输出生成图像。
这些模型的生成过程能够形容如下:
给定一个文本形容 T，模型的指标是生成一个图像 I，使联结概率 P(I,T)最大化。这通常是通过训练模型来最大化条件概率 P(I∣T)来实现的，以确保生成的图像与文本形容统一。
2、生成模型的微调
微调是指在一个新的、通常更小的数据集上调整预训练的模型，以使模型适应特定的需要，而不会失去从原始数据集学习到的泛化性。这种办法在数据稀缺或须要定制的应用程序中至关重要。
在数学上，微调调整模型的参数 θ 以优化新数据集 Dnew 上的损失函数 L，同时避免与原始参数 θ origin 的显著偏差。这能够被表述为一个正则化问题:
3、DreamBooth
DreamBooth 提出了一种新的微调办法，容许生成具备特定主题或对象的图像，同时放弃模型生成不同图像的能力。
传统的微调可能会导致适度拟合或灾难性忘记(遗记原始数据散布)，DreamBooth 则能确保模型保留其个别性能。
该过程包含训练特定于主题的标记以及原始模型参数。这在概念上相似于向模型的词汇表中增加一个代表新主题的新“单词”。训练的指标能够形容为:
其中 φ 示意主题特定参数，1,2λ1，λ2 为正则化参数。
更具体内容请看原论文《DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation》
技术筹备
应用 DreamBooth 微调像 Stable Diffusion 这样的生成模型须要大量的计算，并且须要大量的内存。为了确保训练过程的效率并防止潜在的瓶颈，强烈建议应用高性能 GPU。
如果没有 GPU，咱们能够间接应用 Google Colab，他就能够满足本文的需要。
而后须要装置以下库
Diffusers: 扩散模型库，专门用于微调和利用预训练模型。
Accelerate: 一个用于分布式训练和混合精度的库。
TensorBoard: 用于可视化训练进度和指标。
transformer、FTFY 和 Gradio: 用于模型组件、文本处理和创立用于模型交互的 web ui。
Bitsandbytes: 用于内存高效和疾速训练，特地是用于优化特定 GPU 架构上的模型训练。
!pip install -U -qq git+https://github.com/huggingface/diffusers.git
!pip install -qq accelerate tensorboard transformers ftfy gradio
!pip install -qq "ipywidgets>=7,<8"
!pip install -qq bitsandbytes
!pip install huggingface_hub为了更快、更节俭内存的训练，特地是如果应用特定类型的 gpu (T4、P100、V100、A100)，还能够抉择装置以下组件:
Xformers: 提供高效 transformers 组件的库。
Triton: 用于 gpu 编程。
而后咱们开始进行代码的编写，首先导入库:
importargparse
importitertools
importmath
importos
fromcontextlibimportnullcontext
importrandom
importnumpyasnp
importtorch
importtorch.nn.functionalasF
importtorch.utils.checkpoint
fromtorch.utils.dataimportDataset
importPIL
fromaccelerateimportAccelerator
fromaccelerate.loggingimportget_logger
fromaccelerate.utilsimportset_seed
fromdiffusersimportAutoencoderKL, DDPMScheduler, PNDMScheduler, StableDiffusionPipeline, UNet2DConditionModel
fromdiffusers.optimizationimportget_scheduler
fromdiffusers.pipelines.stable_diffusionimportStableDiffusionSafetyChecker
fromPILimportImage
fromtorchvisionimporttransforms
fromtqdm.autoimporttqdm
fromtransformersimportCLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
importbitsandbytesasbnb
defimage_grid(imgs, rows, cols):
assertlen(imgs) ==rows*cols
w, h=imgs[0].size
grid=Image.new('RGB', size=(cols*w, rows*h))
grid_w, grid_h=grid.size
fori, imginenumerate(imgs):
grid.paste(img, box=(i%cols*w, i//cols*h))
returngrid数据集的筹备
这个过程包含抉择有代表性的图像，对它们进行预处理，并将它们组织成一个结构化的格局以供训练。在本文中，咱们将应用以下 4 个训练图像作为示例。
1、下载和可视化训练图像
上面的 download_image 函数用于从指定的 url 列表中检索图像，而后下载这些图像并将其转换为 RGB 格局以放弃一致性。
咱们将应用他来下载下面的 4 个图片
urls= [
"https://huggingface.co/datasets/Entreprenerdly/finetunestablediffusion/resolve/main/2.jpeg",
"https://huggingface.co/datasets/Entreprenerdly/finetunestablediffusion/resolve/main/3.jpeg",
"https://huggingface.co/datasets/Entreprenerdly/finetunestablediffusion/resolve/main/5.jpeg",
"https://huggingface.co/datasets/Entreprenerdly/finetunestablediffusion/resolve/main/6.jpeg",
## Add additional images here
]
importrequests
importglob
fromioimportBytesIO
defdownload_image(url):
try:
response=requests.get(url)
except:
returnNone
returnImage.open(BytesIO(response.content)).convert("RGB")
images=list(filter(None,[download_image(url) forurlinurls]))
save_path="./my_concept"
ifnotos.path.exists(save_path):
os.mkdir(save_path)
[image.save(f"{save_path}/{i}.jpeg") fori, imageinenumerate(images)]
image_grid(images, 1, len(images))2、创立图像文字对
为了应用 DreamBooth 对生成模型进行微调，须要配置特定的设置以无效地定义新概念。instance_prompt 十分重要，因为它蕴含一个描述性标识符，模型应用它来辨认和生成新概念——在本例中是 cat_toy。
prior_preservation 标记批示模型在训练期间是否应该保留更宽泛的类属性。这有助于提高质量和泛化，但可能会缩短训练工夫。
instance_prompt="<cat-toy> toy" # Descriptive prompt with unique identifier
prior_preservation=False # Flag for enabling class characteristics preservation
prior_preservation_class_prompt="a photo of a cat clay toy" # Prompt for the class of the concept
# Parameters for class image generation and loss weighting
num_class_images=12
sample_batch_size=2
prior_loss_weight=0.5
prior_preservation_class_folder="./class_images"
# Directories for storing class images
class_data_root=prior_preservation_class_folder
class_prompt=prior_preservation_class_prompt3、自定义 DataSet 类
DreamBoothDataset 类扩大了 PyTorch 的 Dataset，用于治理用于训练模型的图像数据。它负责从指定的目录加载图像，利用所需的转换，并应用提供的标记器对提醒进行编码。
另一个类 PromptDataset 被设置为解决类图像提醒的生成。这个简略的数据集构造存储了提醒和要生成的样本数量。
这些类有助于为 DreamBooth 微调过程构建训练数据，确保模型以预期的格局接收数据并进行必要的扩大。
# Initialization of DreamBoothDataset with directory paths and settings
classDreamBoothDataset(Dataset):
...
def__init__(self, instance_data_root, instance_prompt, tokenizer, class_data_root=None, class_prompt=None, size=512, center_crop=False):
...
self.image_transforms=transforms.Compose([...])
def__getitem__(self, index):
...
returnexample
# Class for prompt dataset
classPromptDataset(Dataset):
...优于篇幅太长，咱们就不贴残缺代码了，请在最初的残缺代码中查看
模型加载与配置
咱们将应用 stable-diffusion-2
pretrained_model_name_or_path = "stabilityai/stable-diffusion-2"而后就是与 DreamBooth 微调相干的配置参数包含加载模型体系结构的各种组件: 文本编码器、变分主动编码器 (VAE) 和 U -Net。
每个组件都是从预训练的模型中加载的，以确保兼容性并保留学习到的特色。
# Load models and create wrapper for stable diffusion
text_encoder=CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder")
vae=AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
unet=UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet")
tokenizer=CLIPTokenizer.from_pretrained(
pretrained_model_name_or_path,
subfolder="tokenizer",
)这里的 CLIPTextModel 负责将文本形容编码为嵌入；AutoencoderKL 解决图像的潜在空间示意；UNet2DConditionModel 是模型的次要生成网络；tokenizer 则是解决词元令牌的标记器。
微调
1、设置训练参数
训练参数封装在 argparse 模块的 Namespace 类中。这些参数包含预训练模型的门路、图像的分辨率、是否训练文本编码器、学习率以及与启用之前保留相干的细节。
fromargparseimportNamespace
args=Namespace(
pretrained_model_name_or_path=pretrained_model_name_or_path,
resolution=vae.sample_size,
center_crop=True,
train_text_encoder=False,
instance_data_dir=save_path,
instance_prompt=instance_prompt,
learning_rate=5e-06,
max_train_steps=300,
save_steps=50,
train_batch_size=2, # set to 1 if using prior preservation
gradient_accumulation_steps=2,
max_grad_norm=1.0,
mixed_precision="fp16", # set to "fp16" for mixed-precision training.
gradient_checkpointing=True, # set this to True to lower the memory usage.
use_8bit_adam=True, # use 8bit optimizer from bitsandbytes
seed=3434554,
with_prior_preservation=prior_preservation,
prior_loss_weight=prior_loss_weight,
sample_batch_size=2,
class_data_dir=prior_preservation_class_folder,
class_prompt=prior_preservation_class_prompt,
num_class_images=num_class_images,
lr_scheduler="constant",
lr_warmup_steps=100,
output_dir="dreambooth-concept",
)2、定义训练函数
训练函数负责微调过程。它初始化 Accelerator 以解决分布式训练，为调试复现设置随机种子，配置优化器(应用 8 位精度以进步内存效率)。
该函数还筹备数据集和数据加载器，设置学习率调度器，并定义训练循环，其中包含损失计算和模型更新。
上面是训练函数设置的代码片段:
fromaccelerate.utilsimportset_seed
deftraining_function(text_encoder, vae, unet):
logger=get_logger(__name__)
set_seed(args.seed)
...
ifargs.use_8bit_adam:
optimizer_class=bnb.optim.AdamW8bit
else:
optimizer_class=torch.optim.AdamW
...3、训练
理论的训练是应用 accelerate 库中的 notebook_launcher 启动的，它承受训练函数和先前定义的参数。
importaccelerate
accelerate.notebook_launcher(training_function, args=(text_encoder, vae, unet))
forparaminitertools.chain(unet.parameters(), text_encoder.parameters()):
ifparam.gradisnotNone:
delparam.grad # free some memory
torch.cuda.empty_cache()推理
应用 StableDiffusionPipeline 设置推理管道:
fromdiffusersimportStableDiffusionPipeline, DPMSolverMultistepScheduler
try:
pipe
exceptNameError:
pipe=StableDiffusionPipeline.from_pretrained(
args.output_dir,
scheduler=DPMSolverMultistepScheduler.from_pretrained(args.output_dir, subfolder="scheduler"),
torch_dtype=torch.float16,
).to("cuda")在 Gradio 和 / 或 Google Colab 上运行 Pipeline
importgradioasgr
definference(prompt, num_samples):
images=pipe(prompt, num_images_per_prompt=num_samples, num_inference_steps=25).images
returnimages
withgr.Blocks() asdemo:
prompt=gr.Textbox(label="prompt")
samples=gr.Slider(label="Samples", value=1)
run=gr.Button(value="Run")
gallery=gr.Gallery(show_label=False)
run.click(inference, inputs=[prompt, samples], outputs=gallery)
demo.launch()如果是 Colab，则应用上面代码
prompt = "a <cat-toy> in mad max fury road"
num_samples = 2
all_images = []
images = pipe(prompt, num_images_per_prompt=num_samples, num_inference_steps=25, guidance_scale=9).images
all_images.extend(images)
grid = image_grid(all_images, num_rows, num_samples)总结
dreambooth 通过向模型注入自定义的主题来 fine-tune diffusion model，它简化了咱们微调自定义模型的老本，而 Google Colab 的收费 GPU 能够让咱们进行更多的测试，以下是本文的残缺代码，能够间接在线测试：
https://avoid.overfit.cn/post/1b02c6c3d93749558941777786b64513
作者：Cris Velasquez