transformers/examples/pytorch/image-classification

Image classification example

This directory contains a script, run_image_classification.py, that showcases how to fine-tune any model supported by the AutoModelForImageClassification API (such as ViT, ConvNeXT, ResNet, Swin Transformer...) using PyTorch. It can be used to fine-tune models on both well-known datasets (like CIFAR-10, Fashion MNIST, ...) as well as on your own custom data.

This page includes 2 sections:

• Using datasets from the 🤗 hub
• Using your own data.

Using datasets from Hub

Here we show how to fine-tune a Vision Transformer (ViT) on the beans dataset, to classify the disease type of bean leaves.

👀 See the results here: nateraw/vit-base-beans.

python run_image_classification.py \
    --dataset_name beans \
    --output_dir ./beans_outputs/ \
    --remove_unused_columns False \
    --do_train \
    --do_eval \
    --push_to_hub \
    --push_to_hub_model_id vit-base-beans \
    --learning_rate 2e-5 \
    --num_train_epochs 5 \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 8 \
    --logging_strategy steps \
    --logging_steps 10 \
    --evaluation_strategy epoch \
    --save_strategy epoch \
    --load_best_model_at_end True \
    --save_total_limit 3 \
    --seed 1337

To fine-tune another model, simply provide the --model_name_or_path argument. To train on another dataset, simply set the --dataset_name argument.

Using your own data

To use your own dataset, there are 2 ways:

• you can either provide your own folders as --train_dir and/or --validation_dir arguments
• you can upload your dataset to the hub (possibly as a private repo, if you prefer so), and simply pass the --dataset_name argument.

Below, we explain both in more detail.

Provide them as folders

If you provide your own folders with images, the script expects the following directory structure:

root/dog/xxx.png
root/dog/xxy.png
root/dog/[...]/xxz.png

root/cat/123.png
root/cat/nsdf3.png
root/cat/[...]/asd932_.png

In other words, you need to organize your images in subfolders, based on their class. You can then run the script like this:

python run_image_classification.py \
    --train_dir <path-to-train-root> \
    --output_dir ./outputs/ \
    --remove_unused_columns False \
    --do_train \
    --do_eval

Internally, the script will use the ImageFolder feature which will automatically turn the folders into 🤗 Dataset objects.

💡 The above will split the train dir into training and evaluation sets

• To control the split amount, use the --train_val_split flag.
• To provide your own validation split in its own directory, you can pass the --validation_dir <path-to-val-root> flag.

Upload your data to the hub, as a (possibly private) repo

It's very easy (and convenient) to upload your image dataset to the hub using the ImageFolder feature available in 🤗 Datasets. Simply do the following:

from datasets import load_dataset

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

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

# example 3: remote files (suppoted 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"]})

ImageFolder will create a label column, and the label name is based on the directory name.

Next, push it to the hub!

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)

and that's it! You can now simply train your model simply by setting the --dataset_name argument to the name of your dataset on the hub (as explained in Using datasets from the 🤗 hub).

More on this can also be found in this blog post.

Sharing your model on 🤗 Hub

0. If you haven't already, sign up for a 🤗 account

1. Make sure you have git-lfs installed and git set up.

$ apt install git-lfs
$ git config --global user.email "you@example.com"
$ git config --global user.name "Your Name"

2. Log in with your HuggingFace account credentials using huggingface-cli:

$ huggingface-cli login
# ...follow the prompts

3. When running the script, pass the following arguments:

python run_image_classification.py \
    --push_to_hub \
    --push_to_hub_model_id <name-your-model> \
    ...