Custom group convolution for TensorFlow 2

The group convolution was initially implemented in Alexnet, as the response for limited GPU RAM resources. Back in 2012, Alex Krizknevsky together with Ilya Sutskever and Geoffrey Hinton was working on new CNN architecture for image classification. They designed a state-of-the-art neural network called Alexnet, but faced a training problem. The network was not fitting inside the GPU memory. The model was 3 GB big, but the NVIDIa GPU had only 1,5 GB of available memory. They came up with a group convolution as the solution to reduce the model size and as a side result, they improved the model accuracy. Kill two birds with one stone.

The standard 2D convolution filter has dimensions of HxWxD, where D is the depth of the previous feature tensor (number of channels). For example, each kernel can have 3x3x256 dimensions if the previous feature tensor has 256 channels (depth). In the case of group convolution, the depth of the kernel is divided into groups and as a result, the depth is significantly smaller. For example, if we use a group of size 2, the kernel from the example above would have dimensions of 3x3x128. It is important that the depth of the kernel needs to be divisible by a number of groups. The group convolution is well illustrated in Figure 1.

Figure 1. Group convolution.

The same kernel is applied at the beginning of the features tensor and at the end. Because the kernel is twice smaller, the number of trainable parameters is twice smaller as well and as the result, the memory footprint is twice smaller.

Custom group convolution in TensorFlow

In TensorFlow, a custom layer can be created by subclassing a tf.keras.layers.Layer. It needs to contain 2 essential functions:

• def build(self, inputs_shape):

• def call(self, inputs):

The build function tells how to build a layer while building a model. It contains instructions for weights and bias (if present) initialization. The call function tells how to perform the call operation on the layer. It contains instructions for applying weights and bias. Weight and bias are standard members in neural network layers. However, the custom layer can contain different members or ways of applying those members.

For group convolution in the build function, we need to initialize kernels to the size of HxWx(Depth/groups) and define the group_conv operation. Group convolution operation can be performed with lambda as follows:

self.groupConv = lambda i, k: tf.nn.conv2d(i, k, strides=self.strides, padding=self.padding, data_format=self.data_format, dilations=self.dilation_rate, name=self.name)

Lambda function takes as the argument i and k and performs conv2d operation on them. We will use the input tensor as i and a kernel as k. In the call function, we need to split an input tensor into a number of groups among the depth channel and the number of kernels (new feature tensor depth) into a number of groups. We perform convolution 2D operations on each tensor and corresponding kernel using our previously defined lambda function. After that, we need to concatenate results back into a single tensor. Group convolution of size 2 can be seen as the convolution of two layers with two corresponding filters, performed side-by-side, where each layer has its own filter set.

The demo code for custom group convolution in TensorFlow is presented below.

TensorFlow code for custom layer — GroupConv2D.