Transfer Learning Neural Networks

Let’s bring things back to MobileNet for a moment. The MobileNet model was trained to identify features that distinguish a thousand items from each other. That means there are convolutions to detect fur, metal, round things, ears, and all kinds of crucial differential features. All these features are chewed up and simplified before they are flattened into a neural network, where the combination of various features creates a classification.

The MobileNet model can identify different breeds of dogs, and even distinguish a Maltese terrier from a Tibetan terrier. If you were to make a “dog or cat” classifier, it makes sense that a majority of those advanced features would be reusable in your simpler model.

The previously learned convolutional filters would be extremely useful in identifying key features for brand-new classifications, like our capybara example in Figure 11-2. The trick is to take the feature identification portion of the model and apply your own neural network to the convolutional output, as illustrated in Figure 11-3.

Figure 11-3. CNN transfer learning

So how do you separate and recombine these sections of previously trained models? You’ve got lots of options. Again, we’ll learn a bit more about Graph and Layers models.

TensorFlow Hub Check, Mate!

We’re going to use transfer learning with MobileNet to identify chess pieces like the one shown in Figure 11-5.

Figure 11-5. Simple chess pieces classifier

You’ll only have a few images of each chess piece. That’s not normally enough, but with the magic of transfer learning, you’ll get an efficient model.

console.log("Loading huge CSV - this will take a while");
const numImages = 130; // between 1 and 150
// Get Y values
const labels = await dfd.read_csv("chess_labels.csv", numImages); 
const Y = labels.tensor; 
// Get X values (Chess images)
const chessImages = await dfd.read_csv("chess_images.csv", numImages);
const chessTensor = chessImages.tensor.reshape([
  labels.shape[0], 224, 224, 3, 
]);
console.log("Finished loading CSVs", chessTensor.shape, Y.shape);

Finished loading CSVs (4) [130, 224, 224, 3] (2) [130, 6]

// Load feature model
const tfhubURL =
  "https://oreil.ly/P2t2k";
const featureModel = await tf.loadGraphModel(tfhubURL, {
  fromTFHub: true,
});
const featureX = featureModel.predict(chessTensor);
// Push data through feature detection
console.log(`Features stack ${featureX.shape}`);

Features stack 130,1664

// Create NN
const transferModel = tf.sequential({
  layers: [                              
    tf.layers.dense({
      inputShape: [featureX.shape[1]],   
      units: 64,
      activation: "relu",
    }),
    tf.layers.dense({ units: 6, activation: "softmax" }),
  ],
});

Training results

Nothing is new in the training code. The model trains based on the feature output. Because the feature output is so small compared to the original image tensor, the training happens extremely quickly.

transferModel.compile({
  optimizer: "adam",
  loss: "categoricalCrossentropy",
  metrics: ["accuracy"],
});

await transferModel.fit(featureX, Y, {
  validationSplit: 0.2,
  epochs: 20,
  callbacks: { onEpochEnd: console.log },
});

Within a few epochs, the model has outstanding accuracy. Take a look at Figure 11-6.

Figure 11-6. From 50% to 96% validation accuracy in 20 epochs

Transfer learning using an existing model on TensorFlow Hub relieves you of architectural headaches and rewards you with high accuracy. But it’s not the only way you can implement transfer learning.

Shaving Layers on MobileNet

For this exercise, you will use a version of MobileNet v1.0 that is set up for being a Layers model. This is the model Teachable Machine uses, and while it’s sufficient for small exploratory projects, you’ll notice it’s not as accurate as the MobileNet v2 with 1.30 depth. You’re already well versed in converting models with the wizard, as you learned in Chapter 7, so you can create a larger, newer Layers model when needed. Accuracy is an important metric, but it’s far from the only metric you should evaluate when shopping for a transfer model.

MobileNet has a vast collection of layers, and some of these are layers you’ve never seen before. Let’s take a look. Load the MobileNet model associated with this chapter and review the summary of layers with model.summary(). This prints a huge list of layers. Don’t feel overwhelmed. When you read from the bottom to the top, the last two convolutional layers with activations are called conv_preds and conv_pw_13_relu:

...

conv_pw_13 (Conv2D)          [null,7,7,256]            65536
_________________________________________________________________
conv_pw_13_bn (BatchNormaliz [null,7,7,256]            1024
_________________________________________________________________
conv_pw_13_relu (Activation) [null,7,7,256]            0
_________________________________________________________________
global_average_pooling2d_1 ( [null,256]                0
_________________________________________________________________
 reshape_1 (Reshape)          [null,1,1,256]            0
_________________________________________________________________
dropout (Dropout)            [null,1,1,256]            0
_________________________________________________________________
conv_preds (Conv2D)          [null,1,1,1000]           257000
_________________________________________________________________
act_softmax (Activation)     [null,1,1,1000]           0
_________________________________________________________________
reshape_2 (Reshape)          [null,1000]               0
=================================================================
Total params: 475544
Trainable params: 470072
Non-trainable params: 5472

The last convolution, conv_preds, serves as a flatten layer of the features to the 1,000 possible classes. This is somewhat specific to the model’s trained classes, so because of that, we’ll jump up to the second convolution (conv_pw_13_relu) and cut there.

MobileNet is a complex model, and even though you don’t have to understand all the layers to use it for transfer learning, there’s a bit of art in deciding what to remove. In simpler models, like the one for the upcoming Chapter Challenge, it’s common to keep the entire convolutional workflow and cut at the flatten layer.

You can cut to a layer by knowing its unique name. The code shown in Example 11-4 is available on GitHub.

const featureModel = await tf.loadLayersModel('mobilenet/model.json')
console.log('ORIGINAL MODEL')
featureModel.summary()
const lastLayer = featureModel.getLayer('conv_pw_13_relu')
const shavedModel = tf.model({
  inputs: featureModel.inputs,
  outputs: lastLayer.output,
})
console.log('SHAVED DOWN MODEL')
shavedModel.summary()

The code from Example 11-4 prints out two large models, but the key difference is that the second model suddenly stops at conv_pw_13_relu.

The last layer is now the one we identified. When you review the summary of the shaved-down model, it’s like a feature extractor. There is a key difference that should be noted. The final layer is a convolution, so the first layer of your constructed transfer model should flatten the convolutional input so it can be densely connected to a neural network.

Layers Feature Model

Now you can use the shaved model as a features model. This gets you the same two-model system you had from TFHub. Your second model will need to read the output of conv_pw_13_relu:

// Create NN
const transferModel = tf.sequential({
  layers: [
    tf.layers.flatten({ inputShape: featureX.shape.slice(1) }),
    tf.layers.dense({ units: 64, activation: 'relu' }),
    tf.layers.dense({ units: 6, activation: 'softmax' }),
  ],
})

We are setting the shape as defined by the intermediate features. This could also be directly tied to the shaved model’s output shape (shavedModel.outputs[0].shape.slice(1)).

From here, you’re right back to where you were in the TFHub model. The base model creates features, and the second model interprets those features.

Training with these two layers achieves around 80%+ accuracy. Keep in mind we’re using a completely different model architecture (this is MobileNet v1) and a lower depth multiplier. Getting at least 80% from this rough model is good.

A Unified Model

Just as with the feature vector model, your training only has access to a few layers and does not update the convolutional layers. Now that you’ve trained two models, you can unify their layers again into a single model. You might be wondering why you’re combining the model after training instead of before. It’s a common practice to train your new layers with your feature layers locked or “frozen” to their original weights.

Once the new layers have gotten trained up, you can generally “unfreeze” more layers and train the new and the old together. This phase is often called fine-tuning the model.

So how do you unify these two models now? The answer is surprisingly simple. Create a third sequential model and add the two models with model.add. The code looks like this:

// combine the models
const combo = tf.sequential()
combo.add(shavedModel)
combo.add(transferModel)
combo.compile({
  optimizer: 'adam',
  loss: 'categoricalCrossentropy',
  metrics: ['accuracy'],
})
combo.summary()

The new combo model can be downloaded or trained further.

If you had joined the models before training the new layers, you’d likely see your model overfit the data.

Easy KNN: Bunnies Versus Sports Cars

KNN, like MobileNet, has a JS wrapper provided by Google. We can implement KNN transfer learning quickly by hiding all the complexity details use MobileNet and KNN NPM packages to make a quick transfer learning demo.

Not only are we going to avoid running any training, but we’ll also use existing libraries to avoid any deep dive into TensorFlow.js. We’ll be doing this for a flashy demo, but if you decide to build something more robust with these models, you should probably evaluate avoiding abstract packages that you don’t control. You already understand all the inner workings of transfer learning.

To do this quick demo, you’ll import the three NPM modules:

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/[email protected]/dist/tf.min.js">
</script>
<script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/[email protected]">
</script>
<script
src="https://cdn.jsdelivr.net/npm/@tensorflow-models/[email protected]">
</script>

For simplicity, the example code from this chapter has all the images on the page, so you can directly reference them. Now you can load MobileNet with mobileNet = await mobilenet.load(); and the KNN classifier with knnClassifier.create();.

The KNN classifier needs examples of each class. To simplify this process I’ve created the following helper function:

// domID is the DOM element ID
// classID is the unique class index
function addExample(domID, classID) {
  const features = mobileNet.infer( 
    document.getElementById(domID), 
    true                            
  );
  classifier.addExample(features, classID);
}

The infer method returns values rather than the rich JavaScript object of detections.

The image id on the page will tell MobileNet what image to resize and process. The tensor logic is hidden by JavaScript, but many chapters in this book have explained what is actually happening.

The MobileNet model returns the features (sometimes called embeddings) of the image. If this is not set, then the tensor of 1,000 raw values is returned (sometimes called logits).

Now you can add examples of each class with this helper method. You just name the image element’s unique DOM ID and what class it should be associated with. Adding three examples of each is as simple as this:

// Add examples of two classes
addExample('bunny1', 0)
addExample('bunny2', 0)
addExample('bunny3', 0)
addExample('sport1', 1)
addExample('sport2', 1)
addExample('sport3', 1)

Lastly, it’s the same system to predict. Get the features of an image, and ask the classifier to identify which class it believes the input is based on KNN.

// Moment of truth
const testImage = document.getElementById('test')
const testFeature = mobileNet.infer(testImage, true);
const predicted = await classifier.predictClass(testFeature)
if (predicted.classIndex === 0) { 
  document.getElementById("result").innerText = "A Bunny" 
} else {
  document.getElementById("result").innerText = "A Sports Car"
}

The classIndex is the number as passed in addExample. If a third class is added, that new index would be a possible output.

The web page text is changed from “???” to the result.

The result is that the AI can identify the correct class for a new image by comparing against six examples, as shown in Figure 11-8.

Figure 11-8. With only three images of each class, the KNN model predicts correctly

You can dynamically add more and more classes. KNN is an exciting and expandable way to utilize the experience of advanced models through transfer learning.