Diffusers documentation

LoRA

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

LoRA

Isso é experimental e a API pode mudar no futuro.

LoRA (Low-Rank Adaptation of Large Language Models) é uma técnica de treinamento popular e leve que reduz significativamente o número de parâmetros treináveis. Ela funciona inserindo um número menor de novos pesos no modelo, e apenas esses são treinados. Isso torna o treinamento com LoRA muito mais rápido, eficiente em memória e produz pesos de modelo menores (alguns centenas de MBs), que são mais fáceis de armazenar e compartilhar. O LoRA também pode ser combinado com outras técnicas de treinamento, como DreamBooth, para acelerar o treinamento.

LoRA é muito versátil e é suportado para DreamBooth, Kandinsky 2.2, Stable Diffusion XL e text-to-image.

Este guia explorará o script train_text_to_image_lora.py para ajudá-lo a se familiarizar melhor com ele e como você pode adaptá-lo para o seu próprio caso de uso.

Antes de executar o script, certifique-se de instalar a biblioteca a partir do código-fonte:

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

Navegue até a pasta de exemplo com o script de treinamento e instale as dependências necessárias para o script que você está usando:

cd examples/text_to_image
pip install -r requirements.txt

🤗 Accelerate é uma biblioteca para ajudar você a treinar em múltiplos GPUs/TPUs ou com precisão mista. Ela configura automaticamente seu ambiente de treinamento com base no seu hardware e ambiente. Dê uma olhada no Quick tour do 🤗 Accelerate para saber mais.

Inicialize um ambiente do 🤗 Accelerate:

accelerate config

Para configurar um ambiente padrão do 🤗 Accelerate sem escolher nenhuma configuração:

accelerate config default

Ou, se o seu ambiente não suporta um shell interativo, como um notebook, você pode usar:

from accelerate.utils import write_basic_config

write_basic_config()

Por fim, se você deseja treinar um modelo em seu próprio conjunto de dados, consulte o guia Criar um conjunto de dados para treinamento para aprender a criar um conjunto de dados que funcione com o script de treinamento.

As seções a seguir destacam partes do script de treinamento que são importantes para entender como modificá-lo, mas não cobrem todos os aspectos do script em detalhes. Se você tiver interesse em aprender mais, sinta-se à vontade para ler o script e nos informar caso tenha alguma dúvida ou preocupação.

Parâmetros do script

O script de treinamento possui muitos parâmetros para ajudar você a personalizar sua execução de treinamento. Todos os parâmetros e suas descrições estão na função parse_args(). Valores padrão são fornecidos para a maioria dos parâmetros e funcionam muito bem, mas você também pode definir seus próprios valores no comando de treinamento, se desejar.

Por exemplo, para aumentar o número de épocas de treinamento:

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

Muitos dos parâmetros básicos e importantes são descritos no guia de treinamento Texto-para-imagem, portanto, este guia se concentra apenas nos parâmetros relevantes para o LoRA:

  • --rank: a dimensão interna das matrizes de baixo rank a serem treinadas; um rank mais alto significa mais parâmetros treináveis
  • --learning_rate: a taxa de aprendizado padrão é 1e-4, mas com LoRA, você pode usar uma taxa de aprendizado mais alta

Script de treinamento

O código de pré-processamento do conjunto de dados e o loop de treinamento estão na função main(), e, se você precisar adaptar o script de treinamento, é aqui que fará suas alterações.

Assim como com os parâmetros do script, um passo a passo do script de treinamento é fornecido no guia de treinamento Text-to-image. Em vez disso, este guia examina as partes do script relevantes para o LoRA.

UNet
codificador de texto

Diffusers usa LoraConfig da biblioteca PEFT para configurar os parâmetros do adaptador LoRA, como rank, alpha e em quais módulos inserir os pesos do LoRA. O adaptador é adicionado ao UNet, e apenas as camadas LoRA são filtradas para otimização em lora_layers.

unet_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank,
    init_lora_weights="gaussian",
    target_modules=["to_k", "to_q", "to_v", "to_out.0"],
)

unet.add_adapter(unet_lora_config)
lora_layers = filter(lambda p: p.requires_grad, unet.parameters())

O otimizador é inicializado com as lora_layers, pois são os únicos pesos que serão otimizados:

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,
)

Além de configurar as camadas LoRA, o script de treinamento é mais ou menos o mesmo que train_text_to_image.py!

Inicie o script

Depois de fazer todas as alterações ou se você estiver satisfeito com a configuração padrão, está pronto para executar o script de treinamento! 🚀

Vamos treinar no conjunto de dados Naruto BLIP captions para gerar seus próprios personagens de Naruto. Defina as variáveis de ambiente MODEL_NAME e DATASET_NAME para o modelo e o conjunto de dados, respectivamente. Você também deve especificar onde salvar o modelo em OUTPUT_DIR e o nome do modelo a ser salvo no Hub com HUB_MODEL_ID. O script cria e salva os seguintes arquivos no seu repositório:

  • pontos de verificação de modelos salvos
  • pytorch_lora_weights.safetensors (os pesos do LoRA treinados)

Se você estiver treinando em mais de uma GPU, adicione o parâmetro --multi_gpu ao comando accelerate launch.

Um treinamento completo leva cerca de 5 horas em uma GPU 2080 Ti com 11 GB de VRAM.

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

Depois que o treinamento tiver sido concluído, você pode usar seu modelo para inferência:

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]

Próximas etapas

Parabéns por treinar um novo modelo com LoRA! Para saber mais sobre como usar seu novo modelo, os seguintes guias podem ser úteis:

Update on GitHub