title finetuning a pretrained torchvision model ayush thakurhttps twitter comayushthakur0 soumik rakshithttps twitter comsoumikrakshit96 date created 20230918 last modified 20230918 description finetuning a pretrained torch model from torchvision for image classification using keras introduction torchvisionhttps pytorch orgvisionstableindex html is a library part of the pytorchhttp pytorch org 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 pretrained backbone model from torchvision on the imagenette datasethttps github comfastaiimagenette using kerascore we will also demonstrate the compatibility of kerascore with an input system consisting of torch datasets and dataloadershttps pytorch orgtutorialsbeginnerbasicsdatatutorial html references customizing what happens in fit with pytorchhttps keras iokerasguidescustomtrainstepintorch pytorch datasets and dataloadershttps pytorch orgtutorialsbeginnerbasicsdatatutorial html transfer learning for computer vision using pytorchhttps pytorch orgtutorialsbeginnertransferlearningtutorial html finetuning a torchvision model using keras https wandb aimlcolabskerastorchreportsfinetuningatorchvisionmodelusingkerasvmlldzo1nde5nde1 setup define the hyperparameters creating the torch datasets and dataloaders in this example we would train an image classification model on the imagenette datasethttps github comfastaiimagenette imagenette is a subset of 10 easily classified classes from imagenethttps www imagenet org tench english springer cassette player chain saw church french horn garbage truck gas pump golf ball parachute fetch the imagenette dataset next we define preprocessing and augmentation transforms from torchvision for the train and validation sets finally we will use torchvision and the torch utils datahttps pytorch orgdocsstabledata html packages for creating the dataloaders for trainig and validation define the train and validation datasets define the torch dataloaders corresponding to the train and validation dataset specify the global device let us visualize a few samples from the training dataloader the image classification model we typically define a model in pytorch using torch nn moduleshttps pytorch orgdocsstablenotesmodules html which act as the building blocks of stateful computation let us define the resnet18 model from the torchvision package as a torch nn module pretrained on the imagenet1k datasethttps huggingface codatasetsimagenet1k define the pretrained resnet18 module from torchvision we set the classification head of the pretrained resnet18 module to an identity module even though keras supports pytorch as a backend it does not mean that we can nest torch modules inside a keras modelhttps keras iokerasapimodels because trainable variables inside a keras model is tracked exclusively via keras layershttps keras iokerasapilayers 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 we set the trainable resnet18 backbone to be a keras layer using torchmodulewrapper we set this to false if you want to freeze the backbone now we will build a keras functional model with the backbone layer create exponential decay learning rate scheduler compile the model define the backendagnostic wandb callbacks for kerascore save best model checkpoints train the model by calling model fit 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 we perform inference and detach the predicted probabilities from the torch computation graph with a tensor that does not require gradient computation introduction torchvision https pytorch org vision stable index html is a library part of the pytorch http pytorch org 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 https github com fastai imagenette using kerascore we will also demonstrate the compatibility of kerascore with an input system consisting of torch datasets and dataloaders https pytorch org tutorials beginner basics data_tutorial html references customizing what happens in fit with pytorch https keras io keras guides custom_train_step_in_torch pytorch datasets and dataloaders https pytorch org tutorials beginner basics data_tutorial html transfer learning for computer vision using pytorch https pytorch org tutorials beginner transfer_learning_tutorial html fine tuning a torchvision model using keras https wandb ai ml colabs keras torch reports fine tuning a torchvision model using keras vmlldzo1nde5nde1 setup define the hyperparameters creating the torch datasets and dataloaders in this example we would train an image classification model on the imagenette dataset https github com fastai imagenette imagenette is a subset of 10 easily classified classes from imagenet https www image net org tench english springer cassette player chain saw church french horn garbage truck gas pump golf ball parachute fetch the imagenette dataset 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 https pytorch org docs stable data html packages for creating the dataloaders for trainig and validation define the train and validation datasets define the torch dataloaders corresponding to the train and validation dataset specify the global device let us visualize a few samples from the training dataloader the image classification model we typically define a model in pytorch using torch nn module s https pytorch org docs stable notes modules html 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 https huggingface co datasets imagenet 1k define the pre trained resnet18 module from torchvision we set the classification head of the pre trained resnet18 module to an identity module even though keras supports pytorch as a backend it does not mean that we can nest torch modules inside a keras model https keras io keras api models because trainable variables inside a keras model is tracked exclusively via keras layers https keras io keras api 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 we set the trainable resnet18 backbone to be a keras layer using torchmodulewrapper we set this to false if you want to freeze the backbone now we will build a keras functional model with the backbone layer create exponential decay learning rate scheduler compile the model define the backend agnostic wandb callbacks for kerascore save best model checkpoints train the model by calling model fit 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 we perform inference and detach the predicted probabilities from the torch computation graph with a tensor that does not require gradient computation
import os

os.environ["KERAS_BACKEND"] = "torch"

import numpy as np
import matplotlib.pyplot as plt

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

import torchvision
from torchvision import datasets, models, transforms

import keras
from keras.layers import TorchModuleWrapper



batch_size = 32
image_size = 224
initial_learning_rate = 1e-3
num_epochs = 5




data_dir = keras.utils.get_file(
    fname="imagenette2-320.tgz",
    origin="https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz",
    extract=True,
)
data_dir = data_dir.replace(".tgz", "")



data_transforms = {
    "train": transforms.Compose(
        [
            transforms.RandomResizedCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    ),
    "val": transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    ),
}




image_datasets = {
    x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
    for x in ["train", "val"]
}


dataloaders = {
    x: torch.utils.data.DataLoader(
        image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4
    )
    for x in ["train", "val"]
}
dataset_sizes = {x: len(image_datasets[x]) for x in ["train", "val"]}
class_names = image_datasets["train"].classes


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")



plt.figure(figsize=(10, 10))
sample_images, sample_labels = next(iter(dataloaders["train"]))
sample_images = sample_images.numpy()
sample_labels = sample_labels.numpy()
for idx in range(9):
    ax = plt.subplot(3, 3, idx + 1)
    image = sample_images[idx].transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = std * image + mean
    image = np.clip(image, 0, 1)
    plt.imshow(image)
    plt.title("Ground Truth Label: " + class_names[int(sample_labels[idx])])
    plt.axis("off")






resnet_18 = models.resnet18(weights="IMAGENET1K_V1")



resnet_18.fc = nn.Identity()





backbone = TorchModuleWrapper(resnet_18)


backbone.trainable = True



inputs = keras.Input(shape=(3, image_size, image_size))
x = backbone(inputs)
x = keras.layers.Dropout(0.5)(x)
x = keras.layers.Dense(len(class_names))(x)
outputs = keras.activations.softmax(x, axis=1)
model = keras.Model(inputs, outputs, name="ResNet18_Classifier")

model.summary()


decay_steps = num_epochs * len(dataloaders["train"]) // batch_size
lr_scheduler = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=initial_learning_rate,
    decay_steps=decay_steps,
    decay_rate=0.1,
)


model.compile(
    loss="sparse_categorical_crossentropy",
    optimizer=keras.optimizers.Adam(lr_scheduler),
    metrics=["accuracy"],
)


callbacks = [
    
    keras.callbacks.ModelCheckpoint(
        filepath="model.weights.h5",
        monitor="val_loss",
        save_best_only=True,
        save_weights_only=True,
    )
]


history = model.fit(
    dataloaders["train"],
    validation_data=dataloaders["val"],
    epochs=num_epochs,
    callbacks=callbacks,
)


def plot_history(item):
    plt.plot(history.history[item], label=item)
    plt.plot(history.history["val_" + item], label="val_" + item)
    plt.xlabel("Epochs")
    plt.ylabel(item)
    plt.title("Train and Validation {} Over Epochs".format(item), fontsize=14)
    plt.legend()
    plt.grid()
    plt.show()


plot_history("loss")
plot_history("accuracy")



model.load_weights("model.weights.h5")
_, val_accuracy = model.evaluate(dataloaders["val"])
print("Best Validation Accuracy:", val_accuracy)



plt.figure(figsize=(10, 10))
sample_images, sample_labels = next(iter(dataloaders["train"]))



sample_pred_probas = model(sample_images.to("cuda")).detach()
sample_pred_logits = keras.ops.argmax(sample_pred_probas, axis=1)
sample_pred_logits = sample_pred_logits.to("cpu").numpy()

sample_images = sample_images.numpy()
sample_labels = sample_labels.numpy()

for idx in range(9):
    ax = plt.subplot(3, 3, idx + 1)
    image = sample_images[idx].transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = std * image + mean
    image = np.clip(image, 0, 1)
    plt.imshow(image)
    title = "Ground Truth Label: " + class_names[int(sample_labels[idx])]
    title += "\nPredicted Label: " + class_names[int(sample_pred_logits[idx])]
    plt.title(title)
    plt.axis("off")