Diffusers documentation
Обучение диффузионной модели
Обучение диффузионной модели
Безусловная генерация изображений — популярное применение диффузионных моделей, которое создаёт изображения, похожие на те, что использовались в датасете для обучения. Как правило, наилучшие результаты достигаются при дообучении предобученной модели на конкретном датасете. Многие такие чекпоинты можно найти на Hub, но если вы не найдёте подходящий, всегда можно обучить собственную модель!
В этом учебнике вы узнаете, как обучить UNet2DModel с нуля на подмножестве набора данных Smithsonian Butterflies, чтобы генерировать собственных 🦋 бабочек 🦋.
💡 Это учебное пособие по обучению основано на ноутбуке Обучение с 🧨 Diffusers. Дополнительные детали и контекст о диффузионных моделях, например, о том, как они работают, можно найти в ноутбуке!
Прежде чем начать, убедитесь, что у вас установлены 🤗 Datasets для загрузки и предобработки наборов данных изображений, а также 🤗 Accelerate для упрощения обучения на любом количестве GPU. Следующая команда также установит TensorBoard для визуализации метрик обучения (вы также можете использовать Weights & Biases для отслеживания обучения).
# uncomment to install the necessary libraries in Colab
#!pip install diffusers[training]Мы рекомендуем делиться вашей моделью с сообществом. Для этого необходимо войти в свой аккаунт Hugging Face (создайте его здесь, если у вас его ещё нет!). Вы можете войти из notebook и ввести свой токен, когда система попросит. Убедитесь, что у вашего токена есть роль write.
>>> from huggingface_hub import notebook_login
>>> notebook_login()Или войдите в систему из терминала:
hf auth login
Поскольку контрольные точки моделей довольно большие, установите Git-LFS для версионирования этих крупных файлов:
!sudo apt -qq install git-lfs
!git config --global credential.helper storeКонфигурация обучения
Для удобства создайте класс TrainingConfig, содержащий гиперпараметры обучения (при необходимости их можно изменить):
>>> from dataclasses import dataclass
>>> @dataclass
... class TrainingConfig:
... image_size = 128 # the generated image resolution
... train_batch_size = 16
... eval_batch_size = 16 # how many images to sample during evaluation
... num_epochs = 50
... gradient_accumulation_steps = 1
... learning_rate = 1e-4
... lr_warmup_steps = 500
... save_image_epochs = 10
... save_model_epochs = 30
... mixed_precision = "fp16" # `no` for float32, `fp16` for automatic mixed precision
... output_dir = "ddpm-butterflies-128" # the model name locally and on the HF Hub
... push_to_hub = True # whether to upload the saved model to the HF Hub
... hub_model_id = "<your-username>/<my-awesome-model>" # the name of the repository to create on the HF Hub
... hub_private_repo = None
... overwrite_output_dir = True # overwrite the old model when re-running the notebook
... seed = 0
>>> config = TrainingConfig()Загрузите набор данных
Вы можете легко загрузить набор данных Smithsonian Butterflies с помощью библиотеки 🤗 Datasets:
>>> from datasets import load_dataset
>>> config.dataset_name = "huggan/smithsonian_butterflies_subset"
>>> dataset = load_dataset(config.dataset_name, split="train")💡 Дополнительные наборы данных можно найти в рамках HugGan Community Event или использовать собственный набор данных, создав локальный
ImageFolder. Установитеconfig.dataset_nameв идентификатор репозитория набора данных, если он взят из HugGan Community Event, илиimagefolder, если вы используете собственные изображения.
🤗 Datasets использует функцию Image для автоматического декодирования данных изображения и загрузки его в виде PIL.Image, что позволяет визуализировать его:
>>> import matplotlib.pyplot as plt
>>> fig, axs = plt.subplots(1, 4, figsize=(16, 4))
>>> for i, image in enumerate(dataset[:4]["image"]):
... axs[i].imshow(image)
... axs[i].set_axis_off()
>>> fig.show()
Однако все изображения имеют разные размеры, поэтому сначала необходимо выполнить их предобработку:
Resizeизменяет размер изображения до значения, заданного вconfig.image_size.RandomHorizontalFlipдополняет набор данных, случайно отражая изображения по горизонтали.Normalizeважно для масштабирования значений пикселей в диапазон [-1, 1], который ожидает модель.
>>> from torchvision import transforms
>>> preprocess = transforms.Compose(
... [
... transforms.Resize((config.image_size, config.image_size)),
... transforms.RandomHorizontalFlip(),
... transforms.ToTensor(),
... transforms.Normalize([0.5], [0.5]),
... ]
... )Используйте метод set_transform библиотеки 🤗 Datasets, чтобы применять функцию preprocess на лету во время обучения:
>>> def transform(examples):
... images = [preprocess(image.convert("RGB")) for image in examples["image"]]
... return {"images": images}
>>> dataset.set_transform(transform)Вы можете снова визуализировать изображения, чтобы убедиться, что они были изменены в размере. Теперь вы готовы обернуть набор данных в DataLoader для обучения!
>>> import torch
>>> train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)Создание UNet2DModel
Предобученные модели в 🧨 Diffusers легко создаются из их класса модели с нужными вам параметрами. Например, чтобы создать UNet2DModel:
>>> from diffusers import UNet2DModel
>>> model = UNet2DModel(
... sample_size=config.image_size, # the target image resolution
... in_channels=3, # the number of input channels, 3 for RGB images
... out_channels=3, # the number of output channels
... layers_per_block=2, # how many ResNet layers to use per UNet block
... block_out_channels=(128, 128, 256, 256, 512, 512), # the number of output channels for each UNet block
... down_block_types=(
... "DownBlock2D", # a regular ResNet downsampling block
... "DownBlock2D",
... "DownBlock2D",
... "DownBlock2D",
... "AttnDownBlock2D", # a ResNet downsampling block with spatial self-attention
... "DownBlock2D",
... ),
... up_block_types=(
... "UpBlock2D", # a regular ResNet upsampling block
... "AttnUpBlock2D", # a ResNet upsampling block with spatial self-attention
... "UpBlock2D",
... "UpBlock2D",
... "UpBlock2D",
... "UpBlock2D",
... ),
... )Часто бывает хорошей идеей быстро проверить, что форма образцов совпадает с формой вывода модели:
>>> sample_image = dataset[0]["images"].unsqueeze(0)
>>> print("Input shape:", sample_image.shape)
Input shape: torch.Size([1, 3, 128, 128])
>>> print("Output shape:", model(sample_image, timestep=0).sample.shape)
Output shape: torch.Size([1, 3, 128, 128])Отлично! Далее вам понадобится планировщик, чтобы добавить шум к изображению.
Создание планировщика
Поведение планировщика зависит от того, используется ли модель для обучения или инференса. Во время инференса планировщик генерирует изображение из шума. Во время обучения планировщик принимает выход модели — или образец — из конкретной точки процесса диффузии и добавляет шум к изображению в соответствии с расписанием шума и правилом обновления.
Рассмотрим DDPMScheduler и воспользуемся методом add_noise, чтобы добавить случайный шум к изображению sample_image, полученному ранее:
>>> import torch
>>> from PIL import Image
>>> from diffusers import DDPMScheduler
>>> noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
>>> noise = torch.randn(sample_image.shape)
>>> timesteps = torch.LongTensor([50])
>>> noisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)
>>> Image.fromarray(((noisy_image.permute(0, 2, 3, 1) + 1.0) * 127.5).type(torch.uint8).numpy()[0])
Целью обучения модели является предсказание шума, добавленного к изображению. Потери на этом шаге можно вычислить по формуле:
>>> import torch.nn.functional as F
>>> noise_pred = model(noisy_image, timesteps).sample
>>> loss = F.mse_loss(noise_pred, noise)Обучение модели
К этому моменту у вас есть все необходимые компоненты для начала обучения модели, и осталось лишь собрать всё воедино.
Сначала вам понадобится оптимизатор и планировщик скорости обучения:
>>> from diffusers.optimization import get_cosine_schedule_with_warmup
>>> optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
>>> lr_scheduler = get_cosine_schedule_with_warmup(
... optimizer=optimizer,
... num_warmup_steps=config.lr_warmup_steps,
... num_training_steps=(len(train_dataloader) * config.num_epochs),
... )Затем вам понадобится способ оценить модель. Для оценки можно использовать DDPMPipeline, чтобы сгенерировать пакет примеров изображений и сохранить его в виде сетки:
>>> from diffusers import DDPMPipeline
>>> from diffusers.utils import make_image_grid
>>> import os
>>> def evaluate(config, epoch, pipeline):
... # Sample some images from random noise (this is the backward diffusion process).
... # The default pipeline output type is `List[PIL.Image]`
... images = pipeline(
... batch_size=config.eval_batch_size,
... generator=torch.Generator(device='cpu').manual_seed(config.seed), # Use a separate torch generator to avoid rewinding the random state of the main training loop
... ).images
... # Make a grid out of the images
... image_grid = make_image_grid(images, rows=4, cols=4)
... # Save the images
... test_dir = os.path.join(config.output_dir, "samples")
... os.makedirs(test_dir, exist_ok=True)
... image_grid.save(f"{test_dir}/{epoch:04d}.png")Теперь вы можете объединить все эти компоненты в цикл обучения с помощью 🤗 Accelerate для простого логирования в TensorBoard, накопления градиентов и обучения с смешанной точностью. Чтобы загрузить модель в Hub, напишите функцию для получения имени вашего репозитория и информации о нём, а затем отправьте её в Hub.
💡 Цикл обучения ниже может показаться пугающим и длинным, но позже это окупится, когда вы запустите обучение всего одной строкой кода! Если вы не можете дождаться и хотите начать генерировать изображения, смело копируйте и запускайте код ниже. Вы всегда можете вернуться и более внимательно изучить цикл обучения позже, например, пока ждёте завершения обучения вашей модели. 🤗
>>> from accelerate import Accelerator
>>> from huggingface_hub import create_repo, upload_folder
>>> from tqdm.auto import tqdm
>>> from pathlib import Path
>>> import os
>>> def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler):
... # Initialize accelerator and tensorboard logging
... accelerator = Accelerator(
... mixed_precision=config.mixed_precision,
... gradient_accumulation_steps=config.gradient_accumulation_steps,
... log_with="tensorboard",
... project_dir=os.path.join(config.output_dir, "logs"),
... )
... if accelerator.is_main_process:
... if config.output_dir is not None:
... os.makedirs(config.output_dir, exist_ok=True)
... if config.push_to_hub:
... repo_id = create_repo(
... repo_id=config.hub_model_id or Path(config.output_dir).name, exist_ok=True
... ).repo_id
... accelerator.init_trackers("train_example")
... # Prepare everything
... # There is no specific order to remember, you just need to unpack the
... # objects in the same order you gave them to the prepare method.
... model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
... model, optimizer, train_dataloader, lr_scheduler
... )
... global_step = 0
... # Now you train the model
... for epoch in range(config.num_epochs):
... progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
... progress_bar.set_description(f"Epoch {epoch}")
... for step, batch in enumerate(train_dataloader):
... clean_images = batch["images"]
... # Sample noise to add to the images
... noise = torch.randn(clean_images.shape, device=clean_images.device)
... bs = clean_images.shape[0]
... # Sample a random timestep for each image
... timesteps = torch.randint(
... 0, noise_scheduler.config.num_train_timesteps, (bs,), device=clean_images.device,
... dtype=torch.int64
... )
... # Add noise to the clean images according to the noise magnitude at each timestep
... # (this is the forward diffusion process)
... noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
... with accelerator.accumulate(model):
... # Predict the noise residual
... noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
... loss = F.mse_loss(noise_pred, noise)
... accelerator.backward(loss)
... if accelerator.sync_gradients:
... accelerator.clip_grad_norm_(model.parameters(), 1.0)
... optimizer.step()
... lr_scheduler.step()
... optimizer.zero_grad()
... progress_bar.update(1)
... logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
... progress_bar.set_postfix(**logs)
... accelerator.log(logs, step=global_step)
... global_step += 1
... # After each epoch you optionally sample some demo images with evaluate() and save the model
... if accelerator.is_main_process:
... pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)
... if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:
... evaluate(config, epoch, pipeline)
... if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
... if config.push_to_hub:
... upload_folder(
... repo_id=repo_id,
... folder_path=config.output_dir,
... commit_message=f"Epoch {epoch}",
... ignore_patterns=["step_*", "epoch_*"],
... )
... else:
... pipeline.save_pretrained(config.output_dir)Фух, это было немало кода! Но теперь вы наконец готовы запустить обучение с помощью функции notebook_launcher из 🤗 Accelerate. Передайте в функцию цикл обучения, все аргументы обучения и количество процессов (это значение можно изменить на количество доступных вам GPU), которые будут использоваться для обучения:
>>> from accelerate import notebook_launcher
>>> args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)
>>> notebook_launcher(train_loop, args, num_processes=1)Когда обучение завершено, посмотрите на финальные изображения 🦋, сгенерированные вашей диффузионной моделью! 🦋
>>> import glob
>>> sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png"))
>>> Image.open(sample_images[-1])
Дальнейшие шаги
Безусловная генерация изображений — один из примеров задач, для которых можно обучать модель. Вы можете ознакомиться с другими задачами и методами обучения на странице 🧨 Diffusers Training Examples. Вот несколько примеров того, что вы можете узнать:
- Textual Inversion — алгоритм, который обучает модель конкретному визуальному концепту и интегрирует его в генерируемое изображение.
- DreamBooth — техника генерации персонализированных изображений объекта на основе нескольких входных изображений этого объекта.
- Руководство по дообучению модели Stable Diffusion на собственном наборе данных.
- Руководство по использованию LoRA — эффективного с точки зрения памяти метода для более быстрого дообучения очень больших моделей.