How to apply convolution on the last three dimensi

Question

Usually the input tensor of the Conv2D in Keras is a 4D tensor with the dimension batch_size * n * n * channel_size. Now I have a 5D tensor with the dimension batch_size * N * n * n * channel_size and I want to apply the 2D convolutional layer for the last three dimensions for each i in N. For example, if the kernel size is 1, then I expect that the output will have the dimension batch_size * N * n * n * 1.

Anyone knows some easy ways to implement it with Keras?

For example, for the fully-connected layer Keras can do it automatically. If the input has the shape batch_size * N * n, then the Dense layer in Keras will set a FC layer for each i in N. Hence we will get the output with batch_size * N * m, if we set Dense(m).

Answer 1

You can use the TimeDistributed layer wrapper to apply the same convolution layer on all the images in the 5D tensor. For example:

model = Sequential()
model.add(TimeDistributed(Conv2D(5, (3,3), padding='same'), input_shape=(10, 100, 100, 3)))

model.summary()

Model summary:

Layer (type)                 Output Shape              Param #   
=================================================================
time_distributed_2 (TimeDist (None, 10, 100, 100, 5)   140       
=================================================================
Total params: 140
Trainable params: 140
Non-trainable params: 0
_________________________________________________________________