wgpubs's picture
update code formatting
23d3539

A newer version of the Gradio SDK is available: 5.6.0

Upgrade

Objective

Inspired by Vishnu's excellent Marvel Character classifier, this model is designed to thwart adversarial attacks by DC fans who can only dream of their characters being brought into the superior Marvel universe.

Dataset

The dataset is composed of roughly 200 Marvel and 200 DC character images fetched from https://duckduckgo.com/ using the code in Jeremy Howard's 'Is it a bird? Creating a model from your own data' Kaggle notebook.

Training

With minimal modifications to our DataBlock and parameters passed to fastai's vision_learner, this model demonstrates how we can turn the multi-classification example Jeremy presented in session 1 of the 2022 fastai course into a regression task. These changes include:

  1. Creating a labeling function that returns a float, 0.0 if it is a DC character and 1.0 if it is a Marvel character
def is_marvel(img):
    return 1. if img.parent.name.lower().startswith("marvel") else 0.
  1. Updating our DataBlockto use a RegressionBlock for our targets, and then assigning our labeling function above to the get_y argument.
blocks=(ImageBlock, RegressionBlock)
  1. Updating our call to vision_learner to use a regression friendly metric like RMSE, as well as specifying a y_range to constrain our predictions to the expected range of between 0 and 1.
learn = vision_learner(dls, resnet18, metrics=rmse, y_range=(0, 1))

We'll start with a pre-trained ResNet18 model which we'll train by calling Learner.fine_tune() and a learning rate of 1e-3. This will train the classification head of the model (e.g., it will update the completely randomized weights dedicated to predicting a value) for 1 epoch, and then train all the model weights for 3 epochs. The final results of this process are included below.

Training the classification head only:

epoch train_loss valid_loss rmse time
0 0.344505 0.324276 0.569452 00:02

Training the entire model:

epoch train_loss valid_loss rmse time
0 0.303813 0.292256 0.540607 00:02
1 0.250147 0.272003 0.521539 00:01
2 0.223758 0.270610 0.520202 00:01

Examples

Example Marvel and DC character images from the dataset above are provided as examples for this demo. Feel free to upload your own Marvel, DC, and/or whatever else images to see whether you got a hero worth rooting for (or one to avoid).