Could I implement my own convolution function to replace tf.nn.conv2d in tensorflow? How do I do?

Question

I want to implement my own convolution function just like tf.nn.conv2d function in tensorflow. I want to know whether I could do this with existed interface provided by tensorflow. Someone advised me to implement this according to the "how to add a new op in tensorflow" tutorial at the tensorflow website. Should I really implement my function in C++ firstly, then add it to tensorflow?

Answer 1

You might want to take a look at How to make a custom activation function with only Python in Tensorflow?

patapouf_ai explains how to add an operation using python code (without the need to implement C++ code.

Answer 2

One way with tensorflow API and model subclassing

#Custom class inherited from Layer class

class CustomConv2D(Layer):

def __init__(self, n_filters, kernel_size, n_strides, padding="valid"):
    super(CustomConv2D, self).__init__(name="custom_conv2D")
   
    # From tensorflow API
    self.conv = Conv2D(
        filters=n_filters,
        kernel_size=kernel_size,
        activation="relu",
        strides= n_strides,
        padding=padding
    )

    # For batch normalization and can be removed if not required for use case

    self.batch_norm = BatchNormalization()

def call(self, x, training):

    x = self.conv(x)
    x = self.batch_norm(x, training)

    return x