Diffusers documentation

创建用于训练的数据集

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

创建用于训练的数据集

Hub 上有许多可用于训练模型的数据集,但如果您找不到感兴趣的数据集,或者想使用自己的数据,可以使用 🤗 Datasets 库来创建数据集。数据集的结构取决于您希望训练模型执行的任务。最基本的数据集结构是一个包含图像的目录,适用于无条件图像生成等任务。另一种数据集结构可能是一个包含图像的目录,以及一个包含对应文本描述的文本文件,适用于文本到图像生成等任务。

本指南将向你展示创建用于微调的数据集的两种方法:

  • 将包含图像的文件夹提供给 --train_data_dir 参数
  • 将数据集上传到 Hub,并将数据集仓库 ID 传递给 --dataset_name 参数

💡 有关如何创建用于训练的图片数据集,请参阅 创建图片数据集 指南。

以文件夹形式提供数据集

对于无条件生成,您可以提供一个包含图像的文件夹作为自己的数据集。训练脚本使用 🤗 Datasets 中的 ImageFolder builder 来自动从该文件夹构建数据集。您的目录结构应如下所示:

data_dir/xxx.png
data_dir/xxy.png
data_dir/[...]/xxz.png

将数据集目录的路径传递给 --train_data_dir 参数,然后即可开始训练:

accelerate launch train_unconditional.py \
    --train_data_dir <path-to-train-directory> \
    <other-arguments>

将您的数据上传到 Hub

💡 有关在 Hub 上创建和上传数据集的更多详细信息和背景,请参阅 使用 🤗 Datasets 进行图像搜索 文章。

首先,使用 ImageFolder 功能创建一个数据集,该功能会创建一个包含 PIL 编码图像的 image 列。

您可以使用 data_dir 或 data_files 参数来指定数据集的位置。data_files 参数支持将特定文件映射到 train 或 test 等数据集拆分:

from datasets import load_dataset

# example 1: local folder
dataset = load_dataset("imagefolder", data_dir="path_to_your_folder")

# example 2: local files (supported formats are tar, gzip, zip, xz, rar, zstd)
dataset = load_dataset("imagefolder", data_files="path_to_zip_file")

# example 3: remote files (supported formats are tar, gzip, zip, xz, rar, zstd)
dataset = load_dataset(
    "imagefolder",
    data_files="https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_3367a.zip",
)

# example 4: providing several splits
dataset = load_dataset(
    "imagefolder", data_files={"train": ["path/to/file1", "path/to/file2"], "test": ["path/to/file3", "path/to/file4"]}
)

然后使用 push_to_hub 方法将数据集上传到 Hub:

# assuming you have ran the hf auth login command in a terminal
dataset.push_to_hub("name_of_your_dataset")

# if you want to push to a private repo, simply pass private=True:
dataset.push_to_hub("name_of_your_dataset", private=True)

现在,只需将数据集名称传递给 --dataset_name 参数,即可使用该数据集进行训练:

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path="stable-diffusion-v1-5/stable-diffusion-v1-5" \
  --dataset_name="name_of_your_dataset" \
  <other-arguments>

后续步骤

现在,您已经创建了一个数据集,可以将其传入训练脚本的 train_data_dir(如果数据集是本地文件)或 dataset_name(如果数据集在 Hub 上)参数中。

作为下一步,您可以尝试使用自己的数据集来训练一个用于无条件生成或文本到图像生成的模型!

Update on GitHub