10 TensorFlow Tricks Every ML Practitioner Must Know

 2 years ago
source link: https://mc.ai/10-tensorflow-tricks-every-ml-practitioner-must-know-2/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

10 TensorFlow Tricks Every ML Practitioner Must Know

Why TensorFlow is the complete ML package

Photo by Author, Logo via TensorFlow

TensorFlow 2.x offers a lot of simplicity in building models and in overall TensorFlow usage. So what’s new in TF2?

  • Easy model building with Keras and eager execution.
  • Robust model deployment in production on any platform.
  • Powerful experimentation for research.
  • Simplifying the API by cleaning up deprecated APIs and reducing duplication.

In this article, we will explore 10 features of TF 2.0, that make working with TensorFlow smoother, reduces lines of code and increases efficiency as these functions/classes belong to the TensorFlow API.

1(a). The tf.data API for Building Input Pipelines

The tf.data API offers functions for data pipelining and related operations. We can build pipelines, map preprocessing functions, shuffle or batch a dataset and much more.

Building a pipeline from tensors

>>> dataset = tf.data.Dataset.from_tensor_slices([8, 3, 0, 8, 2, 1])
>>> iter(dataset).next().numpy()

Batch and Shuffle

# Shuffle
>>> dataset = tf.data.Dataset.from_tensor_slices([8, 3, 0, 8, 2, 1]).shuffle(6)
>>> iter(dataset).next().numpy()
# Batch
>>> dataset = tf.data.Dataset.from_tensor_slices([8, 3, 0, 8, 2, 1]).batch(2)
>>> iter(dataset).next().numpy()
array([8, 3], dtype=int32)
# Shuffle and Batch
>>> dataset = tf.data.Dataset.from_tensor_slices([8, 3, 0, 8, 2, 1]).shuffle(6).batch(2)
>>> iter(dataset).next().numpy()
array([3, 0], dtype=int32)

Zipping Two Datsets

>>> dataset0 = tf.data.Dataset.from_tensor_slices([8, 3, 0, 8, 2, 1])
>>> dataset1 = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6])
>>> dataset = tf.data.Dataset.zip((dataset0, dataset1))
>>> iter(dataset).next()
(<tf.Tensor: shape=(), dtype=int32, numpy=8>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)

Mapping External Functions

def into_2(num):
 return num * 2>>> dataset = tf.data.Dataset.from_tensor_slices([8, 3, 0, 8, 2, 1]).map(into_2)
>>> iter(dataset).next().numpy()

1(b). ImageDataGenerator

This is one of the best features of the tensorflow.keras API (in my opinion). The ImageDataGenerator is capable of generating dataset slices within batching and preprocessing along with data augmentation in real-time.

The Generator allows data flow directly from directories or from dataframes.

A misconception about data augmentation in ImageDataGenerator is that, the it adds more data to the existing dataset. Although that is the actual definition of data augmentation, in ImageDataGenerator, the images in the dataset are transformed dynamically at different steps in training so that the model can be trained on noisy data it hasn’t seen.

train_datagen = ImageDataGenerator(

Here, the rescaling is done on all the samples (for normalizing), while the other parameters are for augmentation.

train_generator = train_datagen.flow_from_directory(
 target_size=(150, 150),

We specify the directory for the real-time data flow. This can be done using dataframes as well.

train_generator = flow_from_dataframe(

The x_col parameter defines the full path of the image whereas the y_col parameter defines the label column for classification.

The model can be directly fed with the generator. Although the steps_per_epoch parameter needs to be specified which is essentially number_of_samples // batch_size.

 steps_per_epoch=(num_samples // batch_size),
 validation_steps=(num_val_samples // batch_size)

2. Data Augmentation with tf.image

Data augmentation is necessary. In case of insufficient data, making changes in the data and treating it as a separate datapoint is a very effective way for training under less data.

The tf.image API has tools for transforming images which can be later used for data augmentation with tf.data API discussed earlier.

flipped = tf.image.flip_left_right(image)
visualise(image, flipped)
Output of the above code snippet
saturated = tf.image.adjust_saturation(image, 5)
visualise(image, saturated)
Output of the above code snippet
rotated = tf.image.rot90(image)
visualise(image, rotated)
Output of the above code snippet
cropped = tf.image.central_crop(image, central_fraction=0.5)
visualise(image, cropped)
Output of the above code snippet

3. TensorFlow Datasets

pip install tensorflow-datasets

This is a very useful library as it is a single go-to point for a dump of very well-known datasets from various domains collected by tensorflow.

import tensorflow_datasets as tfdsmnist_data = tfds.load("mnist")
mnist_train, mnist_test = mnist_data["train"], mnist_data["test"]
assert isinstance(mnist_train, tf.data.Dataset)

A detailed list of the datasets available in tensorflow-datasets can be found on the Datasets page of the documentation .

Audio, Image, Image classification, Object Detection, Structured, Summarization, Text, Translate, Video are the genres offered by tfds.

4. Transfer Learning with Pretrained Models

Transfer Learning is the new cool in Machine Learning and it is as important as it sounds. It is infeasible and impractical to train a benchmark model which is already trained by someone else and under generous resources (for eg. multiple expensive GPUs that one may not afford). Transfer Learning, addresses this issue. A pretrained model can be reused for a given use case or can be extended for a different use case.

TensorFlow offers benchmark pretrained models that can be extended easily for the desired use case.

base_model = tf.keras.applications.MobileNetV2(

This base_model can be easily extended with additional layers or with different models. For eg:

model = tf.keras.Sequential([

For detailed list of other models and/or modules under tf.keras.applications, refer the docs page .

5. Estimators

An Estimator is TensorFlow’s high-level representation of a complete model, and it has been designed for easy scaling and asynchronous training

TensorFlow Docs

Premade estimators offer a very high level abstraction of a model, so you can directly focus on training the model without worrying about the lower level intricacies. For example:

linear_est = tf.estimator.LinearClassifier(
result = linear_est.evaluate(eval_input_fn)

This shows how easy it is to build and train an estimator using tf.estimator. Estimators can also be customised.

TensorFlow has many premade estimators including LinearRegressor, BoostedTreesClassifier, etc. A complete, detailed list of estimators can be found at the Tensorflow docs .

6. Custom Layers

Neural Nets are known for many layer deep networks wherein the layers can be of different types. TensorFlow contains many predefined layers (like Dense, LSTM, etc.). But for more complex architectures, the logic of a layer is much more complex than a primary layer. For such instances, TensorFlow allows building custom layers. This can be done by subclassing the tf.keras.layers.Layer class.

class CustomDense(tf.keras.layers.Layer):
 def __init__(self, num_outputs):
 super(CustomDense, self).__init__()
 self.num_outputs = num_outputs

 def build(self, input_shape):
 self.kernel = self.add_weight(

 def call(self, input):
 return tf.matmul(input, self.kernel)

As stated in the documentation , The best way to implement your own layer is extending the tf.keras.Layer class and implementing:

  1. __init__ , where you can do all input-independent initialization.
  2. build , where you know the shapes of the input tensors and can do the rest of the initialization.
  3. call , where you do the forward computation.

Although the kernel initialization can be done in __init__ itself, it is considered better to be initialized in build as otherwise, you would have to explicitly specify the input_shape on every instance of a new layer creation.

7. Custom Training

The tf.keras Sequential and the Model API makes training models easier. However, most of the time while training complex models, custom loss functions are used. Moreover, the model training can also differ from the default (for eg. applying gradients separately to different model components).

TensorFlow’s automatic differentiation helps calculating gradients in an efficient way. These primitives are used in defining custom training loops.

def train(model, inputs, outputs, learning_rate):
with tf.GradientTape() as t:
# Computing Losses from Model Prediction
current_loss = loss(outputs, model(inputs))
# Gradients for Trainable Variables with Obtained Losses
dW, db = t.gradient(current_loss, [model.W, model.b])
# Applying Gradients to Weights
model.W.assign_sub(learning_rate * dW)
model.b.assign_sub(learning_rate * db)

This loop can be repeated for multiple epochs and with a more customised setting as per the use case.

8. Checkpoints

Saving a TensorFlow model can be of two types:

  1. SavedModel : Saving the complete state of the model along with all the parameters. This is independent of source codes. model.save_weights('checkpoint')
  2. Checkpoints

Checkpoints capture the exact values of all the parameters used by a model. Models built with the Sequential API or the Model API can simply be saved in the SavedModel format.

However, for custom models, checkpoints are required.

Checkpoints do not contain any description of the computation defined by the model and thus are typically only useful when source code that will use the saved parameter values is available.

Saving a Checkpoint

checkpoint_path = “save_path”# Defining a Checkpoint
ckpt = tf.train.Checkpoint(model=model, optimizer=optimizer)
# Creating a CheckpointManager Object
ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)
# Saving a Model

Loading from a Checkpoint

TensorFlow matches variables to checkpointed values by traversing a directed graph with named edges, starting from the object being loaded.

Dependency Graph for Model Restoratrion via Docs
if ckpt_manager.latest_checkpoint:

9. Keras Tuner

This is a fairly new feature in TensorFlow.

!pip install keras-tuner

Hyper-parameter tuning or Hypertuning is the process of cherrypicking parameters that define the configuration of a ML model. These factors are the deciding factors for the performance of a model in the aftermath of feature engineering and preprocessing.

# model_builder is a function that builds a model and returns it
tuner = kt.Hyperband(

Along with HyperBand, BayesianOptimization and RandomSearch are also available for tuning.

img_train, label_train,
epochs = 10,

# Get the optimal hyperparameters

best_hps = tuner.get_best_hyperparameters(num_trials=1)[0]

Further, we train the model using the optmial hyper-parameters:

model = tuner.hypermodel.build(best_hps)
 validation_data=(img_test, label_test)

10. Distributed Training

If you have multiple GPUs and wish to optimize training by dispersing the training loop over multiple GPUs, TensorFlow’s various distributed training strategies are capable of optimizing the GPU usage and manipulate training over the GPUs for you.

tf.distribute.MirroredStrategy is the most common strategy used. How does it work anyway? The docs state:

  • All the variables and the model graph is replicated on the replicas.
  • Input is evenly distributed across the replicas.
  • Each replica calculates the loss and gradients for the input it received.
  • The gradients are synced across all the replicas by summing them.
  • After the sync, the same update is made to the copies of the variables on each replica.
strategy = tf.distribute.MirroredStrategy()with strategy.scope():
 model = tf.keras.Sequential([
 32, 3, activation='relu', input_shape=(28, 28, 1)
 tf.keras.layers.Dense(64, activation='relu'),


For other strategies and custom training loops, refer the documentation .


TensorFlow is sufficient for building almost all the components of a ML pipeline. The takeaway from this tutorial is an introduction to the various APIs provided by TensorFlow and a quick guide on how to use them.

Here is a link to the github repository of the code. Feel free to fork it.


The code used in this guide is referred from the following official TensorFlow Documentation:

Audio, Image, Image classification, Object Detection, Structured, Summarization, Text, Translate, Video are the genres offered by tfds.

About Joyk

Aggregate valuable and interesting links.
Joyk means Joy of geeK