Diffusers documentation

LoRA

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
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

O 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