For this model we're going to train twice. For the first round of training we will do feature extraction for 10 epochs by training with the network frozen, only adjusting the fully connected layer weights, as we discussed in the Transfer network architecture section. Then, in the next section we will unfreeze some of the layers and train again, fine-tuning for another 10 epochs, as shown in the following code:

data_dir = "data/train/"
val_dir = "data/val/"
epochs = 10
batch_size = 30
model = build_model_feature_extraction()
train_generator, val_generator = setup_data(data_dir, val_dir)
callbacks_fe = create_callbacks(name='feature_extraction')
# stage 1 fit
model.fit_generator(
    train_generator,
    steps_per_epoch=train_generator.n // batch_size,
    epochs=epochs,
    validation_data=val_generator,
    validation_steps=val_generator.n // batch_size,
    callbacks=callbacks_fe,
    verbose=1)

scores = model.evaluate_generator(val_generator, steps=val_generator.n // batch_size)
print("Step 1 Scores: Loss: " + str(scores[0]) + " Accuracy: " + str(scores[1]))

In the preceding example, we're using the ImageDataGenerator n attribute to know the total number of images available to the generator and define the steps per epoch as that number divided by the batch size.

The rest of this code should be familiar.

As previously mentioned, we're only going to need to train for about 10 epochs. Now, let's take a look at that training process in TensorBoard:

As you can see, the network is performing really well even after a single epoch. We achieve very marginal performance improvements until approximately epoch 7. During epoch 7, we achieve our best performance, resulting in 0.9828 accuracy and 0.0547 loss.