In this section, we will illustrate the TensorFlow code for training the model. The model is trained for a modest 10 epochs, to avoid overfitting. The learning rate used for the optimizer is 0.001, while the training batch size and the validation batch size are set at 250 and 50, respectively. One thing to note is that we are saving the model graph definition in the model.pbtxt file, using the tf.train.write_graph function. Also, once the model is trained, we will save the model weights in the checkpoint file, model_ckpt, using the tf.train.Saver function. The model.pbtxt and model_ckpt files will be used to create an optimized version of the TensorFlow model in the protobuf format, which can be integrated with the Android app:
def _train(self):
self.num_batches = int(self.X_train.shape[0]//self.batch_size)
self._build_model()
...