How to use 2D convolution layer in TensorFlow | tf.keras

gcptutorials.com TensorFlow

The convolution layer uses filters that perform convolution operations as it is scanning the input I with respect to its dimensions. Lets understand working of 2D convolution layer with an example. 2D convolution layer can be used from tf.keras.layers api. Use Colab notebook for executing code snippets.

In this post we will touch following TensorFlow modules and functions.

tensorflow-datasets for loading the image dataset
tf.expand_dims and tf.dtypes.cast for changing image dimensions and dtype
tf.keras.layers.Conv2D for using the 2D convolution layer
plt.imshow for displaying the image

Load `caltech101` dataset with tensorflow datasets


import tensorflow as tf
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt
ds, dsinfo = tfds.load('caltech101', split='train', with_info=True)

Visualize one input image from the dataset


for example in list(ds)[2:3]:
image, label = example['image'], example['label']
print(image.shape)
print(label)
plt.imshow(image)

Example Output:


(225, 300, 3)
tf.Tensor(51, shape=(), dtype=int64)

Change image shape and dtype to make it compatible with `tf.keras.layers.Conv2D`


image  = tf.expand_dims(image, 0)
image = tf.dtypes.cast(image, tf.float32)
print(image.shape)

Example Output:


(1, 225, 300, 3)

Create 2D conv layer with `tf.keras.layers` and provide input image

Number of feature maps generated with 2D convolution layer depends on integer value provided to the filter argument in the layer, in this example we have filter=5, hence 5 feature maps would be generated.

Shape of feature map is controlled with stride argument provided in the layer, in this example we have stride =(1,1) , so shape of feature map would be same as of provided input image, in case we choose stride =(2,2) then shape of feature map would be half of the size of provided input.


x = tf.keras.layers.Conv2D(filters=5, kernel_size=(2,2), strides=(1,1))(image)
print(x.shape)

Example Output:


(1, 224, 299, 5) # notice change in last dimension as it got changed from 3 to 5 because of filter argument

Visualize feature maps generated by 2D convolution layer

Filter value provided to tf.keras.layers.Conv2D is 5, all the feature maps can be viewed with following code plt.imshow(x[0, :, :, 0]), plt.imshow(x[0, :, :, 1]), plt.imshow(x[0, :, :, 2]) , plt.imshow(x[0, :, :, 3]), plt.imshow(x[0, :, :, 4]) . Below is the example output for first feature map.