AI Image Classification with Hyperparameter Optimization | by Brent Larzalere | Towards Data Science

The ability to use artificial intelligence to identify objects and classify images has been one of the fastest growing areas in AI since Alex Krizhevsky et al. developed the AlexNet convolutional neural network (CNN) that won the 2012 ImageNet challenge. The AlexNet network achieved state-of-the-art performance labeling images from the 14,197,122 ImageNet dataset. Since then, the quantity and diversity of CNN architectures has expanded dramatically across deep CNN architectures to inception and residual CNN networks. Today, CNN’s are being used for image classification, medical imaging, facial recognition and identification, sorting your mail, robot vision and coming soon, self-driving cars.
This article will present a code-based example for performing image classification using an AI neural network. While there are many excellent image classification articles on the web, the two key items presented in this article are the use of transfer learning for jump starting your neural network model and a code framework that leverages TensorBoard to assist in optimizing your model and training hyperparameters. For this example, PyTorch will be used as the deep learning framework and Google Colab for the training infrastructure.
This example’s dataset contains four classes of flower images: tulip, rose, daisy, and sunflower. The image dataset can be downloaded here.
The first task is to review the images per class. As shown below, this is an imbalanced classification dataset.
It’s considered best practice to attempt to balance a classification dataset. The below Python code uses an over-sampling approach to balance the dataset at 600 images per class.
The dataset is then split into training, validation and test sets per class.
Ideally, a deep learning model would be trained using a large quantity of training data. Given the small sample size of available data, I will leverage transfer learning to jump start the neural network model’s training effort. Transfer learning is utilizing a neural network model that has been successfully trained on a different, but similar, task and leveraging that pre-trained model as the starting point for your model.
The pre-trained VGG-16 model is used as the transfer learning model. Since this model was trained on the ImageNet dataset, the means and standard deviations are normalized, per channel, using to the ImageNet dataset normalization values. The image size is cropped to 224 x 224.
Image augmentation will also be used to enhance the training effort. A set of image transformations will be randomly applied to the images as they are fed to the neural network during each training run. This exposes the neural network to a wider variation of images during the training process. Image augmentation is not applied to the validation and test datasets.
The approach used for the transfer learning model is to keep the core model intact, so that we can leverage its prior feature map training, and replace the image classifier layers with the output layers for our model.
The following function is used to load the transfer learning model, freeze the model weights from further training and replace the original classifier layers with new layers for the revised model.
During training, only the newly defined neural network layer weights are updated; the frozen (aka pre-trained) weights remain static. The final output layer is a LogSoftmax layer because a NLL (negative log likelihood) loss function is used during training.
To explore various model and hyperparameter settings during the model training process, the model parameters are defined in a Python ordered dictionary.
# Training Hyperparametersparams = OrderedDict(    hidden_units = [256, 512],    dropout = [0.4, 0.5],    num_classes = [4],    lr = [0.001],    batch_size = [25],    n_epochs = [40])
A Python class is defined to generate a list of run definitions across the various parameter permutations.
A RunManager class is created to define a set of methods for managing the run and epoch parameter initializations, to track the training outputs and feed the output data to TensorBoard.
The training loop is defined as a function. The function instantiates an instance of the RunManager class and cycles through a run for each RunBuilder generated list of parameters. The function creates the CNN model using the run parameters and defines the optimizer and data loaders. Note that the RunBuilder class does not need to be instantiated because it contains a single static method.
The function then executes the PyTorch training loop per epoch. PyTorch gradient calculation is turned on for training, but turned off during the processing of the validation data.
Throughout the training process, the loss, accuracy and weight information is tracked by the RunManager methods and sent to TensorBoard. Additionally, the results per run and epoch are saved to an output file.
Training is initiated using the following command; where ‘Run_Results’ is the output file name.
history = execution_loop('Run_Results')
Once the training runs are complete, the saved output file provides tabular data for analysis.
The output file can be sorted to find the best parameter combinations.
The top performing combination was run 1 with a training accuracy of 98.9% and a validation accuracy of 94%. The run parameters were:
256 hidden units
Dropout rate of 0.4
Batch size of 25
Learning rate of 0.001
Here are some samples of the output data sent to TensorBoard during the training runs.
The model graph can be interactively drilled down for confirmation of your neural network model design.
Plots of the training and validation accuracy and loss curves per run. The plots can be toggled to display all or only a select set of runs.
The model weight and bias distributions and histograms for each layer per run are also available for review.
Using transfer learning models can jump start your neural network model creation and greatly enhance your model performance, even when training on datasets that are smaller than ideal.
While this example only included a two pairs of parameters to optimize (hidden units and dropout), it is extensible to any number of parameters / hyperparameters; you can vary: the learning rate, batch size, sample shuffling during training, optimizer selection, loss function selection, setting multiple hyperparameter values within various layers, toggling layers on/off within the model, etc.
It should also be noted that this parameter-based training framework is not specific to PyTorch. It can be used, with a few coding tweaks to the training loop, for pure TensorFlow or Keras/TensorFlow models. A big shout out to the folks at Deeplizard for coming up with RunBuilder / RunManager concept.
I hope you enjoyed this article and found it informative. You can find all of the code for this example in my GitHub repo.