Geração de imagem incondicional

Modelos de geração de imagem incondicional não são condicionados em texto ou imagens durante o treinamento. Eles apenas geram imagens que se assemelham à distribuição dos dados de treinamento.

Este guia explorará o script de treinamento train_unconditional.py para ajudá-lo a se familiarizar 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 .

Em seguida, navegue até a pasta de exemplo que contém o script de treinamento e instale as dependências necessárias:

cd examples/unconditional_image_generation
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.

Parâmetros do script

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.

O script de treinamento fornece muitos parâmetros para ajudar você a personalizar sua execução de treinamento. Todos os parâmetros e suas descrições estão encontrados na função parse_args(). Ele fornece valores padrão para cada parâmetro, como o tamanho do lote de treinamento e a taxa de aprendizado, mas você também pode definir seus próprios valores no comando de treinamento, se desejar.

Por exemplo, para acelerar o treinamento com precisão mista usando o formato bf16, adicione o parâmetro --mixed_precision ao comando de treinamento:

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

Alguns parâmetros básicos e importantes a serem especificados incluem:

• --dataset_name: o nome do conjunto de dados no Hub ou um caminho local para o conjunto de dados em que o treinamento será realizado
• --output_dir: onde salvar o modelo treinado
• --push_to_hub: se o modelo treinado deve ser enviado para o Hub
• --checkpointing_steps: frequência de salvamento de um checkpoint durante o treinamento do modelo; isso é útil caso o treinamento seja interrompido, pois você pode continuar o treinamento a partir desse checkpoint adicionando --resume_from_checkpoint ao seu comando de treinamento

Traga seu conjunto de dados e deixe o script de treinamento cuidar de tudo o mais!

Script de treinamento

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

O script train_unconditional inicializa um UNet2DModel caso você não forneça uma configuração de modelo. Você pode configurar o UNet aqui, se desejar:

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

Em seguida, o script inicializa um scheduler e um 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,
)

Em seguida, ele carrega um conjunto de dados e você pode especificar como pré-processá-lo:

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

Por fim, o training loop cuida de tudo o mais, como adicionar ruído às imagens, prever o resíduo de ruído, calcular a perda, salvar checkpoints em etapas específicas e salvar e enviar o modelo para o Hub. Se você quiser saber mais sobre como o training loop funciona, confira o tutorial Understanding pipelines, models and schedulers, que detalha o padrão básico do processo de remoção de ruído.

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! 🚀

Um treinamento completo leva 2 horas em 4 GPUs V100.

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

O script de treinamento cria e salva um arquivo de checkpoint no seu repositório. Agora você pode carregar e usar seu modelo treinado para inferência:

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]