The repository contains gpu supported dockerized code for training, inferencing and evaluating an image classification model using ensembling techniques. It uses pre-trained VGG16 and ResNet152 models for classification and combines their predictions using an ensemble model. The dataset used is a collection of images divided into four classes: berry, bird, dog, and flower. The code includes data preprocessing, model creation, training, inference for single images, evaluation metrics such as accuracy and loss, and visualization of training results.
- Install docker on your system.
- Clone the repository:
git clone https://github.com/tawhidwasik08/Dockerized-Image-Classification.git
- Move into project directory in command line
cd Dockerized-Image-Classification/
- Run
docker compose up
- Move dataset into docker_mount folder. It directory should look like this :
Dockerized-Image-Classification/docker_mount/dataset/train/berry/
- Jupyter should be running in localhost:8888. Make sure port 8888 is not occupied.
- The model architecture is defined for VGG16 classification using the
Vgg16Classifier
class. - The model is initialized with pretrained VGG16 weights and a linear layer for the final classification.
- The model summary is printed using the
summary
function. - The model's parameters except for the linear layer are set to not require gradients.
- The training loop is implemented using the
training
function. - The model is trained for the specified number of epochs using the provided dataloaders, criterion, optimizer, and learning rate
- The best model based on the lowest validation loss is saved.
- The training loss, training accuracy, validation loss, and validation accuracy are stored for visualization and evaluation.
- The
infer
function is defined to perform inference on a single image using the trained model. - A random image from the test set is selected using the
get_random_image_from_test_set
function. - The actual class and predicted class for the image are displayed.
Hope it helps anyone looking for simple dockerized machine/deep learning notebooks with gpu support.
Peace! ✌️