Building End-to-End Computer Vision Applications with Fastai

ByJordan Brown

Computer vision has seen remarkable progress in recent years, with deep learning models achieving superhuman performance on tasks like image classification, object detection, and semantic segmentation. Thanks to the availability of powerful open source deep learning libraries, it‘s now possible for developers to quickly build and deploy state-of-the-art computer vision applications without being an expert in neural networks.

In this blog post, we‘ll walk through how to build an end-to-end image classification application using fastai, a cutting-edge deep learning library. Fastai provides high-level APIs for quickly training models using best practices, while still providing low-level access to customize models as needed. By the end of this post, you‘ll be able to train a highly accurate classifier and deploy it in an application.

The Fastai Deep Learning Library

Fastai is a deep learning library built on top of PyTorch that allows developers to train state-of-the-art models with minimal code. Some key features of fastai include:

  • High-level APIs for common tasks like image classification, text classification, tabular modeling, and recommendation systems
  • Easy access to best practices like learning rate finding, discriminative layer training, mixed precision training, and progressive resizing
  • Customizable models and low-level access to PyTorch for fine-grained control
  • Rich documentation and tutorials for learning
  • Interoperability with Hugging Face transformers, Weights & Biases, and other popular libraries

Fastai adopts a "layered" approach where high-level APIs for training models are composed of modular, reusable, and customizable building blocks. The core APIs are:

  1. DataLoaders – prepare your data for model training by applying transforms and creating batches
  2. Learner – handles model training and inference by putting together an optimizer, a model, and data
  3. Metrics – evaluate your model performance during training and inference
  4. Callbacks – inject custom functionality into the training loop to monitor, debug, or tweak model training

With just these core components, you can develop a wide range of deep learning applications across different domains and tasks. Let‘s see how to use them to build an image classifier.

Prepare the Data

The first step in any machine learning project is to gather, clean, and prepare your data. For our example, we‘ll use the Oxford-IIIT Pet Dataset which has ~7,000 images of cats and dogs annotated with breed labels.

We can download this dataset directly using fastai‘s built-in dataset collection:

from fastai.vision.all import *

path = untar_data(URLs.PETS)

This downloads the data, extracts it, and returns a Path object with the directory containing the images.

Next, we need to tell fastai how to interpret the directory structure to create a DataLoaders object. We use fastai‘s DataBlock API to define the inputs (images) and outputs (labels):

pets = DataBlock(
    blocks=(ImageBlock, CategoryBlock), 
    get_items=get_image_files, 
    splitter=RandomSplitter(seed=42),
    get_y=using_attr(RegexLabeller(r‘(.+)_\d+.jpg$‘), ‘name‘),
    item_tfms=[Resize(460), ToTensor],
    batch_tfms=aug_transforms(size=224, min_scale=0.75)
)
dls = pets.dataloaders(path/"images")

There‘s a lot going on here, so let‘s break it down:

  • blocks specifies the type of input (images) and output (categories)
  • get_items points to the input image files
  • splitter defines how to split the data into training and validation sets
  • get_y extracts the labels from the filenames using a regular expression
  • item_tfms resizes the images to a standard size and converts them to PyTorch tensors
  • batch_tfms performs data augmentation on batches of images during training

The DataBlock defines a "recipe" for creating a DataLoaders object from a file path. The dataloaders method creates a training and validation set according to the block definitions.

With our data prepared, we‘re ready to train a model! But first, let‘s take a peek at our data:

dls.show_batch(max_n=9, figsize=(7,8))

This shows a batch of images along with their labels:

Sample of cat and dog images

Find the Learning Rate

An important step before training any deep learning model is to find a good learning rate. The learning rate determines the size of the steps the optimizer takes to update the model weights. If the learning rate is too low, training will progress very slowly. If it‘s too high, the loss may fluctuate or even diverge.

Fastai provides a handy learning rate finder utility that trains the model for a few iterations while increasing the learning rate exponentially. By plotting the loss vs. learning rate, we can identify a good learning rate as the point with the steepest slope (strongest decrease in loss).

To use the learning rate finder, we first need to create a Learner object that combines our model, data, optimizer, and loss function:

learn = vision_learner(dls, resnet50, metrics=error_rate)
lr_min,lr_steep = learn.lr_find()
print(f"Suggested learning rates:\nMin numerically stable: {lr_min}\nSteepest point: {lr_steep}")

The vision_learner function is a shortcut that creates a Learner with an image classifier head on top of a pre-trained model, in this case a ResNet50. We pass in the DataLoaders object, the metrics to compute during training, and save the returned Learner object for later.

When we call lr_find, fastai trains the model for a few iterations while exponentially increasing the learning rate from 1e-7 to 10. We can plot the losses with learn.recorder.plot() to visualize:

Learning rate finder plot showing loss vs. learning rate

The learning rate finder suggests two learning rates: the numerically minimum stable one (1e-5) and the "steepest" one with the strongest rate of decrease in the loss (1e-3). We‘ll use the latter since it will help the model converge faster.

Train the Model

Now we‘re ready to train our model! Since we‘re using a pre-trained model, we‘ll use transfer learning to quickly adapt the model to our dataset.

The basic idea is to:

  1. Freeze all the layers except the last classifier layer
  2. Train the classifier head for a few epochs with a high learning rate to quickly converge
  3. Unfreeze all layers and train the entire network at a lower learning rate to fine-tune

We can do this with fastai‘s fine_tune method:

learn.fine_tune(epochs=6, base_lr=1e-3, freeze_epochs=3)

This freezes the model for 3 epochs while training the classifier head at a learning rate of 1e-3, then unfreezes the model and trains all layers for 3 more epochs at a lower learning rate.

On a GPU, this fine-tuning process only takes a few minutes. Here‘s the output:

epoch   train_loss  valid_loss  error_rate  time
0   0.663807    0.185389    0.062622    00:05
1   0.458638    0.152183    0.047479    00:06
2   0.336021    0.128474    0.042607    00:06
3   0.219740    0.082140    0.015385    00:06
4   0.184726    0.075628    0.013746    00:08
5   0.175744    0.068064    0.017024    00:09

After just 6 epochs of training, our model is achieving 98.3% accuracy on the validation set! Let‘s see how well it does on some sample images:

learn.show_results()

Model predictions on sample images

Looking good! The model correctly classifies the breeds of cats and dogs in the sample images.

Evaluate the Model

To get a more comprehensive view of our model‘s performance, we can use fastai‘s ClassificationInterpretation class to generate a confusion matrix:

interp = ClassificationInterpretation.from_learner(learn)
interp.plot_confusion_matrix(figsize=(12,12), dpi=60)

Confusion matrix showing model predictions vs. actual labels

The confusion matrix shows that the model performs well across all breeds, with only a few mistakes. The most confused breeds appear to be the Birman and Bombay cats and the Miniature Poodle and Toy Poodle dogs.

We can also plot the most wrong predictions with interp.plot_top_losses() to debug further:

Top prediction losses

Put the Model in Production

Now that we have a trained model, how can we use it in an application? Fastai provides an export method that saves the model as a serialized file:

learn.export(‘model.pkl‘)

This saves a pickled version of the Learner object that includes the model architecture and trained parameters. We can load this in a separate process with:

learn = load_learner(‘model.pkl‘)
img = PILImage.create(‘my-image.jpg‘)
pred,pred_idx,probs = learn.predict(img)
print(f"Prediction: {pred}; Probability: {probs[pred_idx]:.4f}")

This loads the Learner, creates an image from a file path, passes it through the model, and returns the predicted class, class index, and probabilities.

We can even deploy our trained model to a cloud endpoint on services like AWS Lambda, Google Cloud Functions, or Hugging Face Spaces without having to manage any infrastructure. Check out the FastAPI and Gradio tutorials to learn more.

Next Steps

In this post, we walked through the key steps of building an end-to-end image classification model with fastai:

  1. Preparing the data with DataBlock and DataLoaders
  2. Finding a good learning rate with the lr_find utility
  3. Training the model using transfer learning and fine-tuning
  4. Evaluating performance with a confusion matrix and top losses
  5. Exporting the model for production use

Of course, this just scratches the surface of what‘s possible with computer vision and deep learning. Here are some ideas to explore further:

  • Try different architectures like EfficientNet, Vision Transformers, or ConvNeXt
  • Experiment with more advanced augmentation techniques like MixUp or CutMix
  • Leverage self-supervised learning to train on unlabeled data with techniques like DINO or BEiT
  • Apply your model to tasks like multi-label classification, object detection, or segmentation
  • Explore model distillation to compress the model for edge devices
  • Deploy your model behind an API, build a demo app, or integrate into a product

I‘m really excited by the rapid progress in computer vision and how accessible state-of-the-art models have become thanks to libraries like fastai. Go build some cutting-edge applications!

Resources

How useful was this post?

Click on a star to rate it!

Average rating 0 / 5. Vote count: 0

No votes so far! Be the first to rate this post.

You May Like to Read,

Similar Posts