# Обучение диффузионной модели Безусловная генерация изображений — популярное применение диффузионных моделей, которое создаёт изображения, похожие на те, что использовались в датасете для обучения. Как правило, наилучшие результаты достигаются при дообучении предобученной модели на конкретном датасете. Многие такие чекпоинты можно найти на [Hub](https://huggingface.co/search/full-text?q=unconditional-image-generation&type=model), но если вы не найдёте подходящий, всегда можно обучить собственную модель! В этом учебнике вы узнаете, как обучить [UNet2DModel](/docs/diffusers/v0.43.0/ru/api/models/unet2d#diffusers.UNet2DModel) с нуля на подмножестве набора данных [Smithsonian Butterflies](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset), чтобы генерировать собственных 🦋 бабочек 🦋. > [!TIP] > 💡 Это учебное пособие по обучению основано на ноутбуке [Обучение с 🧨 Diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb). Дополнительные детали и контекст о диффузионных моделях, например, о том, как они работают, можно найти в ноутбуке! Прежде чем начать, убедитесь, что у вас установлены 🤗 Datasets для загрузки и предобработки наборов данных изображений, а также 🤗 Accelerate для упрощения обучения на любом количестве GPU. Следующая команда также установит [TensorBoard](https://www.tensorflow.org/tensorboard) для визуализации метрик обучения (вы также можете использовать [Weights & Biases](https://docs.wandb.ai/) для отслеживания обучения). ```py # uncomment to install the necessary libraries in Colab #!pip install diffusers[training] ``` Мы рекомендуем делиться вашей моделью с сообществом. Для этого необходимо войти в свой аккаунт Hugging Face (создайте его [здесь](https://hf.co/join), если у вас его ещё нет!). Вы можете войти из notebook и ввести свой токен, когда система попросит. Убедитесь, что у вашего токена есть роль write. ```py >>> from huggingface_hub import notebook_login >>> notebook_login() ``` Или войдите в систему из терминала: ```bash hf auth login ``` Поскольку контрольные точки моделей довольно большие, установите [Git-LFS](https://git-lfs.com/) для версионирования этих крупных файлов: ```bash !sudo apt -qq install git-lfs !git config --global credential.helper store ``` ## Конфигурация обучения Для удобства создайте класс `TrainingConfig`, содержащий гиперпараметры обучения (при необходимости их можно изменить): ```py >>> 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 = "/" # 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](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) с помощью библиотеки 🤗 Datasets: ```py >>> from datasets import load_dataset >>> config.dataset_name = "huggan/smithsonian_butterflies_subset" >>> dataset = load_dataset(config.dataset_name, split="train") ``` > [!TIP] > 💡 Дополнительные наборы данных можно найти в рамках [HugGan Community Event](https://huggingface.co/huggan) или использовать собственный набор данных, создав локальный [`ImageFolder`](https://huggingface.co/docs/datasets/image_dataset#imagefolder). Установите `config.dataset_name` в идентификатор репозитория набора данных, если он взят из HugGan Community Event, или `imagefolder`, если вы используете собственные изображения. 🤗 Datasets использует функцию [Image](https://huggingface.co/docs/datasets/v5.2.0/en/package_reference/main_classes#datasets.Image) для автоматического декодирования данных изображения и загрузки его в виде [`PIL.Image`](https://pillow.readthedocs.io/en/stable/reference/Image.html), что позволяет визуализировать его: ```py >>> 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], который ожидает модель. ```py >>> 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](https://huggingface.co/docs/datasets/v5.2.0/en/package_reference/main_classes#datasets.Dataset.set_transform) библиотеки 🤗 Datasets, чтобы применять функцию `preprocess` на лету во время обучения: ```py >>> def transform(examples): ... images = [preprocess(image.convert("RGB")) for image in examples["image"]] ... return {"images": images} >>> dataset.set_transform(transform) ``` Вы можете снова визуализировать изображения, чтобы убедиться, что они были изменены в размере. Теперь вы готовы обернуть набор данных в [DataLoader](https://pytorch.org/docs/stable/data#torch.utils.data.DataLoader) для обучения! ```py >>> import torch >>> train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True) ``` ## Создание UNet2DModel Предобученные модели в 🧨 Diffusers легко создаются из их класса модели с нужными вам параметрами. Например, чтобы создать [UNet2DModel](/docs/diffusers/v0.43.0/ru/api/models/unet2d#diffusers.UNet2DModel): ```py >>> 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", ... ), ... ) ``` Часто бывает хорошей идеей быстро проверить, что форма образцов совпадает с формой вывода модели: ```py >>> 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](/docs/diffusers/v0.43.0/ru/api/schedulers/ddpm#diffusers.DDPMScheduler) и воспользуемся методом `add_noise`, чтобы добавить случайный шум к изображению `sample_image`, полученному ранее: ```py >>> 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]) ``` Целью обучения модели является предсказание шума, добавленного к изображению. Потери на этом шаге можно вычислить по формуле: ```py >>> import torch.nn.functional as F >>> noise_pred = model(noisy_image, timesteps).sample >>> loss = F.mse_loss(noise_pred, noise) ``` ## Обучение модели К этому моменту у вас есть все необходимые компоненты для начала обучения модели, и осталось лишь собрать всё воедино. Сначала вам понадобится оптимизатор и планировщик скорости обучения: ```py >>> 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](/docs/diffusers/v0.43.0/ru/api/pipelines/ddpm#diffusers.DDPMPipeline), чтобы сгенерировать пакет примеров изображений и сохранить его в виде сетки: ```py >>> 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. > [!TIP] > 💡 Цикл обучения ниже может показаться пугающим и длинным, но позже это окупится, когда вы запустите обучение всего одной строкой кода! Если вы не можете дождаться и хотите начать генерировать изображения, смело копируйте и запускайте код ниже. Вы всегда можете вернуться и более внимательно изучить цикл обучения позже, например, пока ждёте завершения обучения вашей модели. 🤗 ```py >>> 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](https://huggingface.co/docs/accelerate/v1.15.0/en/package_reference/launchers#accelerate.notebook_launcher) из 🤗 Accelerate. Передайте в функцию цикл обучения, все аргументы обучения и количество процессов (это значение можно изменить на количество доступных вам GPU), которые будут использоваться для обучения: ```py >>> from accelerate import notebook_launcher >>> args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler) >>> notebook_launcher(train_loop, args, num_processes=1) ``` Когда обучение завершено, посмотрите на финальные изображения 🦋, сгенерированные вашей диффузионной моделью! 🦋 ```py >>> import glob >>> sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png")) >>> Image.open(sample_images[-1]) ``` ## Дальнейшие шаги Безусловная генерация изображений — один из примеров задач, для которых можно обучать модель. Вы можете ознакомиться с другими задачами и методами обучения на странице [🧨 Diffusers Training Examples](../training/overview). Вот несколько примеров того, что вы можете узнать: * [Textual Inversion](../training/text_inversion) — алгоритм, который обучает модель конкретному визуальному концепту и интегрирует его в генерируемое изображение. * [DreamBooth](../training/dreambooth) — техника генерации персонализированных изображений объекта на основе нескольких входных изображений этого объекта. * [Руководство](../training/text2image) по дообучению модели Stable Diffusion на собственном наборе данных. * [Руководство](../training/lora) по использованию LoRA — эффективного с точки зрения памяти метода для более быстрого дообучения очень больших моделей.