Diffusers documentation
Huấn luyện một mô hình khuếch tán
Huấn luyện một mô hình khuếch tán
Tạo ảnh không điều kiện là một ứng dụng phổ biến của mô hình khuếch tán, tạo ra các hình ảnh trông giống với những hình ảnh trong bộ dữ liệu được sử dụng để huấn luyện. Thông thường, kết quả tốt nhất đạt được từ việc tinh chỉnh (finetuning) một mô hình đã được huấn luyện trước trên một bộ dữ liệu cụ thể. Bạn có thể tìm thấy nhiều checkpoint như vậy trên Hub, nhưng nếu không tìm được checkpoint nào ưng ý, bạn luôn có thể tự huấn luyện một checkpoint của riêng mình!
Hướng dẫn này sẽ dạy bạn cách huấn luyện một UNet2DModel từ đầu trên một tập con của bộ dữ liệu Smithsonian Butterflies để tạo ra những chú bướm 🦋 của riêng bạn 🦋.
💡 Bài hướng dẫn huấn luyện này dựa trên notebook Training with 🧨 Diffusers. Để biết thêm chi tiết và ngữ cảnh về các mô hình khuếch tán, chẳng hạn như cách chúng hoạt động, hãy xem notebook!
Trước khi bắt đầu, hãy đảm bảo bạn đã cài đặt 🤗 Datasets để tải và tiền xử lý các bộ dữ liệu ảnh, cùng với 🤗 Accelerate để đơn giản hóa việc huấn luyện trên bất kỳ số lượng GPU nào. Lệnh sau đây cũng sẽ cài đặt TensorBoard để trực quan hóa các chỉ số huấn luyện (bạn cũng có thể sử dụng Weights & Biases để theo dõi quá trình huấn luyện của mình).
# uncomment to install the necessary libraries in Colab
#!pip install diffusers[training]Chúng tôi khuyến khích bạn chia sẻ mô hình của mình với cộng đồng, và để làm được điều đó, bạn cần đăng nhập vào tài khoản Hugging Face của mình (tạo một tài khoản tại đây nếu bạn chưa có!). Bạn có thể đăng nhập từ notebook và nhập token khi được yêu cầu. Hãy đảm bảo token của bạn có vai trò ghi (write).
>>> from huggingface_hub import notebook_login
>>> notebook_login()Hoặc đăng nhập từ terminal:
hf auth login
Vì các checkpoint của mô hình khá lớn, hãy cài đặt Git-LFS để quản lý phiên bản cho các tệp lớn này:
!sudo apt -qq install git-lfs
!git config --global credential.helper storeCấu hình huấn luyện
Để tiện lợi, hãy tạo một lớp TrainingConfig chứa các siêu tham số huấn luyện (bạn có thể tự do điều chỉnh chúng):
>>> 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()Tải bộ dữ liệu
Bạn có thể dễ dàng tải bộ dữ liệu Smithsonian Butterflies bằng thư viện 🤗 Datasets:
>>> from datasets import load_dataset
>>> config.dataset_name = "huggan/smithsonian_butterflies_subset"
>>> dataset = load_dataset(config.dataset_name, split="train")💡 Bạn có thể tìm thêm các bộ dữ liệu từ Sự kiện Cộng đồng HugGan hoặc sử dụng bộ dữ liệu của riêng bạn bằng cách tạo một
ImageFoldercục bộ. Đặtconfig.dataset_namethành id của kho dữ liệu nếu nó đến từ Sự kiện Cộng đồng HugGan, hoặcimagefoldernếu bạn đang sử dụng hình ảnh của riêng mình.
🤗 Datasets sử dụng tính năng Image để tự động giải mã dữ liệu ảnh và nạp nó dưới dạng PIL.Image, cho phép chúng ta trực quan hóa:
>>> 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()
Tuy nhiên, tất cả các hình ảnh đều có kích thước khác nhau, vì vậy bạn sẽ cần tiền xử lý chúng trước:
Resizethay đổi kích thước ảnh thành kích thước được định nghĩa trongconfig.image_size.RandomHorizontalFliptăng cường bộ dữ liệu bằng cách lật ngẫu nhiên các hình ảnh theo chiều ngang.Normalizerất quan trọng để chuẩn hóa lại giá trị pixel vào khoảng [-1, 1], đây là điều mà mô hình mong đợi.
>>> 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]),
... ]
... )Sử dụng phương thức set_transform của 🤗 Datasets để áp dụng hàm preprocess trực tiếp trong quá trình huấn luyện:
>>> def transform(examples):
... images = [preprocess(image.convert("RGB")) for image in examples["image"]]
... return {"images": images}
>>> dataset.set_transform(transform)Hãy tự do trực quan hóa lại các hình ảnh để xác nhận rằng chúng đã được thay đổi kích thước. Bây giờ bạn đã sẵn sàng bọc dataset trong một DataLoader để huấn luyện!
>>> import torch
>>> train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)Tạo một UNet2DModel
Các mô hình được huấn luyện trước trong 🧨 Diffusers có thể dễ dàng được tạo ra từ lớp mô hình của chúng với các tham số bạn mong muốn. Ví dụ, để tạo một 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",
... ),
... )Thường là một ý tưởng hay để nhanh chóng kiểm tra xem hình dạng của hình ảnh mẫu có khớp với hình dạng đầu ra của mô hình không:
>>> 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])Tuyệt! Tiếp theo, bạn sẽ cần một trình lập lịch (scheduler) để thêm nhiễu vào hình ảnh.
Tạo một trình lập lịch
Trình điều phối hoạt động khác nhau tùy thuộc vào việc bạn sử dụng mô hình cho huấn luyện hay suy luận. Trong quá trình suy luận, trình điều phối tạo ra hình ảnh từ nhiễu. Trong quá trình huấn luyện, trình điều phối nhận một đầu ra của mô hình - hoặc một mẫu - từ một điểm cụ thể trong quá trình khuếch tán và áp dụng nhiễu vào hình ảnh theo một lịch trình nhiễu và một quy tắc cập nhật.
Hãy cùng xem xét DDPMScheduler và sử dụng phương thức add_noise để thêm một chút nhiễu ngẫu nhiên vào sample_image từ trước đó:
>>> 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])
Mục tiêu huấn luyện của mô hình là dự đoán nhiễu được thêm vào ảnh. Tổn thất ở bước này có thể được tính bằng:
>>> import torch.nn.functional as F
>>> noise_pred = model(noisy_image, timesteps).sample
>>> loss = F.mse_loss(noise_pred, noise)Huấn luyện mô hình
Đến đây, bạn đã có hầu hết các thành phần cần thiết để bắt đầu huấn luyện mô hình và điều còn lại chỉ là ghép tất cả lại với nhau.
Đầu tiên, bạn sẽ cần một optimizer và một 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),
... )Sau đó, bạn sẽ cần một cách để đánh giá mô hình. Để đánh giá, bạn có thể sử dụng DDPMPipeline để tạo một lô ảnh mẫu và lưu chúng dưới dạng lưới:
>>> 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")Bây giờ bạn có thể kết hợp tất cả các thành phần này trong một vòng lặp huấn luyện với 🤗 Accelerate để dễ dàng ghi log TensorBoard, tích lũy gradient và huấn luyện độ chính xác hỗn hợp. Để tải mô hình lên Hub, hãy viết một hàm để lấy tên và thông tin kho lưu trữ của bạn, sau đó đẩy nó lên Hub.
💡 Vòng lặp huấn luyện bên dưới có thể trông đáng sợ và dài dòng, nhưng nó sẽ xứng đáng khi bạn khởi chạy huấn luyện chỉ bằng một dòng mã sau này! Nếu bạn không thể chờ đợi và muốn bắt đầu tạo hình ảnh, hãy thoải mái sao chép và chạy mã bên dưới. Bạn luôn có thể quay lại và xem xét kỹ hơn vòng lặp huấn luyện sau, chẳng hạn như khi bạn đang chờ mô hình của mình hoàn tất huấn luyện. 🤗
>>> 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, khá nhiều mã! Nhưng cuối cùng bạn đã sẵn sàng để khởi động quá trình huấn luyện bằng hàm notebook_launcher của 🤗 Accelerate. Truyền cho hàm này vòng lặp huấn luyện, tất cả các đối số huấn luyện và số lượng tiến trình (bạn có thể thay đổi giá trị này thành số lượng GPU khả dụng của bạn) sẽ được sử dụng cho huấn luyện:
>>> from accelerate import notebook_launcher
>>> args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)
>>> notebook_launcher(train_loop, args, num_processes=1)Khi quá trình huấn luyện hoàn tất, hãy xem các hình ảnh 🦋 cuối cùng 🦋 do mô hình khuếch tán của bạn tạo ra!
>>> import glob
>>> sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png"))
>>> Image.open(sample_images[-1])
Các bước tiếp theo
Tạo ảnh không điều kiện là một ví dụ về tác vụ có thể được huấn luyện. Bạn có thể khám phá các tác vụ và kỹ thuật huấn luyện khác bằng cách truy cập trang 🧨 Diffusers Training Examples. Dưới đây là một số ví dụ về những gì bạn có thể học:
- Textual Inversion, một thuật toán dạy cho mô hình một khái niệm thị giác cụ thể và tích hợp nó vào hình ảnh được tạo ra.
- DreamBooth, một kỹ thuật để tạo ra các hình ảnh cá nhân hóa của một chủ thể dựa trên một vài hình ảnh đầu vào của chủ thể đó.
- Hướng dẫn tinh chỉnh mô hình Stable Diffusion trên bộ dữ liệu của bạn.
- Hướng dẫn sử dụng LoRA, một kỹ thuật tiết kiệm bộ nhớ để tinh chỉnh các mô hình rất lớn nhanh hơn.