How to use Transfer Learning in Keras for Image Classification

B200C LEGO Image Classification Data Set on Kaggle

Training a model for image classification can be a difficult process. Asking a model to distinguish between two near-identical images is a tall order, especially if you expect to train that model from scratch.

That’s when we introduce Transfer Learning — where we leverage powerful models that have already been trained on similar tasks, but fine-tune them to excel at our specific problem.

Because of how useful and popular transfer learning is in the world of machine learning, Keras has built in tools to assist with the process. Here, we will use Keras to build an image classification model that classifies images of LEGO bricks.

Before we start, you’ll need the data we will be training with. The B200C LEGO Data Set on Kaggle contains 800,000 annotated images belonging to the 200 most popular types of LEGO bricks. Download and unzip the data set from Kaggle to begin.

Import Keras

The only library we need for this project is Keras version 2.6.0, alongside Python 3.9 . To import the sections of Keras that we will use:

Let’s break down some of these lines, and see what these imports are used for:

• ImageDataGenerator will dynamically load images from a given directory, pre-process them, and feed them into our model.
• ModelCheckpoint and TensorBoard will be used to save our model, and graph it’s progress during training.
• Dense, Input, Dropout, and GlobalAveragePooling2D are layers that will be added to the end of the model that we will be transfer learning with.
• applications contains InceptionResNetV2, the model we will be using for transfer learning.
• time will be used to help name our saved models.

Define Training Variables

• “data_dir” should be set to the directory where the B200C data set was unzipped.
• “val_split” will define what percentage of our data we would like to use for validation during training.
• “TARGET_SIZE” defines the image dimensions that all images will be resized to before being fed to the model.
• “INPUT_SHAPE” sets the expected input shape for our model. This includes our channel depth of 3, since our images are all RGB.
• “BATCH_SIZE” defines how many images to feed through the model at the same time, to reduce training time and stabilize training. Start with a small batch size like 64, and progressively increment it until your GPU ram is nearly full during training.
• “EPOCHS” sets how many times we want to iterate over the entire data set while training.

Loading Our Data

We will use the ImageDataGenerator utility from Keras to load and process our data:

As you can see, ImageDataGenerator re-scales all of our pixel values between 0 and 1.

The ImageDataGenerator.flow_from_directory() function lets us specify several attributes in addition to TARGET_SIZE and BATCH_SIZE:

• “color_mode” : this is set to “rgb” because our images are in full color
• “class_mode” : this is set to “categorical”. This tells Keras to remember the name of the folder associated with each image, as each folder name represents the class name for that image. Without this, Keras would not know how to label our data.
• “shuffle” : set to True so that ImageDataGenerator shuffles the order of images randomly, to prevent bias in training.
• “seed” : this represents the random generation seed that is used in shuffle. When seed is set, ImageDataGenerator shuffles images in way every time. This is crucial for reproducibility.
• “subset” : Earlier, we specified “val_split”. subset defines whether a section of data is for training or validation, which is the primary purpose behind having two .flow_from_directory() blocks.

Load and Define Our Model

Here, we will tell Keras to load the InceptionResNetV2 architecture, and populate it with pre-trained weights for ImageNet. After that, we append GlobalAveragePooling2D, Dropout, and Dense layers to handle the outputs from InceptionResNetV2.

• “class_count” refers to the number of different image classes we have. We use this to determine the number of output nodes in our Dense layer. This way, we have a single output node corresponding to each of our LEGO brick types.

Lastly, we compile our model with the Adam optimizer.

Define Callbacks

In Keras, callbacks are objects that perform actions during training. Here, we define a callback to record loss / accuracy data for our model, and define another callback to save our model at the end of each epoch:

Train Model

Because we used ImageDataGenerator to load our data, we will be using .fit_generator() to train our model.

If all has gone well, you should be able to run the script and begin training! Hopefully, something like this comes up:

Let this run until it has reached the level of accuracy you would like. You can change all your hyper-parameters from here and experiment with training.

Some Insights

Even with Transfer Learning, you won’t be able to reach a perfect accuracy on this data set. The B200C LEGO Classification Data Set was designed to be as realistic as possible. As part of that realism, some types of LEGO bricks, viewed at the right angles, cannot be told apart from one another.

This is not a limitation of the model, especially since we ourselves cannot differentiate in these instances. These samples represent a small portion of the full data set, but will nonetheless affect your accuracy, and that is acceptable.

Conclusion

You have successfully learned how to use Keras for transfer learning! Keras makes this easy, so that you don’t have to build image classification models from scratch. While these larger models may take longer to train, they allow you to reach a far higher level of accuracy overall.