nn.Conv2d and F.conv2d with groups == input_channels (DepthWise) generates PartitionedCall in tensorflow frozen_graph

When I define a define a model with DepthWise convolutions (groups == input_channels) the model is converted sucessfully but the tensorflow frozen_graph of this model cannot be converted to tensorflow.js. The problem is that keras.layers.Conv2D generates a PartitionedCall in the frozen_graph that cannot be converted to tensorflow.js.

I provide the python code to reproduce the problem:

import torch.nn as nn
import torch

from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

import tensorflow as tf
import nobuco
from nobuco.commons import ChannelOrder, ChannelOrderingStrategy

class ExampleModel(nn.Module):
    def __init__(self, 
                **kwargs):

        super(ExampleModel, self).__init__()
        self.layer1 = nn.Conv2d(16, 16, (3,3), (1,1), (0,0), (1,1), 16)
        self.layer2 = nn.ReLU()

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x

model = ExampleModel()

# Put model in inference mode
model.eval()

x = torch.randn(1, 16, 113, 113, requires_grad=False)

keras_model = nobuco.pytorch_to_keras(
    model,
    args=[x], kwargs=None)

# Assuming 'model' is your Keras model
full_model = tf.function(lambda x: keras_model(x))
full_model = full_model.get_concrete_function(
    tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))

# Convert Keras model to frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

# Print the input and output tensors
print("Frozen model inputs: ", frozen_func.inputs)
print("Frozen model outputs: ", frozen_func.outputs)

# Save frozen graph to disk
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                logdir='.',
                name='ExampleModel.pb',
                as_text=False)

Inspecting the ExampleModel.pb with Netron this is what happens:

Screenshot from 2024-03-26 14-17-34

In order to fix this error, I made a custom nn.Conv2d converter:

@nobuco.converter(nn.Conv2d)
def converter_Conv2d(self, input: Tensor):
    weight = self.weight
    bias = self.bias
    groups = self.groups
    padding = self.padding
    stride = self.stride
    dilation = self.dilation

    out_filters, in_filters, kh, kw = weight.shape

    weights = weight.cpu().detach().numpy()

    if groups == 1:
        weights = tf.transpose(weights, (2, 3, 1, 0))
    else:
        weights = tf.transpose(weights, (2, 3, 0, 1))

    if bias is not None:
        biases = bias.cpu().detach().numpy()
        params = [weights, biases]
        use_bias = True
    else:
        params = [weights]
        use_bias = False

    if isinstance(dilation, numbers.Number):
        dilation = (dilation, dilation)

    if isinstance(padding, numbers.Number):
        padding = (padding, padding)

    pad_str = 'valid'
    pad_layer = None

    if padding == 'same':
        pad_str = 'same'
    elif padding != (0, 0):
        pad_layer = keras.layers.ZeroPadding2D(padding)

    if groups == 1:
        conv = keras.layers.Conv2D(filters=out_filters,
                                kernel_size=(kh, kw),
                                strides=stride,
                                padding=pad_str,
                                dilation_rate=dilation,
                                groups=groups,
                                use_bias=use_bias,
                                weights=params
                                )
    else:
        conv = keras.layers.DepthwiseConv2D(
                    kernel_size=(kh, kw),
                    strides=stride,
                    padding=pad_str,
                    use_bias=use_bias,
                    activation=None,
                    depth_multiplier=1,
                    weights=params,
                    dilation_rate=dilation,
                )

    def func(input):
        if pad_layer is not None:
            input = pad_layer(input)
        output = conv(input)
        return output
    return func

But I think that probably is better to fix this in the source code.

AlexanderLutsenko / nobuco

nn.Conv2d and F.conv2d with groups == input_channels (DepthWise) generates PartitionedCall in tensorflow frozen_graph #30