Machine Learning in the Browser: Train and Serve a Mobilenet Model for Custom Image Classification | by Erdem Isbilen | Towards Data Science

There are several ways of fine-tuning a deep learning model but doing this on the web browser with WebGL acceleration is something that we experienced not such a long time ago, with the introduction of Tensorflow.js. I will use Tensorflow.js together with Angular to build a Web App that trains a convolutional neural network to detect malaria-infected cells with the help of Mobilenet and Kaggle dataset containing 27.558 infected and uninfected cell images.
Visit the live demo application to see the codes in action. The application runs in the Google Chrome browser without any issue.
Visit my GitHub repository for the full code of this project. You can find the images in the assets folder in this repository as well as the CSV file required to load the images into the browser. You should unzip the images and place them into the assets folder of the Angular project you are working on.
As I mentioned above, I will use the ‘mobilenet’ as a base model for our custom image classifier. The pre-trained ‘mobilenet’ model, which is tensorflow.js compatible, is relatively small (20MB) and can be directly downloaded from the Google API storage folder.
Uninfected and parasitized cells, which are 2 classes I would like to classify with our custom model whereas the original ‘mobilenet’ model is trained to classify 1000 different objects.
I will use all ‘mobilenet’ layers other than the last 5 layers and add a dense layer with 2 units and softmax activation on top of this truncated model to make the modified model suitable for our classification task.
I will not train all layers, as this will require a lot of time and computational power. It is also not necessary for our case as we use a pre-trained model which has a lot of representations learned already. Instead, I will freeze the majority of the layers and keep only the last 2 ones trainable.
Before start coding, we have to think about how we will feed the images to our custom model during the training. The mobilenet model requires specific image sizes (224x224x3) and image pre-processing operations and we have to apply the same pre-processes to our images before feeding them to our model. Besides, not to make our model biased to one class, we have to feed the same amount of images for each class during the training epochs.
After we all clear about the model and the dataset, it is time to use the Angular command-line interface to initialize the Angular web application.
npm install -g @angular/cling new TFJS-CustomImageClassificationcd TFJS-CustomImageClassification
Then I will use the ‘nmp’ package manager to install the tensorflow.js library. To polish the web app, I will use Angular Material classes.
TFJS-CustomImageClassification npm install @tensorflow/tfjs --saveTFJS-CustomImageClassification ng add @angular/material
First things first, let’s start coding with generating a function to modify the pre-trained model so that we can use this modified model for our specific task which is to classify the malaria-infected images. As I do not want to train the model starting from scratch, I will freeze all the layers except the ones I need to re-train to finetune the model.
//-------------------------------------------------------------// modifies the pre-trained mobilenet to detect malaria infected// cells, freezes layers to train only the last couple of layers//-------------------------------------------------------------async getModifiedMobilenet(){ const trainableLayers=     ['denseModified','conv_pw_13_bn','conv_pw_13','conv_dw_13_bn','conv _dw_13']; const mobilenet =  await tf.loadLayersModel('https://storage.googleapis.com/tfjs- models/tfjs/mobilenet_v1_0.25_224/model.json');console.log('Mobilenet model is loaded') const x=mobilenet.getLayer('global_average_pooling2d_1'); const predictions= <tf.SymbolicTensor> tf.layers.dense({units: 2,  activation: 'softmax',name: 'denseModified'}).apply(x.output);  let mobilenetModified = tf.model({inputs: mobilenet.input, outputs:  predictions, name: 'modelModified' }); console.log('Mobilenet model is modified')mobilenetModified = this.freezeModelLayers(trainableLayers,mobilenetModified) console.log('ModifiedMobilenet model layers are freezed')mobilenetModified.compile({loss: categoricalCrossentropy,  optimizer: tf.train.adam(1e-3), metrics:   ['accuracy','crossentropy']});mobilenet.dispose(); x.dispose(); return mobilenetModified}//-------------------------------------------------------------// freezes mobilenet layers to make them untrainable// just keeps final layers trainable with argument trainableLayers//-------------------------------------------------------------freezeModelLayers(trainableLayers,mobilenetModified){ for (const layer of mobilenetModified.layers) {  layer.trainable = false;  for (const tobeTrained of trainableLayers)  {    if (layer.name.indexOf(tobeTrained) === 0)     {      layer.trainable = true;      break;     }   }  } return mobilenetModified;}
To train the model, we need uninfected and infected cell images in 224x224x3 shape tensor and another 1-dimensional tensor containing 1,0 values indicating the classes of the images. What I will do is to read from a CSV file containing the image src and the class info, then generating HTMLImageElement to see them in the browser. Capture() function then will get the image id to generate required image tensors from the images on the browser. Please note that we have to pre-process the image tensors as the mobilenet expects normalized inputs. To be honest, what I used here as a data-pipeline is not the proper way of doing it, as I am loading the whole chunk of image data into the memory at once. It would be much better to use tf.fitDataset and to use the memory iteratively just when it is needed.
//-------------------------------------------------------------// this function generate input and target tensors for the training// input tensor is produced from 224x224x3 image in HTMLImageElement// target tensor shape2 is produced from the class definition//-------------------------------------------------------------generateData (trainData,batchSize){  const imageTensors = [];  const targetTensors = [];let allTextLines = this.csvContent.split(/\r|\n|\r/);  const csvSeparator = ',';  const csvSeparator_2 = '.';for ( let i = 0; i < batchSize; i++)  {    // split content based on comma    const cols: string[] = allTextLines[i].split(csvSeparator);    console.log(cols[0].split(csvSeparator_2)[0])if (cols[0].split(csvSeparator_2)[1]=="png")    {     const imageTensor = this.capture(i);      let targetTensor   =tf.tensor1d([this.label_x1[i],this.label_x2[i]]);      targetTensor.print();      imageTensors.push(imageTensor);      targetTensors.push(targetTensor);    }  }  const images = tf.stack(imageTensors);  const targets = tf.stack(targetTensors);  return {images, targets};}//-------------------------------------------------------------// converts images in HTMLImageElement into the tensors// takes Image Id in HTML as argument//-------------------------------------------------------------capture(imgId){  // Reads the image as a Tensor from the <image> element.  this.picture = <HTMLImageElement> document.getElementById(imgId);  const trainImage = tf.browser.fromPixels(this.picture);  // Normalize the image between -1 and 1. The image comes in  between  0-255, so we divide by 127 and subtract 1.  const trainim =  trainImage.toFloat().div(tf.scalar(127)).sub(tf.scalar(1));  return trainim;}
As we prepared the data and the model, it is now time to fine-tune the model so that it can classify the malaria-infected cell images. To reduce the time for the training process, I will use 120 images in total, a training process of 5 epochs with data batches each containing 24 images. Loss function to be ‘categoricalCrossentropy’ and the ‘adam optimizer’ to be used in training with relatively small learning rate value. I used ‘onBatchEnd’ callback function which is provided by Keras to write the training metrics into the console for each epoch. We have to dispose of the tensors to free up the memory once the training is completed. As a final step, we will save the trained model into our local storage to use it later for inference.
async fineTuneModifiedModel(model,images,targets){  function onBatchEnd(batch, logs)  {    console.log('Accuracy', logs.acc);    console.log('CrossEntropy', logs.ce);    console.log('All', logs);  }  console.log('Finetuning the model...');    await model.fit(images, targets,  {    epochs: 5,    batchSize: 24,    validationSplit: 0.2,    callbacks: {onBatchEnd}    }).then(info => {    console.log    console.log('Final accuracy', info.history.acc);    console.log('Cross entropy', info.ce);    console.log('All', info);    console.log('All', info.history['acc'][0]);       for ( let k = 0; k < 5; k++)    {      this.traningMetrics.push({acc: 0, ce: 0 , loss: 0});      this.traningMetrics[k].acc=info.history['acc'][k];      this.traningMetrics[k].ce=info.history['ce'][k];      this.traningMetrics[k].loss=info.history['loss'][k];    }    images.dispose();    targets.dispose();    model.dispose();  });}
Visit my GitHub repository for the full code of this project. You can find the images in the assets folder in this repository as well as the CSV file required to load the images into the browser. You should unzip the images and place them into the assets folder of the Angular project you are working on.
Visit the live demo application to see the codes in action. The application runs in the Google Chrome browser without any issue.