Skip to content

创建用于训练的数据集

Hub 上有很多数据集可用于训练模型,但如果你找不到你感兴趣的数据集,或者想使用自己的数据集,你可以使用 🤗 Datasets 库来创建数据集。数据集结构取决于你想要训练模型的任务。最基本的数据集结构是用于无条件图像生成的图像目录。另一个数据集结构可能是图像目录和包含其对应文本标题的文本文件,用于文本到图像生成等任务。

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

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

将数据集作为文件夹提供

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

└── train
    └── image1.jpg
    └── image2.jpg
    └── image3.jpg
    └── ...

然后,你可以使用以下命令启动训练:

bash
python train.py --train_data_dir /path/to/your/dataset/train

将数据集上传到 Hub

如果你想将数据集上传到 Hub,你可以使用 🤗 Datasets 库来创建数据集。以下是如何创建数据集的示例:

python
from datasets import load_dataset, ClassLabel, Sequence
import os

def create_dataset(path):
    """
    创建数据集。

    Args:
        path: 数据集路径。

    Returns:
        数据集。
    """

    dataset = load_dataset("imagefolder", data_dir=path)

    # 将数据集上传到 Hub
    dataset.push_to_hub("your_username/your_dataset_name")

    return dataset

# 创建数据集
dataset = create_dataset("/path/to/your/dataset")

# 打印数据集信息
print(dataset)

然后,你可以使用以下命令启动训练:

bash
python train.py --dataset_name your_username/your_dataset_name
bash
data_dir/xxx.png
data_dir/xxy.png
data_dir/[...]/xxz.png

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

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

上传你的数据到 Hub

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

你可以使用 data_dirdata_files 参数来指定数据集的位置。data_files 参数支持将特定文件映射到数据集拆分,例如 traintest

python
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"]}
)

然后使用 [~datasets.Dataset.push_to_hub] 方法将数据集上传到 Hub:

python
# assuming you have ran the huggingface-cli 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 参数来用于训练:

bash
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 上)参数中。

接下来,你可以尝试使用你的数据集来训练一个模型,用于 无条件生成文本到图像生成