Group Conv in TensorFlow 2

Viewed 2397

Does TensorFlow2.x support Group Convolution?
I have seen a lot of posts/blogs/tutorials saying that TensorFlow does not support it, and I have also seen some posts saying that tf.keras.layers.DepthwiseConv2D is equivalent to group convolution. However, I noticed that there is a group parameter in tf.keras.layers.Conv2d, is this the group convolution described in many papers (for example, ResNeXt paper)? Or am I misunderstanding it?
Any help and explanation would be great!

edit: an example of group conv (the third one) and equivalent parrellel conv (first two). Example from ResNeXt paper group conv

a group conv of 32 groups with depth 4 in pytorch, which means total output channel is 128:
torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=(3,3), groups=32)

more specifically, a group conv of n groups and depth of d with input channel i will split i input channels into n groups of equal size, each group will be a normal convolution with the same kernel size, stride that has i/n channels as input, and d channels as output. The output of all groups will be concatenated into n*d channels and passed to the next layer as input.

2 Answers

Yes, tensorflow does support the Group Conv directly with the groups argument. From Conv2D arguments in the official docs of TF2:

groups: A positive integer specifying the number of groups in which the input is split along the channel axis. Each group is convolved separately with filters / groups filters. The output is the concatenation of all the groups results along the channel axis. Input channels and filters must both be divisible by groups.

Tensorflow 2 definitely does NOT support grouped convolution!

While the Doc claims that it is supported by tf.keras.layers.Conv2D, as the 'groups' argument, when you try it you get the oft-reported error:

“UnimplementedError: Fused conv implementation does not support grouped convolutions for now.”

This error has been repeatedly reported since 2019.

The special case of Depthwise Convolution, where the number of groups equals the number of channels, is separately supported in tf.keras.layers.DepthwiseConv2D.

Related