Diffusers documentation
Melatih model difusi
Melatih model difusi
Generasi gambar tanpa syarat adalah aplikasi populer dari model difusi yang menghasilkan gambar yang mirip dengan gambar dalam dataset yang digunakan untuk pelatihan. Biasanya, hasil terbaik diperoleh dari penyesuaian halus (finetuning) model pra-terlatih pada dataset tertentu. Anda dapat menemukan banyak checkpoint ini di Hub, tetapi jika Anda tidak menemukan yang Anda sukai, Anda selalu dapat melatih model Anda sendiri!
Tutorial ini akan mengajarkan Anda cara melatih UNet2DModel dari awal pada subset dataset Smithsonian Butterflies untuk menghasilkan kupu-kupu π¦ Anda sendiri π¦.
π‘ Tutorial pelatihan ini didasarkan pada notebook Training with 𧨠Diffusers. Untuk detail dan konteks tambahan tentang model difusi, seperti cara kerjanya, silakan lihat notebook tersebut!
Sebelum memulai, pastikan Anda telah menginstal π€ Datasets untuk memuat dan mempreproses dataset gambar, serta π€ Accelerate, untuk menyederhanakan pelatihan pada jumlah GPU berapa pun. Perintah berikut juga akan menginstal TensorBoard untuk memvisualisasikan metrik pelatihan (Anda juga dapat menggunakan Weights & Biases untuk melacak pelatihan Anda).
# uncomment to install the necessary libraries in Colab
#!pip install diffusers[training]Kami mendorong Anda untuk membagikan model Anda kepada komunitas, dan untuk melakukannya, Anda perlu masuk ke akun Hugging Face Anda (buat satu di sini jika Anda belum memilikinya!). Anda dapat masuk dari notebook dan memasukkan token Anda saat diminta. Pastikan token Anda memiliki peran tulis.
>>> from huggingface_hub import notebook_login
>>> notebook_login()Atau masuk dari terminal:
hf auth login
Karena checkpoint model cukup besar, instal Git-LFS untuk memverifikasi file-file besar ini:
!sudo apt -qq install git-lfs
!git config --global credential.helper storeKonfigurasi pelatihan
Untuk kenyamanan, buat kelas TrainingConfig yang berisi hiperparameter pelatihan (silakan sesuaikan):
>>> from dataclasses import dataclass
>>> @dataclass
... class TrainingConfig:
... image_size = 128 # the generated image resolution
... train_batch_size = 16
... eval_batch_size = 16 # how many images to sample during evaluation
... num_epochs = 50
... gradient_accumulation_steps = 1
... learning_rate = 1e-4
... lr_warmup_steps = 500
... save_image_epochs = 10
... save_model_epochs = 30
... mixed_precision = "fp16" # `no` for float32, `fp16` for automatic mixed precision
... output_dir = "ddpm-butterflies-128" # the model name locally and on the HF Hub
... push_to_hub = True # whether to upload the saved model to the HF Hub
... hub_model_id = "<your-username>/<my-awesome-model>" # the name of the repository to create on the HF Hub
... hub_private_repo = None
... overwrite_output_dir = True # overwrite the old model when re-running the notebook
... seed = 0
>>> config = TrainingConfig()Muat dataset
Anda dapat dengan mudah memuat dataset Smithsonian Butterflies dengan pustaka π€ Datasets:
>>> from datasets import load_dataset
>>> config.dataset_name = "huggan/smithsonian_butterflies_subset"
>>> dataset = load_dataset(config.dataset_name, split="train")π‘ Anda dapat menemukan dataset tambahan dari HugGan Community Event atau Anda dapat menggunakan dataset Anda sendiri dengan membuat
ImageFolderlokal. Aturconfig.dataset_nameke id repositori dataset jika berasal dari HugGan Community Event, atauimagefolderjika Anda menggunakan gambar Anda sendiri.
π€ Datasets menggunakan fitur Image untuk secara otomatis mendekode data gambar dan memuatnya sebagai PIL.Image yang dapat kita visualisasikan:
>>> import matplotlib.pyplot as plt
>>> fig, axs = plt.subplots(1, 4, figsize=(16, 4))
>>> for i, image in enumerate(dataset[:4]["image"]):
... axs[i].imshow(image)
... axs[i].set_axis_off()
>>> fig.show()
Namun, semua gambar memiliki ukuran yang berbeda, jadi Anda perlu memprosesnya terlebih dahulu:
Resizemengubah ukuran gambar menjadi ukuran yang didefinisikan dalamconfig.image_size.RandomHorizontalFlipmemperkaya dataset dengan cara membalik gambar secara acak.Normalizepenting untuk menskalakan ulang nilai piksel ke rentang [-1, 1], yang merupakan nilai yang diharapkan oleh model.
>>> from torchvision import transforms
>>> preprocess = transforms.Compose(
... [
... transforms.Resize((config.image_size, config.image_size)),
... transforms.RandomHorizontalFlip(),
... transforms.ToTensor(),
... transforms.Normalize([0.5], [0.5]),
... ]
... )Gunakan metode set_transform dari π€ Datasets untuk menerapkan fungsi preprocess secara langsung selama pelatihan:
>>> def transform(examples):
... images = [preprocess(image.convert("RGB")) for image in examples["image"]]
... return {"images": images}
>>> dataset.set_transform(transform)Silakan visualisasikan gambar-gambar tersebut sekali lagi untuk memastikan bahwa ukurannya telah diubah. Sekarang Anda siap membungkus dataset dalam DataLoader untuk pelatihan!
>>> import torch
>>> train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)Buat UNet2DModel
Model pretrained di 𧨠Diffusers dengan mudah dibuat dari kelas modelnya dengan parameter yang Anda inginkan. Misalnya, untuk membuat UNet2DModel:
>>> from diffusers import UNet2DModel
>>> model = UNet2DModel(
... sample_size=config.image_size, # the target image resolution
... in_channels=3, # the number of input channels, 3 for RGB images
... out_channels=3, # the number of output channels
... layers_per_block=2, # how many ResNet layers to use per UNet block
... block_out_channels=(128, 128, 256, 256, 512, 512), # the number of output channels for each UNet block
... down_block_types=(
... "DownBlock2D", # a regular ResNet downsampling block
... "DownBlock2D",
... "DownBlock2D",
... "DownBlock2D",
... "AttnDownBlock2D", # a ResNet downsampling block with spatial self-attention
... "DownBlock2D",
... ),
... up_block_types=(
... "UpBlock2D", # a regular ResNet upsampling block
... "AttnUpBlock2D", # a ResNet upsampling block with spatial self-attention
... "UpBlock2D",
... "UpBlock2D",
... "UpBlock2D",
... "UpBlock2D",
... ),
... )Sering kali merupakan ide yang baik untuk dengan cepat memeriksa apakah bentuk gambar sampel cocok dengan bentuk keluaran model:
>>> sample_image = dataset[0]["images"].unsqueeze(0)
>>> print("Input shape:", sample_image.shape)
Input shape: torch.Size([1, 3, 128, 128])
>>> print("Output shape:", model(sample_image, timestep=0).sample.shape)
Output shape: torch.Size([1, 3, 128, 128])Bagus! Selanjutnya, Anda akan membutuhkan scheduler untuk menambahkan beberapa noise ke gambar.
Buat scheduler
Perilaku scheduler berbeda tergantung apakah Anda menggunakan model untuk pelatihan atau inferensi. Selama inferensi, scheduler menghasilkan gambar dari noise. Selama pelatihan, scheduler mengambil output model - atau sampel - dari titik tertentu dalam proses difusi dan menerapkan noise ke gambar sesuai dengan noise schedule dan update rule.
Mari kita lihat DDPMScheduler dan gunakan metode add_noise untuk menambahkan beberapa noise acak pada sample_image dari sebelumnya:
>>> import torch
>>> from PIL import Image
>>> from diffusers import DDPMScheduler
>>> noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
>>> noise = torch.randn(sample_image.shape)
>>> timesteps = torch.LongTensor([50])
>>> noisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)
>>> Image.fromarray(((noisy_image.permute(0, 2, 3, 1) + 1.0) * 127.5).type(torch.uint8).numpy()[0])
Tujuan pelatihan model adalah untuk memprediksi noise yang ditambahkan ke gambar. Loss pada langkah ini dapat dihitung dengan:
>>> import torch.nn.functional as F
>>> noise_pred = model(noisy_image, timesteps).sample
>>> loss = F.mse_loss(noise_pred, noise)Latih model
Saat ini, Anda sudah memiliki sebagian besar komponen untuk mulai melatih model, dan yang tersisa hanyalah merangkai semuanya.
Pertama, Anda memerlukan optimizer dan learning rate scheduler:
>>> from diffusers.optimization import get_cosine_schedule_with_warmup
>>> optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
>>> lr_scheduler = get_cosine_schedule_with_warmup(
... optimizer=optimizer,
... num_warmup_steps=config.lr_warmup_steps,
... num_training_steps=(len(train_dataloader) * config.num_epochs),
... )Lalu, Anda memerlukan cara untuk mengevaluasi model. Untuk evaluasi, Anda dapat menggunakan DDPMPipeline untuk menghasilkan batch gambar sampel dan menyimpannya sebagai grid:
>>> from diffusers import DDPMPipeline
>>> from diffusers.utils import make_image_grid
>>> import os
>>> def evaluate(config, epoch, pipeline):
... # Sample some images from random noise (this is the backward diffusion process).
... # The default pipeline output type is `List[PIL.Image]`
... images = pipeline(
... batch_size=config.eval_batch_size,
... generator=torch.Generator(device='cpu').manual_seed(config.seed), # Use a separate torch generator to avoid rewinding the random state of the main training loop
... ).images
... # Make a grid out of the images
... image_grid = make_image_grid(images, rows=4, cols=4)
... # Save the images
... test_dir = os.path.join(config.output_dir, "samples")
... os.makedirs(test_dir, exist_ok=True)
... image_grid.save(f"{test_dir}/{epoch:04d}.png")Sekarang Anda dapat menggabungkan semua komponen ini dalam loop pelatihan dengan π€ Accelerate untuk logging TensorBoard yang mudah, akumulasi gradien, dan pelatihan presisi campuran. Untuk mengunggah model ke Hub, tulis fungsi untuk mendapatkan nama dan informasi repositori Anda, lalu dorong ke Hub.
π‘ Loop pelatihan di bawah ini mungkin terlihat mengintimidasi dan panjang, tetapi akan sepadan nanti saat Anda meluncurkan pelatihan hanya dengan satu baris kode! Jika Anda tidak sabar dan ingin mulai menghasilkan gambar, silakan salin dan jalankan kode di bawah ini. Anda selalu bisa kembali dan memeriksa loop pelatihan secara lebih detail nanti, misalnya saat Anda menunggu model Anda selesai dilatih. π€
>>> from accelerate import Accelerator
>>> from huggingface_hub import create_repo, upload_folder
>>> from tqdm.auto import tqdm
>>> from pathlib import Path
>>> import os
>>> def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler):
... # Initialize accelerator and tensorboard logging
... accelerator = Accelerator(
... mixed_precision=config.mixed_precision,
... gradient_accumulation_steps=config.gradient_accumulation_steps,
... log_with="tensorboard",
... project_dir=os.path.join(config.output_dir, "logs"),
... )
... if accelerator.is_main_process:
... if config.output_dir is not None:
... os.makedirs(config.output_dir, exist_ok=True)
... if config.push_to_hub:
... repo_id = create_repo(
... repo_id=config.hub_model_id or Path(config.output_dir).name, exist_ok=True
... ).repo_id
... accelerator.init_trackers("train_example")
... # Prepare everything
... # There is no specific order to remember, you just need to unpack the
... # objects in the same order you gave them to the prepare method.
... model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
... model, optimizer, train_dataloader, lr_scheduler
... )
... global_step = 0
... # Now you train the model
... for epoch in range(config.num_epochs):
... progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
... progress_bar.set_description(f"Epoch {epoch}")
... for step, batch in enumerate(train_dataloader):
... clean_images = batch["images"]
... # Sample noise to add to the images
... noise = torch.randn(clean_images.shape, device=clean_images.device)
... bs = clean_images.shape[0]
... # Sample a random timestep for each image
... timesteps = torch.randint(
... 0, noise_scheduler.config.num_train_timesteps, (bs,), device=clean_images.device,
... dtype=torch.int64
... )
... # Add noise to the clean images according to the noise magnitude at each timestep
... # (this is the forward diffusion process)
... noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
... with accelerator.accumulate(model):
... # Predict the noise residual
... noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
... loss = F.mse_loss(noise_pred, noise)
... accelerator.backward(loss)
... if accelerator.sync_gradients:
... accelerator.clip_grad_norm_(model.parameters(), 1.0)
... optimizer.step()
... lr_scheduler.step()
... optimizer.zero_grad()
... progress_bar.update(1)
... logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
... progress_bar.set_postfix(**logs)
... accelerator.log(logs, step=global_step)
... global_step += 1
... # After each epoch you optionally sample some demo images with evaluate() and save the model
... if accelerator.is_main_process:
... pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)
... if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:
... evaluate(config, epoch, pipeline)
... if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
... if config.push_to_hub:
... upload_folder(
... repo_id=repo_id,
... folder_path=config.output_dir,
... commit_message=f"Epoch {epoch}",
... ignore_patterns=["step_*", "epoch_*"],
... )
... else:
... pipeline.save_pretrained(config.output_dir)Phew, itu cukup banyak kode! Tapi Anda akhirnya siap untuk meluncurkan pelatihan dengan fungsi notebook_launcher dari π€ Accelerate. Berikan fungsi tersebut loop pelatihan, semua argumen pelatihan, dan jumlah proses (Anda dapat mengubah nilai ini menjadi jumlah GPU yang tersedia untuk Anda) yang akan digunakan untuk pelatihan:
>>> from accelerate import notebook_launcher
>>> args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)
>>> notebook_launcher(train_loop, args, num_processes=1)Setelah pelatihan selesai, lihatlah gambar π¦ akhir π¦ yang dihasilkan oleh model difusi Anda!
>>> import glob
>>> sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png"))
>>> Image.open(sample_images[-1])
Langkah selanjutnya
Generasi gambar tanpa syarat adalah salah satu contoh tugas yang dapat dilatih. Anda dapat menjelajahi tugas dan teknik pelatihan lainnya dengan mengunjungi halaman 𧨠Contoh Pelatihan Diffusers. Berikut adalah beberapa contoh yang dapat Anda pelajari:
- Textual Inversion, sebuah algoritma yang mengajarkan model konsep visual tertentu dan mengintegrasikannya ke dalam gambar yang dihasilkan.
- DreamBooth, sebuah teknik untuk menghasilkan gambar personalisasi dari suatu subjek berdasarkan beberapa gambar input dari subjek tersebut.
- Panduan untuk fine-tuning model Stable Diffusion pada dataset Anda sendiri.
- Panduan untuk menggunakan LoRA, sebuah teknik hemat memori untuk finetuning model yang sangat besar dengan lebih cepat.