How to Flatten layer input in TensorFlow

This tutorial explains how to flatten a input layer in TensorFlow. With the use of tf.keras.layers.Flatten input can be flattened without affecting batch size.

Lets understand flattening of input with below example

Import TensorFlow


import tensorflow as tf

Create a sequential model with tf.keras.


model = tf.keras.Sequential()

Add a conv2D layer to the model


model.add(tf.keras.layers.Conv2D(64, 3, 1, input_shape=(3, 64, 64)))

Check the model output before applying tf.keras.layers.Flatten


print(model.output)

#### Output #####
Tensor("conv2d/BiasAdd:0", shape=(None, 1, 62, 64), dtype=float32)

Add the Flatten layer to the model for flattening the input


model.add(tf.keras.layers.Flatten())

Check the model output after applying tf.keras.layers.Flatten


print(model.output)

#### Output ####
Tensor("flatten/Reshape:0", shape=(None, 3968), dtype=float32)