TensorFlow Computer Vision & Deep Learning Examples

Reading code is one effective way to get professional in TensorFlow (TF). In this article, we reuse the examples in the TensorFlow tutorial but we keep them condense and take away codes that are for tracing or demonstration purposes. We also keep our discussion minimum so you can browse through as many examples as possible to get a complete picture. If you have problems to follow, in particular after reading the first example, please read the articles in this series first.

Yet, in most examples, we keep all the boilerplate codes required by the examples. So skip them if needed. We use snapshots for the code so you cannot cut-and-paste them because of potential TF API changes. Refer to the link on the top of the example instead. Here is the list of examples, tested with TF 2.4.0 released in Dec 2020.

  • Keras MNIST data: Sequential Model using Dense Layers,
  • Keras MINST data: Custom CNN Model Class trained with GradientTape & Dataset,
  • Custom layer creating new model parameters,
  • Dataset performance,
  • Overfitting,
  • Save and load,
  • Classify flowers with data augmentation layers,
  • Classify flowers with data augmentation using dataset mapping,
  • Transfer learning,
  • Image segmentation,
  • Regression on CSV file: Using Pandas to process data,
  • CSV preprocessing,

Keras MNIST data: Sequential Model using Dense Layers

Classify MNIST NumPy data using a Sequential Model with dense layers.

Model summary:

Keras MINST data: Custom CNN Model Class trained with GradientTape & Dataset

In this example,

  • train with a Dataset (in particular, if samples cannot fit into the memory),
  • a custom CNN model class, and
  • a custom training with GraidentTape.

First, we create the dataset and define the custom CNN model.

Custom training with GradientTape and testing step:

Training loop:

Custom layer creating new model parameters

A custom layer can contain other layers and/or have its own layer parameters. However, the input shape may not be known at the time of instantiation. A separate method “build” is called to create parameters once TF knows the input shape, say by invoking the layer with an input for the first time.


In this example:

  • Load Higgs CSV dataset using tf.data,
  • Train with a custom learning rate decay for the optimizer,
  • Record data later for TensorBoard,
  • Apply regularization and dropout to avoid overfitting, and
  • Apply early stopping.

The detection of Higgs boson particles proves the presence of Higgs Field that gives mass to the fundamental particles (quarks, leptons, etc …). This example predicts the class of an event — a “signal” or a “background”. “Signal” indicates the 4-lepton events have occurred from decays involving a Higgs boson. Otherwise, it is “background” from decays not involving a Higgs boson. First, we load a CSV dataset using tf.data. Each row contains the label and 28 features — 21 are kinematic properties measured by the particle detectors and 7 derived from these measurements.

Create a dataset with features and labels separated — ds is batched so the mapping is done in a single operation, and then later un-batched.

Create the training and validation dataset — it uses “take” and “skip” to partition the original training samples.

The optimizer will use a custom learning rate for the optimizer.

compile_and_fit will configure the training and fit the model. We also set up callbacks to log data for the TensorBoard.

We built a “combined” model with regularization and dropout to avoid overfitting.

Finally, we train the model.

We store the TensorBoard data for this training and model (called combined) under $tmp/tensorboard_logs/regularizers/combined. We can train a different model under a different director, say regularizers/other_model.

We can review the data using the TensorBoard for this training and other training under “regularizers”.

Dataset Performance

To improve Dataset performance, we should cache and prefetch data. But we may not need to shuffle validation or testing dataset.

Save & Load

In this example, we save and restore a model. First, we have the boilerplate code in loading the MNIST data and create a dense model.

Next, we create a callback for model.fit to save a model. For these checkpoints, we save the weights only.

Here is the checkpoint created.

We can instantiate a new model and reload the weights again.

We can add the epoch number as part of the checkpoint filename and we can change the frequency of the checkpoint (every 5 epochs below).

Here are the checkpoints after 50 epochs.

And we can load the latest checkpoint with:

To salve the weights manually instead:

Or save the whole model. In the latter case, we don't need the Python code to restore the model.


If we have access to the optimizer and the data iterator, like in GradientTape, we can use CheckpointManager to save checkpoints. Here is the boilerplate code first.

Then, we can configure a CheckpointManager.

This includes an option to keep max_to_keep checkpoints only.

And we can use it to save and restore the model.

Classify flowers with data augmentation layers

In this example, we use data augmentation to improve model performance. Prepare training and validation dataset:

Build and train a model with Keras pre-processing layers:

Perform a test:

Classify flowers with data augmentation using dataset mapping

We perform the data augmentation using the mapping in dataset pipelining. In this example, we perform resizing, rescaling, flipping, and rotation. Here is the boilerplate code for creating dataset and augmentation layers.

Then, we use dataset mapping for data augmentation.

This is the model and the training.

Transfer Learning

Here is the boilerplate code in loading cats and dogs pictures in creating datasets.

Then, we construct a model with data augmentation layers, a preprocessing layer from the MobileNet v2, a pre-trained MobileNet v2, and a classification head. In the first part of the transfer learning, we freeze the MobileNet v2 layers (line 54) and train the classification head only. We also set the training to False for the MobileNet v2 (line 65) such that it will run in inference mode: dropout layers will not be applied and it uses the original training means and variances for the Batch Normalization.

Next, we train the model,

Now, we go to the second phase of the training. We will only freeze the first 100 layers of the MobileNet v2 and fine-tune the rest of the model again.

Once it is trained again, we use it to make predictions for testing images.

Image Segmentation

Image segmentation creates a mask to segment an object.


First, we load the dataset and create the preprocessing layers.

Then, we create the datasets.

The image segmentation model contains downsampling layers followed by upsampling layers. Here is the downsampling.

and the upsampling part.

The more complicated logic above deals with the skip connections between the downsampling and upsampling layers — the horizontal connections between the downsampling and the upsampling layers below that have the same spatial resolution.

Here is the more precise diagram.

Modified from source

At last, we train the model and make predictions.

Regression on CSV file: Using Pandas to process data

This example is a simple regression problem in predicting the MPG of a car from features containing in a CSV file. Prepare datasets for CSV files using Pandas:

Visualize features co-relationships.

Create training/testing features/labels.

Create a model with a Normalization layer to normalize each feature.

Fit, evaluate & predict:

Save & load a model:

And the model will be saved as:

CSV preprocessing

In this example, we will predict who will survive the Titanic accident. This example demonstrates how to preprocess CSV data.

First, we will load the data from a CSV file.

We collect all the numeric fields, normalize it, and put them into preprocessed_inputs.

Convert all category fields into one-hot vectors and append them in preprocessed_inputs.

Concatenate all input and create a pre-processing model.

Create a classification model. Train and make predictions.

Credits and References

All the source code is originated or modified from the TensorFlow tutorial.

Deep Learning