PyTorch FastAI Integration

Introduction

FastAI is a high-level deep learning library built on top of PyTorch that simplifies the process of creating state-of-the-art deep learning models. Created by Jeremy Howard and Sylvain Gugger, fastai introduces a layered API that makes deep learning more accessible while still allowing access to the powerful underlying PyTorch capabilities when needed.

In this tutorial, we'll explore how fastai integrates with PyTorch and how you can leverage this combination to build sophisticated machine learning models with less code and faster experimentation cycles.

Why Use FastAI with PyTorch?

FastAI offers several advantages while preserving all the power of PyTorch:

1. Simplified API: FastAI provides high-level functions that reduce boilerplate code
2. Best Practices Built-in: Incorporates cutting-edge techniques automatically
3. Flexible Design: Access to PyTorch's low-level functions when needed
4. Domain-specific Modules: Specialized libraries for computer vision, NLP, tabular data, and more
5. Data Preprocessing: Powerful tools for preparing your datasets

Setting Up FastAI

First, let's install fastai and its dependencies:

bash

pip install fastai

To verify the installation and check that it's properly linked with PyTorch:

python

import fastai
import torch

print(f"FastAI version: {fastai.__version__}")
print(f"PyTorch version: {torch.__version__}")

Output:

FastAI version: 2.7.12
PyTorch version: 2.0.1

Basic FastAI Workflow

The typical fastai workflow follows these steps:

1. Create a DataLoaders object
2. Create a Learner model
3. Call fit or fine_tune to train the model
4. Evaluate and interpret results

Let's see an example using a simple image classification task:

python

from fastai.vision.all import *

# 1. Load and transform data
path = untar_data(URLs.PETS)
dls = ImageDataLoaders.from_name_func(
    path,
    get_image_files(path/"images"),
    valid_pct=0.2,
    seed=42,
    label_func=lambda x: x.name.split('_')[0],
    item_tfms=Resize(224),
    batch_tfms=aug_transforms()
)

# 2. Create a learner with a pre-trained model
learn = vision_learner(dls, resnet34, metrics=error_rate)

# 3. Fine-tune the model
learn.fine_tune(4)

Output:

epoch   train_loss   valid_loss   error_rate   time
0       1.171082     0.289069     0.089965     00:22
epoch   train_loss   valid_loss   error_rate   time
0       0.414144     0.226155     0.071813     00:29
1       0.288562     0.218122     0.064734     00:29
2       0.218479     0.188817     0.054715     00:29
3       0.173999     0.181131     0.051646     00:28

Understanding the FastAI-PyTorch Connection

FastAI is built directly on top of PyTorch, which means:

1. Every fastai model is fundamentally a PyTorch model
2. You can access and manipulate the underlying PyTorch components
3. FastAI components can be integrated into existing PyTorch code

Let's see how to access the PyTorch model inside a fastai Learner:

python

# Get the underlying PyTorch model
pytorch_model = learn.model

# Check what type of model it is
print(type(pytorch_model))

# Look at model architecture
print(pytorch_model)

Output:

<class 'fastai.vision.learner.xresnet34'>
Sequential(
  (0): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
    ...
  )
  (1): Sequential(
    ...
  )
  (2): Linear(in_features=512, out_features=37, bias=True)
)

Customizing PyTorch Models in FastAI

FastAI allows you to use custom PyTorch models within its ecosystem. Here's how to create a custom PyTorch model and integrate it with fastai:

python

import torch.nn as nn
import torch.nn.functional as F

# Define a custom PyTorch model
class MyCustomModel(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(32 * 56 * 56, 128)
        self.fc2 = nn.Linear(128, num_classes)
        
    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Use the custom model in fastai
custom_model = MyCustomModel(num_classes=len(dls.vocab))
learn = Learner(dls, custom_model, loss_func=nn.CrossEntropyLoss(), metrics=error_rate)

# Train the custom model
learn.fit_one_cycle(5)

Transfer Learning with FastAI and PyTorch

One of fastai's strengths is making transfer learning simple. Here's how to apply transfer learning using a pre-trained PyTorch model:

python

from torchvision import models

# Load a pre-trained model from torchvision
pretrained_model = models.resnet18(weights='IMAGENET1K_V1')

# Modify the final layer for our classification task
num_classes = len(dls.vocab)
pretrained_model.fc = nn.Linear(pretrained_model.fc.in_features, num_classes)

# Create a fastai learner with the modified PyTorch model
learn = Learner(dls, pretrained_model, loss_func=nn.CrossEntropyLoss(), metrics=error_rate)

# Set different learning rates for different layers
learn.freeze()  # Freeze all but the last layer
learn.fit_one_cycle(3)

# Unfreeze and continue training
learn.unfreeze()
learn.fit_one_cycle(3, max_lr=slice(1e-6, 1e-4))

Accessing FastAI's Preprocessing from PyTorch

You can use fastai's powerful data preprocessing capabilities even if you want to work directly with PyTorch for model training:

python

# Create fastai DataLoaders
dls = ImageDataLoaders.from_folder(
    path/"images", 
    valid_pct=0.2,
    item_tfms=Resize(224),
    batch_tfms=aug_transforms()
)

# Extract PyTorch DataLoaders
pytorch_train_dl = dls.train.to_pytorch()
pytorch_valid_dl = dls.valid.to_pytorch()

# Define a PyTorch training loop
def train_pytorch_model(model, train_dl, valid_dl, epochs=3):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    
    for epoch in range(epochs):
        # Training phase
        model.train()
        train_losses = []
        for x, y in train_dl:
            x, y = x.to('cuda'), y.to('cuda')
            optimizer.zero_grad()
            output = model(x)
            loss = loss_fn(output, y)
            loss.backward()
            optimizer.step()
            train_losses.append(loss.item())
        
        # Validation phase
        model.eval()
        valid_losses = []
        correct = 0
        total = 0
        with torch.no_grad():
            for x, y in valid_dl:
                x, y = x.to('cuda'), y.to('cuda')
                output = model(x)
                loss = loss_fn(output, y)
                valid_losses.append(loss.item())
                _, predicted = torch.max(output.data, 1)
                total += y.size(0)
                correct += (predicted == y).sum().item()
        
        print(f"Epoch {epoch+1}/{epochs}, Train Loss: {sum(train_losses)/len(train_losses):.4f}, "
              f"Valid Loss: {sum(valid_losses)/len(valid_losses):.4f}, "
              f"Accuracy: {100 * correct / total:.2f}%")

# Load a PyTorch model and train it
pytorch_model = models.resnet18(weights='IMAGENET1K_V1')
pytorch_model.fc = nn.Linear(pytorch_model.fc.in_features, len(dls.vocab))
pytorch_model = pytorch_model.to('cuda')

train_pytorch_model(pytorch_model, pytorch_train_dl, pytorch_valid_dl)

Real-World Example: Image Classification with FastAI and PyTorch

Let's build a practical image classification model to categorize different types of furniture using fastai and PyTorch. We'll use transfer learning and demonstrate how to export the final model for production use.

python

from fastai.vision.all import *
import torch
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path

# Assume we have a dataset with furniture images in folders by category
# Download a sample dataset (using CIFAR-10 subset for demonstration)
path = untar_data(URLs.CIFAR_100)

# Let's focus only on furniture-related classes
furniture_classes = ['bed', 'chair', 'table', 'wardrobe', 'sofa']
furniture_path = Path('furniture_dataset')
furniture_path.mkdir(exist_ok=True)

for cls in furniture_classes:
    (furniture_path/cls).mkdir(exist_ok=True)
    
    # Copy images (simulated for this example)
    # In real scenarios, you would have actual furniture images

# Create DataLoaders with data augmentation
furniture_data = DataBlock(
    blocks=(ImageBlock, CategoryBlock),
    get_items=get_image_files,
    splitter=RandomSplitter(valid_pct=0.2, seed=42),
    get_y=parent_label,
    item_tfms=[Resize(224)],
    batch_tfms=[*aug_transforms(size=224, max_warp=0), Normalize.from_stats(*imagenet_stats)]
)

dls = furniture_data.dataloaders(furniture_path)

# Show a batch of images with labels
dls.show_batch(max_n=4, figsize=(10, 8))

# Create a learner with a pre-trained model
learn = vision_learner(dls, models.resnet34, metrics=[error_rate, accuracy])

# Find optimal learning rate
learn.lr_find()

# Train the model
learn.fine_tune(5, base_lr=1e-3, freeze_epochs=3)

# Confusion matrix
interp = ClassificationInterpretation.from_learner(learn)
interp.plot_confusion_matrix()

# Export the model for production
learn.export('furniture_classifier.pkl')

# Load the model in PyTorch-friendly format
learn = load_learner('furniture_classifier.pkl')

# Convert to pure PyTorch for deployment
pytorch_model = learn.model.eval()

# Example prediction with PyTorch
img_path = furniture_path/'chair'/'some_chair_image.jpg'
img = PILImage.create(img_path)
tensor_img = torch.tensor(fastai.vision.core.image2tensor(img)).unsqueeze(0)

with torch.no_grad():
    prediction = pytorch_model(tensor_img)
    predicted_class = dls.vocab[prediction.argmax().item()]

print(f"Predicted class: {predicted_class}")

Summary

FastAI and PyTorch form a powerful combination: PyTorch provides the flexible, dynamic computational framework, while fastai adds high-level abstractions and best practices to accelerate deep learning development. Key takeaways from this tutorial include:

1. FastAI simplifies the PyTorch workflow without sacrificing flexibility
2. All fastai models are PyTorch models at their core
3. You can easily mix and match fastai and PyTorch components
4. FastAI's data preprocessing tools work well even in pure PyTorch workflows
5. Transfer learning is exceptionally streamlined with fastai

By understanding how these libraries integrate, you can choose the right level of abstraction for each part of your deep learning project, using high-level fastai functions for standard tasks and dropping down to PyTorch when you need custom functionality.

Additional Resources

• FastAI Documentation
• FastAI Book
• PyTorch Documentation
• FastAI Course

Exercises

1. Modify the furniture classifier to include additional categories
2. Create a custom loss function in PyTorch and use it with a fastai Learner
3. Build an image segmentation model using fastai's UNet implementation and export it for use in a PyTorch pipeline
4. Implement a feature extraction pipeline using a pre-trained fastai model and custom PyTorch classifier
5. Create a mixed training routine that uses fastai for initial training and PyTorch for fine-tuning