Безусловная генерация изображений

Модели безусловной генерации изображений не используют в качестве условий текст или изображения во время обучения. Они генерируют только изображения, распределение которых сходно с распределением обучающих данных.

В этом руководстве мы рассмотрим скрипт обучения train_unconditional.py, чтобы вы могли ознакомиться с ним и понять, как адаптировать его под свои задачи.

Перед запуском скрипта убедитесь, что вы установили библиотеку из исходного кода:

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

Затем перейдите в папку с примером, содержащую скрипт обучения, и установите необходимые зависимости:

cd examples/unconditional_image_generation
pip install -r requirements.txt

🤗 Accelerate — это библиотека, которая помогает обучать модели на нескольких GPU/TPU или с использованием смешанной точности. Она автоматически настраивает среду обучения в зависимости от вашего оборудования и окружения. Подробнее можно узнать в Кратком обзоре 🤗 Accelerate.

Инициализируйте среду 🤗 Accelerate:

accelerate config

Чтобы настроить среду по умолчанию 🤗 Accelerate без выбора каких-либо конфигураций:

accelerate config default

Или, если ваша среда не поддерживает интерактивную оболочку, например, ноутбук, вы можете использовать:

from accelerate.utils import write_basic_config

write_basic_config()

Наконец, если вы хотите обучить модель на собственном наборе данных, ознакомьтесь с руководством Create a dataset for training, чтобы узнать, как создать набор данных, совместимый со скриптом обучения.

Параметры скрипта

В следующих разделах выделены части скрипта обучения, важные для понимания того, как его модифицировать, но не все аспекты скрипта описаны подробно. Если вы хотите узнать больше, не стесняйтесь прочитать скрипт и сообщить нам, если у вас возникнут вопросы или замечания.

Скрипт обучения предоставляет множество параметров, которые помогают настроить процесс обучения. Все параметры и их описания находятся в функции parse_args(). Для каждого параметра заданы значения по умолчанию, например, размер обучающей выборки и скорость обучения, но при необходимости вы можете указать свои значения в команде запуска обучения.

Например, чтобы ускорить обучение с использованием смешанной точности в формате bf16, добавьте параметр --mixed_precision к команде обучения:

accelerate launch train_unconditional.py \
  --mixed_precision="bf16"

К числу основных и важных параметров, которые необходимо указать, относятся:

• --dataset_name: название набора данных на Hub или локальный путь к набору данных для обучения
• --output_dir: место для сохранения обученной модели
• --push_to_hub: нужно ли отправить обученную модель в Hub
• --checkpointing_steps: частота сохранения контрольных точек в процессе обучения модели; это полезно, если обучение прерывается, вы можете продолжить обучение с этой контрольной точки, добавив --resume_from_checkpoint к команде обучения

Принесите свой набор данных, а остальное скрипт обучения сделает сам!

Скрипт обучения

Код для предобработки набора данных и цикла обучения находится в функции main(). Если вам нужно адаптировать скрипт обучения, именно здесь следует внести изменения.

Скрипт train_unconditional инициализирует UNet2DModel, если вы не предоставили конфигурацию модели. При необходимости вы можете настроить UNet здесь:

model = UNet2DModel(
    sample_size=args.resolution,
    in_channels=3,
    out_channels=3,
    layers_per_block=2,
    block_out_channels=(128, 128, 256, 256, 512, 512),
    down_block_types=(
        "DownBlock2D",
        "DownBlock2D",
        "DownBlock2D",
        "DownBlock2D",
        "AttnDownBlock2D",
        "DownBlock2D",
    ),
    up_block_types=(
        "UpBlock2D",
        "AttnUpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
        "UpBlock2D",
    ),
)

Далее скрипт инициализирует scheduler и optimizer:

# Initialize the scheduler
accepts_prediction_type = "prediction_type" in set(inspect.signature(DDPMScheduler.__init__).parameters.keys())
if accepts_prediction_type:
    noise_scheduler = DDPMScheduler(
        num_train_timesteps=args.ddpm_num_steps,
        beta_schedule=args.ddpm_beta_schedule,
        prediction_type=args.prediction_type,
    )
else:
    noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule)

# Initialize the optimizer
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Затем он загружает набор данных, и вы можете указать, как его предобработать:

dataset = load_dataset("imagefolder", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split="train")

augmentations = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
        transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)

Наконец, training loop выполняет все остальные задачи: добавление шума к изображениям, предсказание остатка шума, вычисление функции потерь, сохранение контрольных точек на указанных шагах, а также сохранение модели и её публикацию в Hub. Если вы хотите узнать больше о том, как работает цикл обучения, обратитесь к учебнику Understanding pipelines, models and schedulers, в котором подробно разбирается базовый паттерн процесса устранения шума.

Запустите скрипт

Когда вы внесёте все изменения или решите оставить конфигурацию по умолчанию, можно запускать скрипт обучения! 🚀

Полный цикл обучения занимает 2 часа на 4xV100 GPU.

accelerate launch train_unconditional.py \
  --dataset_name="huggan/flowers-102-categories" \
  --output_dir="ddpm-ema-flowers-64" \
  --mixed_precision="fp16" \
  --push_to_hub

accelerate launch --multi_gpu train_unconditional.py \
  --dataset_name="huggan/flowers-102-categories" \
  --output_dir="ddpm-ema-flowers-64" \
  --mixed_precision="fp16" \
  --push_to_hub

Скрипт обучения создаёт и сохраняет файл контрольной точки в вашем репозитории. Теперь вы можете загрузить и использовать обученную модель для инференса:

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("anton-l/ddpm-butterflies-128").to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline().images[0]