Skip to main content

Insights on Training and Prediction

We use fastai to train majority of the CCTV AI models. The fastai framework was chosen as it a gave a nice wrapper around pytorch and allowed us to use best practices with the training routines.

info

fastai is fully compatible with pytorch and you can use pretrained models, loss functions, transforms and dataloaders directly with fastai. In most cases, you can use the classes directly out of the box from pytorch but in rare cases, you will need to wrap the pytorch classes with some fastai so that they can work.

Warnings in running script_runner for compute_msi

warning WandbCallback was not able to get prediction samples -> Object of type MultiCategory is not JSON serializable

Data loading (Use of dataloaders)

  1. Dataloaders are used to load the data into the model.
  2. They handle the batching, shuffling, and other data related tasks.
  3. Using the dataloaders class, training and validation and test dataloaders are created.
  4. It is in the dataloader that you specify the image transforms (augmentation) that you will apply during training.

Image Transforms (augmentation)

danger

When we apply an image transform, we are not changing the original image, we are creating a new image from the original image at run-time. Hence, making the transforms determinstic is required. This is done by setting a seed.

info

Consider a case of 1000 images and you have transforms that randomly crop, rotate and change the contrast of the image. When these images will go into training, your dataloader will only contain 1000 images and will not contain additional transformed images. Transforms are applied in-place and at-runtime.

You can directly choose from the available transforms in fastai and pytorch and use them as you wish. You can also create your own transforms.

Choosing model

Please refer to the document for Backbones

Loss function

Please refer to the document for Loss functions

Learning rate

Finding the ideal learning rate for training a neural network can be tricky. Hence, we use the lr_find function to find the ideal learning rate.

LR Finder

fastai natively has a function to find the ideal learning rate for training a model. You can use it as follows:

ideal_lr = learn.lr_find()
print(ideal_lr[0])

This will plot the learning rate vs loss and you can choose the learning rate that you want to use.

LR Scheduler

During training, we usually don't train the entire network with the same learning rate. We use lr_scheduler to change the learning rate during training.

Refer the training section for more information.

You can also use different LR schedulers from pytorch which require little to no modification.

Layer freezing and unfreezing

The default idea with the training routine of fastai is that, they wish to leverage transfer learning. The idea is start training from pretrained weights (weights obtained by training on the Imagenet dataset) and then freeze the body and then tune the head.

This workflow can be customized by using the freeze and unfreeze methods.

For more information, refer the following links,

  1. https://forums.fast.ai/t/fine-tune-vs-fit-one-cycle-unfreeze-fit-one-cycle/66704
  2. https://forums.fast.ai/c/part1-v3/20
  3. https://forums.fast.ai/t/fine-tune-vs-fit-one-cycle/66029

Actual training

This is done using various methods in fastai

  1. fit method
  2. fit_one_cycle calls the fit method with OneCycle policy
  3. fine_tune calls fit_one_cycle internally with layer freeze and unfreeze
  4. fit_flat_cos calls fit method with Cosine Annealing

You can try these different methods to see which ones give better results.

Saving/exporting the model

fastai supports exporting the models both as pkl files and pth files.

  1. Learn.export method exports the model as a pkl file
  2. Learn.save method exports the model as a pth file

For our use case, we generally use the export method as it saves the weights and creates empty dataloaders which can be used directly in production.

The save method exports the pth file which only contain the weights, so to use them, you need to create the dataloader again.

Load model

You can load the model using the load_learner method.

learn = load_learner('path-to-the-pkl-file')
caution

A pkl file also runs code behind the scenes, it is generally not safe to load a pkl file from an untrusted source. It can run malicious code

Inference

You can use the predict method to get the predictions.

learn.predict('single-img') # to predict on a single image

test_path = Path('path-to-test-folder')
test_files = get_image_files(test_path)
dl = dls.test_dl(test_files) # to predict on a batch of images
learn.get_preds(dl=dl) # to handle batch prediction