Fine-Tune ViT for Image Classification with 🤗 Transformers

Summary

Vision Transformer models bring transformer architectures to computer vision by splitting images into grids of sub-image patches and projecting them into token sequences. To classify healthy and diseased leaves using the beans dataset, the Hugging Face datasets and transformers libraries enable streamlined data ingestion and model fine-tuning. Preprocessing relies on ViTImageProcessor paired with lazy on-the-fly dataset transforms to dynamically generate normalized pixel tensors. Training leverages ViTForImageClassification with the Trainer API, requiring remove_unused_columns set to False so raw image data is preserved for batch collation. Over four training epochs, fine-tuning the google/vit-base-patch16-224-in21k checkpoint achieves an evaluation accuracy of 98.5% alongside an evaluation loss of 0.0637.

Takeaways

1. 1

   ViT tokenizes an image by splitting it into a grid of sub-image patches and projecting each patch linearly into a sequence of input tokens for the transformer.

2. 2

   Using ds.with_transform applies image processor transformations lazily upon indexing, avoiding the slower execution of mapping preprocessing eagerly across large datasets.

3. 3

   Setting remove_unused_columns=False in TrainingArguments is required when raw image features must be kept in the batch to generate pixel_values for the model.