Crie um conjunto de dados para treinamento

Existem muitos conjuntos de dados no Hub para treinar um modelo, mas, se você não encontrar um que lhe interesse ou quiser usar o seu próprio, pode criar um conjunto de dados com a biblioteca 🤗 Datasets. A estrutura do conjunto de dados depende da tarefa na qual você deseja treinar seu modelo. A estrutura mais básica de um conjunto de dados é um diretório de imagens para tarefas como a geração de imagens incondicional. Outra estrutura de conjunto de dados pode ser um diretório de imagens e um arquivo de texto contendo as legendas de texto correspondentes para tarefas como a geração de imagem a partir de texto.

Este guia mostrará duas maneiras de criar um conjunto de dados para ajuste fino:

• fornecer uma pasta de imagens para o argumento --train_data_dir
• envie um conjunto de dados para o Hub e passe o id do repositório do conjunto de dados para o argumento --dataset_name

💡 Saiba mais sobre como criar um conjunto de dados de imagens para treinamento no guia Criar um conjunto de dados de imagens.

Forneça um conjunto de dados como uma pasta

Para geração incondicional, você pode fornecer seu próprio conjunto de dados como uma pasta de imagens. O script de treinamento usa o construtor ImageFolder do 🤗 Datasets para construir automaticamente um conjunto de dados a partir da pasta. A estrutura do seu diretório deve ser:

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

Passe o caminho para o diretório do conjunto de dados para o argumento --train_data_dir e, em seguida, você pode iniciar o treinamento:

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

Envie seus dados para o Hub

💡 Para mais detalhes e contexto sobre a criação e o envio de um dataset para o Hub, confira o post Image search with 🤗 Datasets.

Comece criando um conjunto de dados com o recurso ImageFolder, que cria uma coluna image contendo as imagens codificadas em PIL.

Você pode usar os parâmetros data_dir ou data_files para especificar a localização do conjunto de dados. O parâmetro data_files permite mapear arquivos específicos para divisões do conjunto de dados, como train ou 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"]}
)

Em seguida, use o método push_to_hub para enviar o conjunto de dados para o 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)

Agora o conjunto de dados está disponível para treinamento, passando o nome do conjunto de dados para o argumento --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>

Próximas etapas

Agora que você criou um conjunto de dados, pode inseri-lo nos argumentos train_data_dir (se o seu conjunto de dados for local) ou dataset_name (se o seu conjunto de dados estiver no Hub) de um script de treinamento.

Para os seus próximos passos, sinta-se à vontade para tentar usar seu conjunto de dados para treinar um modelo de geração incondicional ou geração de imagem a partir de texto!