LoRA

Это экспериментальная функция, и API может измениться в будущем.

LoRA (Low-Rank Adaptation of Large Language Models) — популярный и лёгкий метод обучения, который значительно сокращает количество обучаемых параметров. Он работает за счёт вставки в модель меньшего числа новых весов, и обучаются только они. Это делает обучение с LoRA намного быстрее, экономнее по памяти и приводит к получению меньших весов модели (несколько сотен МБ), которые проще хранить и передавать. LoRA также можно комбинировать с другими методами обучения, такими как DreamBooth, чтобы ускорить процесс обучения.

LoRA очень универсален и поддерживается для DreamBooth, Kandinsky 2.2, Stable Diffusion XL и text-to-image.

This guide will explore the train_text_to_image_lora.py script to help you become more familiar with it, and how you can adapt it for your own use-case.

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

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

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

cd examples/text_to_image
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(). Для большинства параметров заданы значения по умолчанию, которые хорошо работают, но при необходимости вы можете указать свои значения в команде запуска обучения.

Например, чтобы увеличить количество эпох обучения:

accelerate launch train_text_to_image_lora.py \
  --num_train_epochs=150 \

Многие базовые и важные параметры описаны в руководстве по обучению Text-to-image, поэтому в данном руководстве рассматриваются только параметры, относящиеся к LoRA:

--rank: внутреннее измерение матриц низкого ранга, подлежащих обучению; чем выше ранг, тем больше обучаемых параметров
--learning_rate: скорость обучения по умолчанию составляет 1e-4, но при использовании LoRA можно применять более высокую скорость обучения

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

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

Как и в случае с параметрами скрипта, пошаговое описание скрипта обучения приведено в руководстве по обучению Text-to-image. В данном руководстве рассматриваются только те части скрипта, которые относятся к LoRA.

UNet

text encoder

Оптимизатор инициализируется с lora_layers, поскольку именно эти веса подлежат оптимизации:

optimizer = optimizer_cls(
    lora_layers,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Aside from setting up the LoRA layers, the training script is more or less the same as train_text_to_image.py!

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

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

Обучимся на датасете Naruto BLIP captions, чтобы генерировать собственных персонажей Наруто. Установите переменные окружения MODEL_NAME и DATASET_NAME для указания модели и датасета соответственно. Также укажите, куда сохранять модель, в OUTPUT_DIR, и имя модели для сохранения на Hub в HUB_MODEL_ID. Скрипт создаёт и сохраняет в ваш репозиторий следующие файлы:

сохранённые контрольные точки модели
pytorch_lora_weights.safetensors (обученные веса LoRA)

Если вы обучаете модель на нескольких GPU, добавьте параметр --multi_gpu к команде accelerate launch.

Полный цикл обучения занимает около 5 часов на GPU 2080 Ti с 11 ГБ видеопамяти.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/naruto"
export HUB_MODEL_ID="naruto-lora"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME \
  --dataloader_num_workers=8 \
  --resolution=512 \
  --center_crop \
  --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-04 \
  --max_grad_norm=1 \
  --lr_scheduler="cosine" \
  --lr_warmup_steps=0 \
  --output_dir=${OUTPUT_DIR} \
  --push_to_hub \
  --hub_model_id=${HUB_MODEL_ID} \
  --report_to=wandb \
  --checkpointing_steps=500 \
  --validation_prompt="A naruto with blue eyes." \
  --seed=1337

После завершения обучения вы можете использовать свою модель для инференса:

from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.load_lora_weights("path/to/lora/model", weight_name="pytorch_lora_weights.safetensors")
image = pipeline("A naruto with blue eyes").images[0]

Дальнейшие шаги

Поздравляем с обучением новой модели с помощью LoRA! Чтобы узнать, как использовать новую модель, вам могут помочь следующие руководства:

Узнайте, как загружать различные форматы LoRA, обученные с помощью комьюнити-тренажеров, таких как Kohya и TheLastBen.
Узнайте, как использовать и комбинировать несколько LoRA с PEFT для инференса.

Update on GitHub

Diffusers

LoRA

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

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

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

Дальнейшие шаги