Fine-Tune a ViT Food Classifier

Fine-tune a Vision Transformer to classify 101 food types using Hugging Face.

Introduction

30 Second Summary

Ever wonder how your phone's camera can identify a dish at a restaurant, or how food delivery apps auto-tag menu photos? Behind the scenes, a machine learning model has been trained to recognize visual patterns in food images.

In this project, you will fine-tune a pre-trained Vision Transformer model to classify food images into 101 categories. You'll do it all inside a free Google Colab notebook with a GPU.

What You'll Build

A food image classifier powered by Google's ViT model that can distinguish 101 types of food, from sushi to steak to tiramisu.

By the end of this project, you'll have:

  • A Google Colab notebook containing your complete ML pipeline from data loading to prediction.
  • A fine-tuned ViT model trained on 4,000 food images across 101 categories.
  • An inference pipeline that takes any food photo and returns predictions with confidence scores.
  • Secret Mission: Upload your own food photos and build a visual confidence chart showing the model's top-5 predictions.

Are there any prerequisites?

You just need a Google account. Everything else (GPU, Python libraries, dataset) is provided for free inside Google Colab. No local installation required.

Set Up Your Colab Environment

Before you can train a machine learning model, you need the right environment. Training a Vision Transformer requires a GPU to process thousands of images efficiently.

Google Colab gives you a free GPU right in your browser. In this step, you'll set up your Colab notebook, install the Python libraries for the project, and verify everything is ready to go.

In this step, get ready to:
  • Open Google Colab and enable the free GPU runtime.
  • Install the required Python libraries and verify GPU access.
  • Import all necessary modules for the project.
Open Google Colab and enable the GPU runtime

Google Colab is a free, browser-based notebook environment that gives you access to GPUs. No installation required on your own machine.

  • Open colab.research.google.com in your browser and sign in with your Google account.
  • Click New Notebook to create a fresh notebook.
  • In the top menu, click Runtime.
  • Select Change runtime type.
  • Under Hardware accelerator, choose T4 GPU, then click Save.

Why do you need a GPU?

A GPU (Graphics Processing Unit) can process thousands of image calculations in parallel. Training a model on images with a CPU alone would take hours. The T4 GPU Colab provides for free brings that down to minutes.

Don't see T4 GPU as an option?

Colab's free GPUs are not guaranteed and depend on availability. If T4 is not listed, try again in a few minutes or select any available GPU option.

Help me troubleshoot GPU access in Google Colab.

Install libraries and verify GPU access

Your Colab notebook starts with PyTorch pre-installed, but you need a few more libraries for this project. You'll install them with a single command, then confirm the GPU is accessible.

  • In the first code cell of your notebook, install the project libraries by running this command:
!pip install -q transformers datasets evaluate accelerate

What does this command do?

  • transformers provides the Vision Transformer model and training utilities from Hugging Face.
  • datasets makes it easy to load and process the Food-101 image dataset.
  • evaluate provides accuracy metrics for measuring model performance.
  • accelerate handles GPU optimization during training.
  • Run the cell by pressing Shift+Enter.
  • You should see installation progress messages followed by a success confirmation. This takes about 30 seconds.

Now verify that your notebook can access the GPU. This quick check confirms the T4 is connected and ready.

  • Click the + Code button to create a new code cell, then add this GPU check:
import torch

# Check if the GPU is available and print its name
print(f"GPU available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU name: {torch.cuda.get_device_name(0)}")

What does this code do?

torch.cuda.is_available() checks whether PyTorch can see a GPU. torch.cuda.get_device_name(0) prints the specific GPU model name so you can confirm it's a T4.

  • Run the cell. You should see output like:
GPU available: True
GPU name: Tesla T4

Seeing GPU available: False?

Go back to Runtime > Change runtime type and make sure a GPU is selected. After changing the runtime type, Colab restarts the session. You'll need to re-run Cell 1 (the install command) before running the GPU check again.

Help me fix my GPU not being detected in Colab.

Import all necessary modules

With the libraries installed and GPU confirmed, you can now import everything the project needs. These imports bring in the dataset loader, model classes, training utilities, and visualization tools.

  • In the same cell (Cell 2), add the following imports below the GPU check code:
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
    DefaultDataCollator,
    pipeline,
)
import evaluate
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
import matplotlib.pyplot as plt

What are these imports for?

  • AutoImageProcessor and AutoModelForImageClassification load the ViT model and its image preprocessor.
  • TrainingArguments and Trainer handle the entire training loop (batching, optimization, evaluation).
  • RandomResizedCrop, Normalize, and ToTensor prepare images for the model (resize, scale, convert to tensors).
  • matplotlib.pyplot displays images and charts inside the notebook.
  • Run Cell 2 by pressing Shift+Enter.
  • You should see the GPU info print again, followed by no errors from the imports. If all imports succeed silently, you're good to go.

Seeing a ModuleNotFoundError?

This means the install in Cell 1 didn't complete successfully. Go back to Cell 1 and re-run it. Make sure the command starts with ! (which tells Colab to run it as a terminal command, not Python code).

Help me fix a ModuleNotFoundError in my Colab notebook.

✔️ Awesome, I've got everything!

Great. Make sure your notebook is saved (Colab auto-saves, but you can press Ctrl+S to be sure).

ⓧ I'd like to double check the full code

Here's what your notebook should contain across both cells:

# Cell 1: Install Libraries
# !pip install -q transformers datasets evaluate accelerate

# Cell 2: Verify GPU and Import Libraries
import torch

print(f"GPU available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU name: {torch.cuda.get_device_name(0)}")

import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
    DefaultDataCollator,
    pipeline,
)
import evaluate
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
import matplotlib.pyplot as plt

Your environment is locked and loaded. Next up, you'll load the Food-101 dataset and explore what 101 categories of food images look like.

Load and Explore the Food-101 Dataset

Your Colab environment is set up with GPU access and all the libraries you need. Now it's time to bring in the data your model will learn from.

Every machine learning project starts with a dataset. You'll load the Food-101 dataset from Hugging Face, explore what's inside it, and visualize sample images so you know exactly what your model will be classifying.

In this step, get ready to:
  • Load a 5,000-image subset of Food-101 and split it into training and test sets.
  • Explore the dataset structure, labels, and sample images.
  • Visualize a grid of food images with their category names.
Load and split the Food-101 dataset

The full Food-101 dataset contains over 100,000 images across 101 food categories. That's too large to train quickly on a free GPU, so you'll load a manageable 5,000-image subset instead.

  • Create a new code cell in your Colab notebook and load the dataset by running this code:
# ============================================================
# Cell 3: Load and Split the Food-101 Dataset
# ============================================================
food = load_dataset("ethz/food101", split="train[:5000]")
food = food.train_test_split(test_size=0.2)

print(f"Training examples: {len(food['train'])}")
print(f"Test examples: {len(food['test'])}")

What does this code do?

load_dataset("ethz/food101", split="train[:5000]") downloads only the first 5,000 images from the Food-101 training split on Hugging Face. This keeps the download fast and ensures training finishes in about 10-15 minutes on a free T4 GPU.

train_test_split(test_size=0.2) randomly divides those 5,000 images into 4,000 for training and 1,000 for testing. The test set lets you measure how well your model performs on images it has never seen during training.

  • Run the cell by pressing Shift+Enter.
  • You should see output showing Training examples: 4000 and Test examples: 1000.

Download taking too long?

The first download may take 1-2 minutes depending on your internet connection. You'll see progress bars as the images download. If the cell appears stuck, check that your Colab runtime is still connected (green checkmark in the top-right area).

help me debug a slow dataset download in Google Colab

Explore the dataset structure

Before training, you should understand what your data looks like. Each example in the dataset has two fields: an image (the photo) and a label (an integer representing the food category).

  • Create a new code cell and explore the dataset labels by running this code:
# ============================================================
# Cell 4: Explore the Dataset
# ============================================================
labels = food["train"].features["label"].names
print(f"Number of food categories: {len(labels)}")
print(f"First 10 categories: {labels[:10]}")

example = food["train"][0]
print(f"\nExample label: {labels[example['label']]}")
example["image"]

What does this code do?

food["train"].features["label"].names extracts the human-readable category names (like "apple_pie", "baby_back_ribs", "baklava") from the dataset metadata. There are 101 categories total.

The last line, example["image"], displays the actual food photo inline in Colab. Each image is stored as a PIL Image object that Colab renders automatically.

  • Run the cell by pressing Shift+Enter.
  • You should see Number of food categories: 101 printed, followed by the first 10 category names, and a food image displayed below the cell.

Why 101 categories?

Food-101 was designed as a realistic classification challenge. With 101 visually distinct food types (from sushi to steak to tiramisu), your model needs to learn subtle visual differences between similar-looking dishes. This is much harder than binary classification (yes/no) and closer to real-world problems.

Visualize a grid of sample images

Seeing multiple examples at once helps you understand the visual diversity in your dataset.

  • Create a new code cell and visualize a grid of food images by running this code:
# ============================================================
# Cell 5: Visualize Sample Images
# ============================================================
fig, axes = plt.subplots(2, 4, figsize=(16, 8))
for i, ax in enumerate(axes.flat):
    example = food["train"][i]
    ax.imshow(example["image"])
    ax.set_title(labels[example["label"]], fontsize=10)
    ax.axis("off")
plt.tight_layout()
plt.show()

What does this code do?

plt.subplots(2, 4) creates a 2-row by 4-column grid of image slots. The loop pulls the first 8 training images and displays each one with its food category name as the title.

ax.axis("off") hides the axis numbers so you see clean images. This is a common pattern for visualizing image datasets before training.

  • Run the cell by pressing Shift+Enter.
  • You should see a 2x4 grid of food photos with category labels above each image.

Images not displaying?

Make sure you ran Cell 3 first (the cell that loads the food variable). If you see a NameError about food or labels not being defined, run the previous cells in order from top to bottom.

help me fix image display issues in my Colab notebook

✔️ Awesome, I've got everything!

Great. Double check you've run all cells in order (Cells 1 through 5) and that your grid of food images is displaying.

ⓧ I'd like to double check the full code

Here is the complete notebook code up to this point. Compare it with your cells to make sure everything matches.

# Cell 1: Install Libraries
# !pip install -q transformers datasets evaluate accelerate

# Cell 2: Verify GPU and Import Libraries
import torch

print(f"GPU available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU name: {torch.cuda.get_device_name(0)}")

import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
    DefaultDataCollator,
    pipeline,
)
import evaluate
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
import matplotlib.pyplot as plt

# Cell 3: Load and Split the Food-101 Dataset
food = load_dataset("ethz/food101", split="train[:5000]")
food = food.train_test_split(test_size=0.2)

print(f"Training examples: {len(food['train'])}")
print(f"Test examples: {len(food['test'])}")

# Cell 4: Explore the Dataset
labels = food["train"].features["label"].names
print(f"Number of food categories: {len(labels)}")
print(f"First 10 categories: {labels[:10]}")

example = food["train"][0]
print(f"\nExample label: {labels[example['label']]}")
example["image"]

# Cell 5: Visualize Sample Images
fig, axes = plt.subplots(2, 4, figsize=(16, 8))
for i, ax in enumerate(axes.flat):
    example = food["train"][i]
    ax.imshow(example["image"])
    ax.set_title(labels[example["label"]], fontsize=10)
    ax.axis("off")
plt.tight_layout()
plt.show()

You now have your data loaded and understand its structure. Next up, you'll preprocess these images into the exact format your Vision Transformer expects.

Preprocess Images for the Vision Transformer

You've loaded and explored the Food-101 dataset. Now you know what the images look like and how the 101 categories are organized.

The Vision Transformer can't work with raw photos directly. It expects every image to be the exact same size (224x224 pixels) with pixel values normalized to a specific range. In this step, you'll build the preprocessing pipeline that converts raw food photos into the format the model needs.

In this step, get ready to:
  • Load the ViT image processor and inspect its expected settings.
  • Define image transforms for training augmentation and normalization.
  • Apply the transforms to your dataset and set up the data collator for batching.
Load the ViT image processor

Every pre-trained model comes with an image processor that knows the exact image size and normalization values the model was trained with. Loading this processor ensures your images are prepared to match what the model expects.

  • Create a new code cell in your Colab notebook and add the following code to load the image processor:
# Point to the pre-trained ViT model on Hugging Face
checkpoint = "google/vit-base-patch16-224-in21k"
image_processor = AutoImageProcessor.from_pretrained(checkpoint)

# Print the processor's expected image settings
print(f"Image size: {image_processor.size}")
print(f"Image mean: {image_processor.image_mean}")
print(f"Image std: {image_processor.image_std}")

What does this code do?

`checkpoint` stores the Hugging Face model ID for Google's ViT-base model, pre-trained on ImageNet-21k (14 million images across 21,843 classes).

`AutoImageProcessor.from_pretrained()` downloads the processor configuration that knows exactly how this model wants its images formatted. The print statements reveal the target size (224x224), mean, and standard deviation values used for normalization.

  • Run the cell.
  • You should see the image size as 224x224, and the mean and std values both as [0.5, 0.5, 0.5].

Seeing a download error?

Make sure you're connected to the internet in Colab. Check the top-right corner for a green checkmark on the connection indicator. If disconnected, click **Reconnect**.

help me fix a download error when loading AutoImageProcessor from Hugging Face

Define image transforms and create preprocessing functions

Now that you know the model's requirements, you need to build a transform pipeline that resizes, converts, and normalizes each image before it reaches the model.

  • In the same cell (or a new cell directly below), add this code to define the transform pipelines:
# Build the normalization transform using the processor's values
normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
size = (image_processor.size["height"], image_processor.size["width"])

# Training transforms: random crop adds variety to help the model generalize
train_transforms = Compose([
    RandomResizedCrop(size),
    ToTensor(),
    normalize,
])

# Validation transforms: same pipeline without random augmentation differences
val_transforms = Compose([
    RandomResizedCrop(size),
    ToTensor(),
    normalize,
])

What does this code do?

`Normalize` scales pixel values using the mean and standard deviation the ViT model was trained with. This puts all pixel values into the same numerical range the model expects.

`RandomResizedCrop(size)` randomly crops a portion of the image and resizes it to 224x224. This is a form of data augmentation. Each time the model sees the same image during training, it gets a slightly different crop, which helps prevent overfitting.

`ToTensor()` converts the PIL image to a PyTorch tensor and rescales pixel values from 0-255 to 0-1. `Compose` chains these three steps into a single callable pipeline.

Next, you need wrapper functions that apply these transforms to entire batches of images from the dataset.

  • Create a new code cell and add the preprocessing functions:
# Preprocessing function for training images
def preprocess_train(examples):
    examples["pixel_values"] = [
        train_transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples

# Preprocessing function for validation/test images
def preprocess_val(examples):
    examples["pixel_values"] = [
        val_transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples

What does this code do?

Each function takes a batch of examples from the dataset. It converts each image to RGB (ensuring consistent 3-channel input), applies the transform pipeline, and stores the result in a new pixel_values column.

The original image column is deleted because the model only needs the processed pixel_values tensor. This saves memory during training.

  • Run the cell. You won't see output, but the functions are now defined and ready to use.
Apply the transforms and set up the data collator

Instead of preprocessing all 5,000 images upfront (which would use a lot of memory), you'll use on-the-fly processing. This applies the transforms only when each batch is loaded during training.

  • In the same cell, add this code below the functions to wire everything together:
# Apply transforms on-the-fly when data is accessed
food["train"].set_transform(preprocess_train)
food["test"].set_transform(preprocess_val)

# Data collator batches processed examples together for the model
data_collator = DefaultDataCollator()

What does this code do?

`set_transform()` tells the dataset to apply your preprocessing function lazily. Instead of transforming all images at once, each image is processed only when it's needed during a training step. This keeps memory usage low.

`DefaultDataCollator()` handles the mechanics of grouping individual processed examples into batches that PyTorch can feed to the model. It stacks the pixel_values tensors and labels into properly shaped batch tensors.

  • Run the cell.
  • You won't see printed output, but the transforms are now connected to your dataset. Every time a batch is drawn during training, images will be automatically preprocessed.

Getting an error about missing columns?

Make sure you ran the cell that loaded the dataset (Cell 3) earlier in this session. Colab sessions reset if you disconnect, so all previous cells need to be re-run in order.

help me debug a missing column error in my Hugging Face dataset transforms

✔️ Awesome, I've got everything!

Great. Double check you've run all cells through Cell 7 successfully before moving on.

ⓧ I'd like to double check the full code

Here's everything your notebook should contain through the end of this step:

# Cell 1: Install Libraries
# !pip install -q transformers datasets evaluate accelerate

# Cell 2: Verify GPU and Import Libraries
import torch

print(f"GPU available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU name: {torch.cuda.get_device_name(0)}")

import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
    DefaultDataCollator,
    pipeline,
)
import evaluate
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
import matplotlib.pyplot as plt

# Cell 3: Load and Split the Food-101 Dataset
food = load_dataset("ethz/food101", split="train[:5000]")
food = food.train_test_split(test_size=0.2)

print(f"Training examples: {len(food['train'])}")
print(f"Test examples: {len(food['test'])}")

# Cell 4: Explore the Dataset
labels = food["train"].features["label"].names
print(f"Number of food categories: {len(labels)}")
print(f"First 10 categories: {labels[:10]}")

example = food["train"][0]
print(f"\nExample label: {labels[example['label']]}")
example["image"]

# Cell 5: Visualize Sample Images
fig, axes = plt.subplots(2, 4, figsize=(16, 8))
for i, ax in enumerate(axes.flat):
    example = food["train"][i]
    ax.imshow(example["image"])
    ax.set_title(labels[example["label"]], fontsize=10)
    ax.axis("off")
plt.tight_layout()
plt.show()

# Cell 6: Load Image Processor and Define Transforms
checkpoint = "google/vit-base-patch16-224-in21k"
image_processor = AutoImageProcessor.from_pretrained(checkpoint)

print(f"Image size: {image_processor.size}")
print(f"Image mean: {image_processor.image_mean}")
print(f"Image std: {image_processor.image_std}")

normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
size = (image_processor.size["height"], image_processor.size["width"])

train_transforms = Compose([
    RandomResizedCrop(size),
    ToTensor(),
    normalize,
])

val_transforms = Compose([
    RandomResizedCrop(size),
    ToTensor(),
    normalize,
])

# Cell 7: Create Preprocessing Functions and Apply Transforms
def preprocess_train(examples):
    examples["pixel_values"] = [
        train_transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples

def preprocess_val(examples):
    examples["pixel_values"] = [
        val_transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples

food["train"].set_transform(preprocess_train)
food["test"].set_transform(preprocess_val)

data_collator = DefaultDataCollator()

Your preprocessing pipeline is locked in. Every food image will now be automatically resized, augmented, and normalized on-the-fly during training. Next up, you'll load the pre-trained ViT model and start fine-tuning it on your food data.

Fine-Tune the ViT Model

Your data is preprocessed and ready. Now it's time for the main event: teaching the model to recognize food.

You'll load Google's pre-trained Vision Transformer and swap its classification head for one that outputs 101 food categories. Then you'll configure the training loop and let the Trainer handle the heavy lifting of backpropagation, gradient updates, and evaluation.

In this step, get ready to:
  • Load the pre-trained ViT model with a 101-class classification head.
  • Configure training settings and create the Trainer.
  • Run training and watch accuracy improve across epochs.
Load the pre-trained model and set up metrics

The ViT model you're loading was pre-trained on ImageNet-21k (14 million images, 21,843 categories). It already knows how to extract visual features like edges, textures, and shapes. You just need to add a new classification head that maps those features to your 101 food categories.

  • Create a new code cell and paste this code to load the model:
# Set up label-to-id mappings for the model
label2id = {label: str(i) for i, label in enumerate(labels)}
id2label = {str(i): label for i, label in enumerate(labels)}

# Load the pre-trained ViT with a new 101-class head
model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes=True,
)

print(f"Model parameters: {model.num_parameters():,}")

What does this code do?

label2id and id2label are dictionaries that map between food names (like "sushi") and numeric IDs (like "0"). The model needs these to know which number corresponds to which food.

ignore_mismatched_sizes=True is critical here. The original model has a classification head for 21,843 ImageNet categories. You're replacing it with one for 101 food categories. This flag tells the library to allow that size change without throwing an error.

  • Run the cell. You should see a warning about weights not being used (this is expected since you're replacing the classification head) followed by the total parameter count, approximately 86 million.

Seeing an out-of-memory error?

Make sure you enabled the T4 GPU runtime earlier. Go to Runtime > Change runtime type and confirm T4 GPU is selected. If GPU is not available, Colab's free resources may be exhausted. Try again later or restart your runtime.

help me fix an out-of-memory error when loading my ViT model in Google Colab

  • In a new code cell, paste this code to set up the evaluation metric:
# Load accuracy metric for evaluation
accuracy = evaluate.load("accuracy")

# Define how to compute metrics during training
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

What does this code do?

evaluate.load("accuracy") loads a pre-built accuracy calculator from the Hugging Face evaluate library.

compute_metrics is a function that the Trainer calls after each evaluation round. It takes the model's raw output (logits), picks the highest-scoring class for each image using argmax, and compares those predictions against the true labels.

  • Run the cell. You won't see visible output here. The compute_metrics function is used automatically during training in the next substep.
Configure training and create the Trainer

Training a model involves dozens of settings: how fast it learns, how many times it sees the data, how often it evaluates. The TrainingArguments object bundles all of these into one configuration.

  • Create a new code cell and paste the training configuration:
# Configure all training hyperparameters
training_args = TrainingArguments(
    output_dir="food-classifier",
    remove_unused_columns=False,
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    warmup_ratio=0.1,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=False,
    report_to="none",
)

What do these settings mean?

  • learning_rate=5e-5: A small learning rate so the model makes gentle updates to its pre-trained weights.
  • per_device_train_batch_size=8 with gradient_accumulation_steps=4: Processes 8 images at a time but accumulates gradients over 4 steps, giving an effective batch size of 32 while staying within GPU memory.
  • num_train_epochs=3: The model sees all 4,000 training images three times.
  • eval_strategy="epoch": Runs evaluation on the test set after each full pass through the data.
  • load_best_model_at_end=True: After training, automatically loads the checkpoint with the highest accuracy.
  • remove_unused_columns=False: Keeps the pixel_values column that your transforms created. Without this, the Trainer would drop it.
  • In a new code cell, create the Trainer that ties everything together:
# Create the Trainer with all components
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=food["train"],
    eval_dataset=food["test"],
    processing_class=image_processor,
    compute_metrics=compute_metrics,
)

What does the Trainer do?

The Trainer is Hugging Face's high-level training API. It handles the entire training loop: feeding batches to the model, computing loss, running backpropagation, updating weights, and evaluating performance.

processing_class=image_processor tells the Trainer which processor to save alongside the model. This ensures anyone who loads your trained model later gets the correct preprocessing automatically.

  • Run both cells. You should see no errors. The Trainer is now ready to start training.

Seeing a NameError?

If you see an error like NameError: name 'data_collator' is not defined, make sure you ran all the cells from the previous steps in order. Colab cells must be executed sequentially for variables to be available.

help me fix a NameError in my Trainer setup in Google Colab

Run training and observe accuracy improving

Everything is configured. Time to start the training loop. This will take approximately 10-15 minutes on a T4 GPU.

  • Create a new code cell and run this single line to start training:
# Start the fine-tuning process
trainer.train()

What should I see during training?

A progress bar appears showing how many training steps have completed. Every 10 steps, the current loss is logged (you'll see it decreasing over time).

After each epoch (one full pass through the 4,000 images), the Trainer evaluates on your 1,000 test images and prints the accuracy. You should see accuracy increase from epoch to epoch as the model learns to distinguish food categories.

  • Wait for training to finish. You should see a table at the end showing training loss and evaluation accuracy for each epoch. The final accuracy will vary, but you should see clear improvement from epoch 1 to epoch 3.

Training is extremely slow or stuck?

If each step takes more than a few seconds, you're likely running on CPU instead of GPU. Go to Runtime > Change runtime type and select T4 GPU. You'll need to re-run all previous cells after switching.

If Colab says no GPU is available, the free tier resources may be temporarily exhausted. Try again later or reduce num_train_epochs to 1 for a faster but less accurate result.

help me fix slow training in my Google Colab notebook

✔️ Awesome, I've got everything!

Great. Make sure you let training finish completely before moving on. The best checkpoint is saved automatically.

ⓧ I'd like to double check the full code

Here is the full code for your notebook up to and including this step:

# Cell 1: Install Libraries
# !pip install -q transformers datasets evaluate accelerate

# Cell 2: Verify GPU and Import Libraries
import torch

print(f"GPU available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU name: {torch.cuda.get_device_name(0)}")

import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
    DefaultDataCollator,
    pipeline,
)
import evaluate
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
import matplotlib.pyplot as plt

# Cell 3: Load and Split the Food-101 Dataset
food = load_dataset("ethz/food101", split="train[:5000]")
food = food.train_test_split(test_size=0.2)

print(f"Training examples: {len(food['train'])}")
print(f"Test examples: {len(food['test'])}")

# Cell 4: Explore the Dataset
labels = food["train"].features["label"].names
print(f"Number of food categories: {len(labels)}")
print(f"First 10 categories: {labels[:10]}")

example = food["train"][0]
print(f"\nExample label: {labels[example['label']]}")
example["image"]

# Cell 5: Visualize Sample Images
fig, axes = plt.subplots(2, 4, figsize=(16, 8))
for i, ax in enumerate(axes.flat):
    example = food["train"][i]
    ax.imshow(example["image"])
    ax.set_title(labels[example["label"]], fontsize=10)
    ax.axis("off")
plt.tight_layout()
plt.show()

# Cell 6: Load Image Processor and Define Transforms
checkpoint = "google/vit-base-patch16-224-in21k"
image_processor = AutoImageProcessor.from_pretrained(checkpoint)

print(f"Image size: {image_processor.size}")
print(f"Image mean: {image_processor.image_mean}")
print(f"Image std: {image_processor.image_std}")

normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
size = (image_processor.size["height"], image_processor.size["width"])

train_transforms = Compose([
    RandomResizedCrop(size),
    ToTensor(),
    normalize,
])

val_transforms = Compose([
    RandomResizedCrop(size),
    ToTensor(),
    normalize,
])

# Cell 7: Create Preprocessing Functions and Apply Transforms
def preprocess_train(examples):
    examples["pixel_values"] = [
        train_transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples

def preprocess_val(examples):
    examples["pixel_values"] = [
        val_transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples

food["train"].set_transform(preprocess_train)
food["test"].set_transform(preprocess_val)

data_collator = DefaultDataCollator()

# Cell 8: Set Up the Model and Metrics
label2id = {label: str(i) for i, label in enumerate(labels)}
id2label = {str(i): label for i, label in enumerate(labels)}

model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes=True,
)

print(f"Model parameters: {model.num_parameters():,}")

accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

# Cell 9: Configure Training and Create Trainer
training_args = TrainingArguments(
    output_dir="food-classifier",
    remove_unused_columns=False,
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    warmup_ratio=0.1,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=False,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=food["train"],
    eval_dataset=food["test"],
    processing_class=image_processor,
    compute_metrics=compute_metrics,
)

# Cell 10: Train the Model
trainer.train()

Your model is trained and the best checkpoint is saved. Next up, you'll load that checkpoint and use it to classify food images you've never shown it before.

Test Your Food Classifier on New Images

Your model has trained across 3 epochs and the accuracy metric climbed with each pass. Now it is time for the real test: feeding your model food images it has never seen before.

In this step, you will build an inference pipeline that takes any food photo and returns predictions with confidence scores. You will also peek under the hood to understand exactly what happens when the model classifies an image.

In this step, get ready to:
  • Load fresh images from the Food-101 validation set and run the pipeline API.
  • Perform manual inference to understand what happens under the hood.
  • Test on multiple images and display predictions in a visual grid.
Load a food image and run inference with the pipeline API

The pipeline API from Hugging Face wraps model loading, image preprocessing, and post-processing into a single callable object. You pass in an image and get back human-readable predictions.

  • Create a new code cell in your notebook (Cell 11) and paste the following code to load validation images and run inference:
# ============================================================
# Cell 11: Inference with Pipeline API
# ============================================================
# Load 10 images from the validation split (never seen during training)
ds = load_dataset("ethz/food101", split="validation[:10]")
test_image = ds[0]["image"]
plt.imshow(test_image)
plt.axis("off")
plt.title("Test Image")
plt.show()

# Create a high-level inference pipeline pointing at your trained checkpoint
classifier = pipeline("image-classification", model="food-classifier/checkpoint-375", device=0)
predictions = classifier(test_image, top_k=5)

# Print the top 5 predictions with confidence percentages
print("\nTop 5 Predictions:")
for pred in predictions:
    print(f"  {pred['label']}: {pred['score']:.2%}")

What does this code do?

  • load_dataset("ethz/food101", split="validation[:10]") loads 10 images from the validation split. These are images your model never saw during training.
  • pipeline("image-classification", model="food-classifier/checkpoint-375", device=0) creates a high-level classifier that handles image preprocessing, model forward pass, and output formatting in one call.
  • top_k=5 tells the pipeline to return the 5 most confident predictions, each with a label name and a probability score.

Your checkpoint path may differ

The path food-classifier/checkpoint-375 depends on your exact training steps. If you get a folder-not-found error, check the food-classifier/ directory in Colab's file browser for the actual checkpoint folder name (e.g. checkpoint-375 or checkpoint-750).

  • Run the cell. You should see the test image displayed, followed by 5 food category predictions with percentage confidence scores printed below.

Not seeing predictions?

  • If you see an error about the model path, open the file browser on the left sidebar and navigate into the food-classifier folder to find your actual checkpoint name.
  • If you see a CUDA out-of-memory error, click Runtime then Restart runtime and re-run all cells from the top.

Why isn't my pipeline producing predictions?

Run inference manually to understand what happens under the hood

The pipeline is convenient, but understanding the individual steps helps you debug issues and customize behavior. In this chunk you will reproduce the same inference step-by-step.

  • Create a new code cell (Cell 12) and paste the following code to perform manual inference:
# ============================================================
# Cell 12: Manual Inference (Under the Hood)
# ============================================================
# Preprocess the image into normalized tensors and move to GPU
inputs = image_processor(test_image, return_tensors="pt").to("cuda")

# Switch model to evaluation mode (disables dropout)
model.eval()
# Disable gradient tracking (saves memory during inference)
with torch.no_grad():
    logits = model(**inputs).logits

# Convert raw logits to probabilities using softmax
probs = torch.nn.functional.softmax(logits, dim=-1)
# Extract the top 5 highest probabilities and their indices
top5_probs, top5_indices = torch.topk(probs, 5)

print("Manual Inference - Top 5 Predictions:")
for prob, idx in zip(top5_probs[0], top5_indices[0]):
    print(f"  {model.config.id2label[str(idx.item())]}: {prob.item():.2%}")

What happens at each stage?

  • image_processor(test_image, return_tensors="pt") resizes the image to 224x224, normalizes pixel values, and returns a PyTorch tensor.
  • model.eval() switches off dropout and batch normalization training behavior. This ensures consistent predictions.
  • torch.no_grad() disables gradient tracking because you are not training. This saves GPU memory.
  • softmax converts raw logits (unbounded numbers) into probabilities between 0 and 1 that sum to 1.0 across all 101 classes.
  • torch.topk(probs, 5) picks the 5 highest probability predictions from the 101 classes.
  • Run the cell. You should see the same top-5 predictions as the pipeline API, confirming that both methods produce identical results.

Seeing different results from the pipeline?

  • Small differences in confidence scores (e.g. 0.01%) are normal due to floating-point precision. The ranked order of predictions should match.
  • If predictions are wildly different, make sure you ran model.eval() before the forward pass. Without it, dropout layers remain active and produce inconsistent results.

Why are my manual inference results different from the pipeline?

Test on multiple images and display predictions in a grid

A single test image is useful, but testing across multiple images gives you a better sense of how your model performs overall.

  • Create a new code cell (Cell 13) and paste the following code to classify 6 images and display them in a grid:
# ============================================================
# Cell 13: Test on Multiple Images
# ============================================================
# Create a 2x3 grid showing predictions on 6 different images
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
for i, ax in enumerate(axes.flat):
    image = ds[i]["image"]
    # Get the single top prediction for each image
    preds = classifier(image, top_k=1)
    ax.imshow(image)
    ax.set_title(f"Predicted: {preds[0]['label']}\nConfidence: {preds[0]['score']:.1%}", fontsize=10)
    ax.axis("off")
plt.tight_layout()
plt.show()

What does this code do?

  • plt.subplots(2, 3) creates a 2-row, 3-column grid of subplots for displaying 6 images side by side.
  • classifier(image, top_k=1) returns only the single most confident prediction for each image, keeping the display clean.
  • Each subplot shows the food image with its predicted label and confidence score as the title.
  • Run the cell. You should see a 2x3 grid of food images, each labeled with the model's top prediction and confidence percentage.

Grid looks empty or images are missing?

  • Make sure Cell 11 ran successfully first. The ds variable and classifier pipeline are both defined in that cell.
  • If only some images show, your ds dataset may have fewer than 6 entries. Confirm by running len(ds) in a new cell.

Why is my matplotlib grid not displaying correctly?

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook with all cells from Cell 1 through Cell 13.

ⓧ I'd like to double check the full code

Here is the complete notebook code up to this point. Compare each cell against yours:

# ============================================================
# Cell 1: Install Libraries
# ============================================================
# !pip install -q transformers datasets evaluate accelerate

# ============================================================
# Cell 2: Verify GPU and Import Libraries
# ============================================================
import torch

print(f"GPU available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU name: {torch.cuda.get_device_name(0)}")

import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
    DefaultDataCollator,
    pipeline,
)
import evaluate
from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
import matplotlib.pyplot as plt

# ============================================================
# Cell 3: Load and Split the Food-101 Dataset
# ============================================================
food = load_dataset("ethz/food101", split="train[:5000]")
food = food.train_test_split(test_size=0.2)

print(f"Training examples: {len(food['train'])}")
print(f"Test examples: {len(food['test'])}")

# ============================================================
# Cell 4: Explore the Dataset
# ============================================================
labels = food["train"].features["label"].names
print(f"Number of food categories: {len(labels)}")
print(f"First 10 categories: {labels[:10]}")

example = food["train"][0]
print(f"\nExample label: {labels[example['label']]}")
example["image"]

# ============================================================
# Cell 5: Visualize Sample Images
# ============================================================
fig, axes = plt.subplots(2, 4, figsize=(16, 8))
for i, ax in enumerate(axes.flat):
    example = food["train"][i]
    ax.imshow(example["image"])
    ax.set_title(labels[example["label"]], fontsize=10)
    ax.axis("off")
plt.tight_layout()
plt.show()

# ============================================================
# Cell 6: Load Image Processor and Define Transforms
# ============================================================
checkpoint = "google/vit-base-patch16-224-in21k"
image_processor = AutoImageProcessor.from_pretrained(checkpoint)

print(f"Image size: {image_processor.size}")
print(f"Image mean: {image_processor.image_mean}")
print(f"Image std: {image_processor.image_std}")

normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
size = (image_processor.size["height"], image_processor.size["width"])

train_transforms = Compose([
    RandomResizedCrop(size),
    ToTensor(),
    normalize,
])

val_transforms = Compose([
    RandomResizedCrop(size),
    ToTensor(),
    normalize,
])

# ============================================================
# Cell 7: Create Preprocessing Functions and Apply Transforms
# ============================================================
def preprocess_train(examples):
    examples["pixel_values"] = [
        train_transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples

def preprocess_val(examples):
    examples["pixel_values"] = [
        val_transforms(img.convert("RGB")) for img in examples["image"]
    ]
    del examples["image"]
    return examples

food["train"].set_transform(preprocess_train)
food["test"].set_transform(preprocess_val)

data_collator = DefaultDataCollator()

# ============================================================
# Cell 8: Set Up the Model and Metrics
# ============================================================
label2id = {label: str(i) for i, label in enumerate(labels)}
id2label = {str(i): label for i, label in enumerate(labels)}

model = AutoModelForImageClassification.from_pretrained(
    checkpoint,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes=True,
)

print(f"Model parameters: {model.num_parameters():,}")

accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

# ============================================================
# Cell 9: Configure Training and Create Trainer
# ============================================================
training_args = TrainingArguments(
    output_dir="food-classifier",
    remove_unused_columns=False,
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    warmup_ratio=0.1,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=False,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=food["train"],
    eval_dataset=food["test"],
    processing_class=image_processor,
    compute_metrics=compute_metrics,
)

# ============================================================
# Cell 10: Train the Model
# ============================================================
trainer.train()

# ============================================================
# Cell 11: Inference with Pipeline API
# ============================================================
ds = load_dataset("ethz/food101", split="validation[:10]")
test_image = ds[0]["image"]
plt.imshow(test_image)
plt.axis("off")
plt.title("Test Image")
plt.show()

classifier = pipeline("image-classification", model="food-classifier/checkpoint-375", device=0)
predictions = classifier(test_image, top_k=5)

print("\nTop 5 Predictions:")
for pred in predictions:
    print(f"  {pred['label']}: {pred['score']:.2%}")

# ============================================================
# Cell 12: Manual Inference (Under the Hood)
# ============================================================
inputs = image_processor(test_image, return_tensors="pt").to("cuda")

model.eval()
with torch.no_grad():
    logits = model(**inputs).logits

probs = torch.nn.functional.softmax(logits, dim=-1)
top5_probs, top5_indices = torch.topk(probs, 5)

print("Manual Inference - Top 5 Predictions:")
for prob, idx in zip(top5_probs[0], top5_indices[0]):
    print(f"  {model.config.id2label[str(idx.item())]}: {prob.item():.2%}")

# ============================================================
# Cell 13: Test on Multiple Images
# ============================================================
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
for i, ax in enumerate(axes.flat):
    image = ds[i]["image"]
    preds = classifier(image, top_k=1)
    ax.imshow(image)
    ax.set_title(f"Predicted: {preds[0]['label']}\nConfidence: {preds[0]['score']:.1%}", fontsize=10)
    ax.axis("off")
plt.tight_layout()
plt.show()

Secret mission

Build a Visual Prediction Dashboard

Your model works on dataset images, but what about your own food photos? In this secret mission, you'll upload your own food pictures, run them through your fine-tuned classifier, and build a visual bar chart that shows the model's top-5 predictions with confidence scores.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely in Google Colab's free tier, so there are no ongoing costs. Decide whether to keep your notebook for future reference, pause the GPU runtime to free up resources, or delete everything entirely.

Resources you used:

  • Google Colab notebook (saved in Google Drive).
  • GPU runtime session (ephemeral, auto-disconnects after inactivity).
  • Model checkpoint files in Colab's temporary storage (food-classifier/ directory).
  • Uploaded food photos in Colab's temporary storage.

Keep everything running

No action needed. Choose this if you want to keep experimenting with your food classifier or try new images.

  • Your notebook stays in Google Drive and you can reopen it any time.
  • The GPU runtime will automatically disconnect after a period of inactivity (or after 12 hours maximum).
  • When the runtime disconnects, your model checkpoints and uploaded photos are deleted from Colab's temporary storage. You would need to re-run the training cells to regenerate them.

Pause - I'll come back to this later

Disconnect the GPU runtime to free up the shared resource for other users. Your notebook code stays saved in Google Drive.

  • In your Colab notebook, click Runtime in the top menu.
  • Click Disconnect and delete runtime.
  • This frees the GPU and clears temporary storage (model checkpoints, uploaded photos). Your notebook code and outputs remain saved.
  • When you return, re-enable the GPU runtime and re-run all cells from the top to restore your model.

Delete - I don't want to use this again

Remove the notebook from Google Drive and disconnect the runtime to delete everything.

  • In your Colab notebook, click Runtime in the top menu, then click Disconnect and delete runtime. This deletes all temporary storage (model checkpoints, uploaded food photos).
  • Open Google Drive at drive.google.com.
  • Find your Colab notebook in the Colab Notebooks folder.
  • Right-click the notebook file and select Move to trash.
  • Empty your Google Drive trash to permanently delete the file.

Nice Work!

Nice Work!

You just trained your first machine learning model! You took Google's pre-trained Vision Transformer and fine-tuned it to classify 101 types of food.

You've learned how to:

  • Fine-tuned a Vision Transformer with 86 million parameters on real food images from the Food-101 dataset.
  • Executed the complete ML workflow from data loading and preprocessing through training, evaluation, and inference.
  • Built an inference pipeline that takes any food photo and returns predictions with confidence scores.
  • Created a visual prediction dashboard with your own food photos and confidence bar charts.

Ready to quiz yourself?