Tạo ảnh không điều kiện

Các mô hình tạo ảnh không điều kiện không được điều kiện hóa dựa trên văn bản hay hình ảnh trong quá trình huấn luyện. Chúng chỉ tạo ra các hình ảnh có phân phối tương tự như dữ liệu huấn luyện của chúng.

Hướng dẫn này sẽ khám phá script huấn luyện train_unconditional.py để giúp bạn làm quen với nó, cũng như cách bạn có thể điều chỉnh nó cho trường hợp sử dụng của riêng mình.

Trước khi chạy script, hãy đảm bảo bạn cài đặt thư viện từ mã nguồn:

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

Sau đó, điều hướng đến thư mục ví dụ chứa script huấn luyện và cài đặt các phụ thuộc cần thiết:

cd examples/unconditional_image_generation
pip install -r requirements.txt

🤗 Accelerate là một thư viện giúp bạn huấn luyện trên nhiều GPU/TPU hoặc với độ chính xác hỗn hợp. Nó sẽ tự động cấu hình thiết lập huấn luyện của bạn dựa trên phần cứng và môi trường của bạn. Hãy xem 🤗 Accelerate Quick tour để tìm hiểu thêm.

Khởi tạo một môi trường 🤗 Accelerate:

accelerate config

Để thiết lập môi trường 🤗 Accelerate mặc định mà không cần chọn bất kỳ cấu hình nào:

accelerate config default

Hoặc nếu môi trường của bạn không hỗ trợ shell tương tác như notebook, bạn có thể sử dụng:

from accelerate.utils import write_basic_config

write_basic_config()

Cuối cùng, nếu bạn muốn huấn luyện một mô hình trên bộ dữ liệu của riêng mình, hãy xem hướng dẫn Tạo bộ dữ liệu để huấn luyện để tìm hiểu cách tạo một bộ dữ liệu tương thích với script huấn luyện.

Tham số của script

Các phần tiếp theo làm nổi bật những phần quan trọng trong script huấn luyện để hiểu cách chỉnh sửa nó, nhưng không đi sâu vào mọi khía cạnh của script. Nếu bạn muốn tìm hiểu thêm, hãy tự do đọc qua script và cho chúng tôi biết nếu bạn có bất kỳ câu hỏi hoặc lo ngại nào.

Tệp kịch bản huấn luyện cung cấp nhiều tham số để giúp bạn tùy chỉnh phiên huấn luyện của mình. Tất cả các tham số và mô tả của chúng được tìm thấy trong hàm parse_args(). Hàm này cung cấp các giá trị mặc định cho từng tham số, chẳng hạn như kích thước lô huấn luyện và tốc độ học, nhưng bạn cũng có thể tự đặt giá trị của riêng mình trong lệnh huấn luyện nếu muốn.

Ví dụ, để tăng tốc độ huấn luyện với độ chính xác hỗn hợp sử dụng định dạng bf16, hãy thêm tham số --mixed_precision vào lệnh huấn luyện:

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

Một số tham số cơ bản và quan trọng cần chỉ định bao gồm:

--dataset_name: tên của bộ dữ liệu trên Hub hoặc đường dẫn cục bộ đến bộ dữ liệu để huấn luyện
--output_dir: nơi lưu trữ mô hình đã được huấn luyện
--push_to_hub: có đẩy mô hình đã huấn luyện lên Hub hay không
--checkpointing_steps: tần suất lưu checkpoint trong quá trình huấn luyện mô hình; điều này hữu ích nếu quá trình huấn luyện bị gián đoạn, bạn có thể tiếp tục huấn luyện từ checkpoint đó bằng cách thêm --resume_from_checkpoint vào lệnh huấn luyện của mình

Hãy chuẩn bị sẵn bộ dữ liệu của bạn và để script huấn luyện xử lý mọi thứ còn lại!

Script huấn luyện

Mã để tiền xử lý bộ dữ liệu và vòng lặp huấn luyện nằm trong hàm main(). Nếu bạn cần điều chỉnh script huấn luyện, đây là nơi bạn sẽ thực hiện các thay đổi.

Script train_unconditional sẽ khởi tạo một UNet2DModel nếu bạn không cung cấp cấu hình mô hình. Bạn có thể cấu hình UNet tại đây nếu muốn:

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

Tiếp theo, script khởi tạo một scheduler và 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,
)

Sau đó, nó tải một bộ dữ liệu và bạn có thể chỉ định cách xử lý trước nó:

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

Cuối cùng, vòng lặp huấn luyện xử lý mọi thứ còn lại như thêm nhiễu vào ảnh, dự đoán phần dư nhiễu, tính toán loss, lưu checkpoint tại các bước được chỉ định, cũng như lưu và đẩy mô hình lên Hub. Nếu bạn muốn tìm hiểu thêm về cách vòng lặp huấn luyện hoạt động, hãy xem hướng dẫn Hiểu về pipelines, models và schedulers trong đó phân tích chi tiết mô hình cơ bản của quá trình khử nhiễu.

Chạy script

Khi bạn đã thực hiện tất cả các thay đổi hoặc hài lòng với cấu hình mặc định, bạn đã sẵn sàng để chạy script huấn luyện! 🚀

Một lần chạy huấn luyện đầy đủ mất 2 giờ trên 4 GPU V100.

một GPU

đa GPU

Tệp kịch bản huấn luyện tạo và lưu một tệp checkpoint trong kho của bạn. Bây giờ bạn có thể tải và sử dụng mô hình đã được huấn luyện của mình để suy luận:

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]

Update on GitHub

Diffusers

Tạo ảnh không điều kiện

Tham số của script

Script huấn luyện

Chạy script