# 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](https://huggingface.co/search/full-text?q=unconditional-image-generation&type=model), tetapi jika Anda tidak menemukan yang Anda sukai, Anda selalu dapat melatih model Anda sendiri!

Tutorial ini akan mengajarkan Anda cara melatih [UNet2DModel](/docs/diffusers/main/id/api/models/unet2d#diffusers.UNet2DModel) dari awal pada subset dataset [Smithsonian Butterflies](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) untuk menghasilkan kupu-kupu 🦋 Anda sendiri 🦋.

> [!TIP]
> 💡 Tutorial pelatihan ini didasarkan pada notebook [Training with 🧨 Diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb). 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](https://www.tensorflow.org/tensorboard) untuk memvisualisasikan metrik pelatihan (Anda juga dapat menggunakan [Weights & Biases](https://docs.wandb.ai/) untuk melacak pelatihan Anda).

```py
# 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](https://hf.co/join) jika Anda belum memilikinya!). Anda dapat masuk dari notebook dan memasukkan token Anda saat diminta. Pastikan token Anda memiliki peran tulis.

```py
>>> from huggingface_hub import notebook_login

>>> notebook_login()
```

Atau masuk dari terminal:

```bash
hf auth login
```

Karena checkpoint model cukup besar, instal [Git-LFS](https://git-lfs.com/) untuk memverifikasi file-file besar ini:

```bash
!sudo apt -qq install git-lfs
!git config --global credential.helper store
```

## Konfigurasi pelatihan

Untuk kenyamanan, buat kelas `TrainingConfig` yang berisi hiperparameter pelatihan (silakan sesuaikan):

```py
>>> 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](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) dengan pustaka 🤗 Datasets:

```py
>>> from datasets import load_dataset

>>> config.dataset_name = "huggan/smithsonian_butterflies_subset"
>>> dataset = load_dataset(config.dataset_name, split="train")
```

> [!TIP]
> 💡 Anda dapat menemukan dataset tambahan dari [HugGan Community Event](https://huggingface.co/huggan) atau Anda dapat menggunakan dataset Anda sendiri dengan membuat [`ImageFolder`](https://huggingface.co/docs/datasets/image_dataset#imagefolder) lokal. Atur `config.dataset_name` ke id repositori dataset jika berasal dari HugGan Community Event, atau `imagefolder` jika Anda menggunakan gambar Anda sendiri.

🤗 Datasets menggunakan fitur [Image](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Image) untuk secara otomatis mendekode data gambar dan memuatnya sebagai [`PIL.Image`](https://pillow.readthedocs.io/en/stable/reference/Image.html) yang dapat kita visualisasikan:

```py
>>> 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:

* `Resize` mengubah ukuran gambar menjadi ukuran yang didefinisikan dalam `config.image_size`.
* `RandomHorizontalFlip` memperkaya dataset dengan cara membalik gambar secara acak.
* `Normalize` penting untuk menskalakan ulang nilai piksel ke rentang [-1, 1], yang merupakan nilai yang diharapkan oleh model.

```py
>>> 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](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.set_transform) dari 🤗 Datasets untuk menerapkan fungsi `preprocess` secara langsung selama pelatihan:

```py
>>> 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](https://pytorch.org/docs/stable/data#torch.utils.data.DataLoader) untuk pelatihan!

```py
>>> 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](/docs/diffusers/main/id/api/models/unet2d#diffusers.UNet2DModel):

```py
>>> 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:

```py
>>> 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](/docs/diffusers/main/id/api/schedulers/ddpm#diffusers.DDPMScheduler) dan gunakan metode `add_noise` untuk menambahkan beberapa noise acak pada `sample_image` dari sebelumnya:

```py
>>> 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:

```py
>>> 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:

```py
>>> 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](/docs/diffusers/main/id/api/pipelines/ddpm#diffusers.DDPMPipeline) untuk menghasilkan batch gambar sampel dan menyimpannya sebagai grid:

```py
>>> 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.

> [!TIP]
> 💡 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 lebih dekat nanti, misalnya saat menunggu model Anda selesai dilatih. 🤗

```py
>>> 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](https://huggingface.co/docs/accelerate/main/en/package_reference/launchers#accelerate.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:

```py
>>> 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!

```py
>>> 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](../training/overview). Berikut adalah beberapa contoh yang dapat Anda pelajari:

* [Textual Inversion](../training/text_inversion), sebuah algoritma yang mengajarkan model konsep visual tertentu dan mengintegrasikannya ke dalam gambar yang dihasilkan.
* [DreamBooth](../training/dreambooth), sebuah teknik untuk menghasilkan gambar personalisasi dari suatu subjek berdasarkan beberapa gambar input dari subjek tersebut.
* [Panduan](../training/text2image) untuk fine-tuning model Stable Diffusion pada dataset Anda sendiri.
* [Panduan](../training/lora) untuk menggunakan LoRA, sebuah teknik hemat memori untuk finetuning model yang sangat besar dengan lebih cepat.