🔥 Fine-tune a TorchVision Model with Keras and WandB 🦄

Introduction

TorchVision is a library part of the PyTorch project that consists of popular datasets, model architectures, and common image transformations for computer vision. This example demonstrates how we can perform transfer learning for image classification using a pre-trained backbone model from TorchVision on the Imagenette dataset using KerasCore. We will also demonstrate the compatibility of KerasCore with an input system consisting of Torch Datasets and Dataloaders.

References:

• Customizing what happens in fit() with PyTorch
• PyTorch Datasets and Dataloaders
• Transfer learning for Computer Vision using PyTorch

Setup

• We install the main branch of KerasCore, this lets us use the latest feature merged in KerasCore.
• We also install wandb-addons, a library that hosts the backend-agnostic callbacks compatible with KerasCore

Define the Hyperparameters

Creating the Torch Datasets and Dataloaders

In this example, we would train an image classification model on the Imagenette dataset. Imagenette is a subset of 10 easily classified classes from Imagenet (tench, English springer, cassette player, chain saw, church, French horn, garbage truck, gas pump, golf ball, parachute).

Next, we define pre-processing and augmentation transforms from TorchVision for the train and validation sets.

Finally, we will use TorchVision and the torch.utils.data packages for creating the dataloaders for trainig and validation.

Let us visualize a few samples from the training dataloader.

The Image Classification Model

We typically define a model in PyTorch using torch.nn.Modules which act as the building blocks of stateful computation. Let us define the ResNet18 model from the TorchVision package as a torch.nn.Module pre-trained on the Imagenet1K dataset.

ven though Keras supports PyTorch as a backend, it does not mean that we can nest torch modules inside a keras_core.Model, because trainable variables inside a Keras Model is tracked exclusively via Keras Layers.

KerasCore provides us with a feature called TorchModuleWrapper which enables us to do exactly this. The TorchModuleWrapper is a Keras Layer that accepts a torch module and tracks its trainable variables, essentially converting the torch module into a Keras Layer. This enables us to put any torch modules inside a Keras Model and train them with a single model.fit()!

Now, we will build a Keras functional model with the backbone layer.

Evaluation and Inference

Now, we let us load the best model weights checkpoint and evaluate the model.

Finally, let us visualize the some predictions of the model