Easy Image Classification with TensorFlow 2.0 | by Cameron Cruz | Towards Data Science

At the TensorFlow Dev Summit 2019, Google introduced the alpha version of TensorFlow 2.0. Historically, TensorFlow is considered the “industrial lathe” of machine learning frameworks: a powerful tool with intimidating complexity and a steep learning curve. If you’ve used TensorFlow 1.x in the past, you know what I’m talking about. This 2.0 release represents a concerted effort to improve the usability, clarity and flexibility of TensorFlow. Here are some highlights:
Eager execution is enabled by default, without sacrificing the performance optimizations of graph-based execution.
APIs are cleaner, more consistent and less redundant.
Tighter Keras integration as the high-level API.
And more
As a result, TensorFlow 2.0 is more Pythonic and less daunting to learn, while retaining the lower-level customization and complexity if you need it. Let’s get started with TensorFlow 2.0 by exploring how to apply its high-level APIs in a classic image classification setting.
Google Colaboratory makes it really easy to setup Python notebooks in the cloud. With free access to a GPU for up to 12 hours at a time, Colab has quickly become my go-to platform for performing machine learning experiments.
Let’s install the TensorFlow 2.0 alpha release (GPU version) on a Colab notebook via pip.
!pip install tensorflow-gpu==2.0.0-alpha0
To verify it installed properly:
import tensorflow as tfprint(tf.__version)# Output: 2.0.0-alpha0
You should be good to go! If you run into issues, double check that your Colab runtime has “GPU” as the runtime accelerator in Edit>Notebook settings.
Let’s use the dataset from the Aerial Cactus Identification competition on Kaggle. Our task is to build a classifier capable of determining whether an aerial image contains a columnar cactus or not. This dataset is a modified version of the Cactus Aerial Photos dataset [1](Kaggle resized each image to 32x32 pixels).
For example code on downloading/unzipping datasets from Kaggle, see the full notebook here.
Let’s load the image file paths and their corresponding labels into lists using pandas, then create a train-validation split of 90–10 using sklearn.model_selection.
train_csv = pd.read_csv('data/train.csv')# Prepend image filenames in train/ with relative pathfilenames = ['train/' + fname for fname in train_csv['id'].tolist()]labels = train_csv['has_cactus'].tolist()train_filenames, val_filenames, train_labels, val_labels =   train_test_split(filenames,                 labels,                 train_size=0.9,                 random_state=42)
Now that we have image filenames and labels split into training and validation sets, we can create respective tf.data.Dataset objects.
train_data = tf.data.Dataset.from_tensor_slices(  (tf.constant(train_filenames), tf.constant(train_labels)))val_data = tf.data.Dataset.from_tensor_slices(  (tf.constant(val_filenames), tf.constant(val_labels)))
However, our dataset still contains just the image filenames, not the actual images. We need to define a function that can load the images from file and perform any necessary preprocessing. While we’re at it, let’s shuffle and batch the datasets, too.
IMAGE_SIZE = 96 # Minimum image size for use with MobileNetV2BATCH_SIZE = 32# Function to load and preprocess each imagedef _parse_fn(filename, label):    img = tf.io.read_file(img)    img = tf.image.decode_jpeg(img)    img = (tf.cast(img, tf.float32)/127.5) - 1    img = tf.image.resize(img, (IMAGE_SIZE, IMAGE_SIZE))    return img, label# Run _parse_fn over each example in train and val datasets# Also shuffle and create batchestrain_data = (train_data.map(_parse_fn)             .shuffle(buffer_size=10000)             .batch(BATCH_SIZE)             )val_data = (val_data.map(_parse_fn)           .shuffle(buffer_size=10000)           .batch(BATCH_SIZE)           )
Transfer learning speeds up training by enabling us to reuse existing pre-trained image classification models, only retraining the top layer of the network that determines the classes an image can belong to [2].
Let’s use TensorFlow 2.0’s high-level Keras API to quickly build our image classification model. For transfer learning, we can use a pre-trained MobileNetV2 model as the feature detector. MobileNetV2 is the second iteration of MobileNet released by Google with the goal of being smaller and more lightweight than models like ResNet and Inception for running on mobile devices [3]. Let’s load the MobileNetV2 model pre-trained on ImageNet without the top layer, freeze its weights, and add a new classification head.
IMG_SHAPE = (IMAGE_SIZE, IMAGE_SIZE, 3)# Pre-trained model with MobileNetV2base_model = tf.keras.applications.MobileNetV2(    input_shape=IMG_SHAPE,    include_top=False,    weights='imagenet')# Freeze the pre-trained model weightsbase_model.trainable = False# Trainable classification headmaxpool_layer = tf.keras.layers.GlobalMaxPooling2D()prediction_layer = tf.keras.layers.Dense(1, activation='sigmoid')# Layer classification head with feature detectormodel = tf.keras.Sequential([    base_model,    maxpool_layer,    prediction_layer])learning_rate = 0.0001# Compile the modelmodel.compile(optimizer=tf.keras.optimizers.Adam(lr=learning_rate),               loss='binary_crossentropy',              metrics=['accuracy'])
Note that it is recommended to use TensorFlow optimizers to train tf.keras models. In TensorFlow 2.0, the optimizers in the former tf.train and tf.keras.optimizers APIs have been unified under tf.keras.optimizers , where the original tf.keras optimizers have been replaced with upgraded TensorFlow optimizers [4]. Thus, applying TensorFlow optimizers is now a simpler and more consistent experience, fully supporting usage with the tf.keras API and not sacrificing performance.
The tf.keras API in TensorFlow 2.0 now has full support for the tf.data API, so we can easily use our tf.data.Dataset objects when training the model [5]. Again, training is executed eagerly by default now without sacrificing the performance benefits of graph-based execution.
num_epochs = 30steps_per_epoch = round(num_train)//BATCH_SIZEval_steps = 20model.fit(train_data.repeat(),          epochs=num_epochs,          steps_per_epoch = steps_per_epoch,          validation_data=val_data.repeat(),           validation_steps=val_steps)
After 30 epochs, the model’s validation accuracy increases from around 0.63 to 0.94.
Let’s try to improve the accuracy of our model even further. When we applied transfer learning, we only trained the new classification head of the model, freezing the weights from MobileNetV2. If we didn’t freeze those weights initially, the model would “forget” all the knowledge it started with since the new classification head was randomly initialized [2]. However, now that we’ve trained the classification head first, we can unfreeze layers in the pre-trained model to fine-tune on this particular dataset.
# Unfreeze all layers of MobileNetV2base_model.trainable = True# Refreeze layers until the layers we want to fine-tunefor layer in base_model.layers[:100]:  layer.trainable =  False# Use a lower learning ratelr_finetune = learning_rate / 10# Recompile the modelmodel.compile(loss='binary_crossentropy',              optimizer = tf.keras.optimizers.Adam(lr=lr_finetune),              metrics=['accuracy'])# Increase training epochs for fine-tuningfine_tune_epochs = 30total_epochs =  num_epochs + fine_tune_epochs# Fine-tune model# Note: Set initial_epoch to begin training after epoch 30 since we# previously trained for 30 epochs.model.fit(train_data.repeat(),           steps_per_epoch = steps_per_epoch,          epochs=total_epochs,           initial_epoch = num_epochs,          validation_data=val_data.repeat(),           validation_steps=val_steps)
After an additional 30 epochs of fine-tuning, the model achieved a validation accuracy of 0.986. Based on the accuracy and loss graphs, more epochs may result in even greater improvements.
In this article, we looked at how TensorFlow 2.0's focus on usability, clarity and flexibility makes it much less daunting to get started with TensorFlow for machine learning experiments. Eager execution and improved high-level APIs abstract away much of TensorFlow’s usual complexity, making it much easier to quickly implement and run a quintessential image classification experiment.
At the time of writing, this is just the alpha release of TensorFlow 2.0, with a final release expected sometime later this year. Already, it’s apparent that the TensorFlow team is forging a significantly more intuitive iteration of TensorFlow. This is likely to improve the productivity of machine learning engineers across the board, reducing the usual amount of complexity while retaining lower-level control for those that need it. In addition, while TensorFlow was already a popular choice among machine learning experts, the smoother learning curve makes it a more attractive option for beginners as well.
Let me know what you think of TensorFlow 2.0 so far in the comments! Also, if you’re interested in the full code for this tutorial, see the notebook here. If you have any questions, don’t hesitate to ask in the comments section, too.
This tutorial is inspired from the Transfer Learning Using Pretrained ConvNets on TensorFlow.org. To read the original tutorial, see here.
[1] J. Vasquez-Gomez, Cactus Aerial Photos (2019), Kaggle.com
[2] Transfer Learning Using Pretrained ConvNets (2019), TensorFlow.org
[3] M. Sandler, A. Howard, M. Zhu, A. Zhmonginov, L. C. Chen, MobileNetV2: Inverted Residuals and Linear Bottlenecks (2019), Google Inc.
[4] F. Chollet, TensorFlow 2.0: Optimizer Unification (2018), TensorFlow GitHub
[5] Roadmap (2019), TensorFlow.org