title customizing the convolution operation of a conv2d layer lukewoodhttps lukewood xyz date created 11032021 last modified 11032021 description this example shows how to implement custom convolution layers using the conv convolutionop api accelerator gpu introduction you may sometimes need to implement custom versions of convolution layers like conv1d and conv2d keras enables you do this without implementing the entire layer from scratch you can reuse most of the base convolution layer and just customize the convolution op itself via the convolutionop method this method was introduced in keras 2 7 so before using the convolutionop api ensure that you are running keras version 2 7 0 or greater a simple standardizedconv2d implementation there are two ways to use the conv convolutionop api the first way is to override the convolutionop method on a convolution layer subclass using this approach we can quickly implement a standardizedconv2dhttps arxiv orgabs1903 10520 as shown below the other way to use the conv convolutionop api is to directly call the convolutionop method from the call method of a convolution layer subclass a comparable class implemented using this approach is shown below example usage both of these layers work as dropin replacements for conv2d the following demonstration performs classification on the mnist dataset model data parameters the data split between train and test sets scale images to the 0 1 range make sure images have shape 28 28 1 convert class vectors to binary class matrices conclusion the conv convolutionop api provides an easy and readable way to implement custom convolution layers a standardizedconvolution implementation using the api is quite terse consisting of only four lines of code introduction you may sometimes need to implement custom versions of convolution layers like conv1d and conv2d keras enables you do this without implementing the entire layer from scratch you can reuse most of the base convolution layer and just customize the convolution op itself via the convolution_op method this method was introduced in keras 2 7 so before using the convolution_op api ensure that you are running keras version 2 7 0 or greater a simple standardizedconv2d implementation there are two ways to use the conv convolution_op api the first way is to override the convolution_op method on a convolution layer subclass using this approach we can quickly implement a standardizedconv2d https arxiv org abs 1903 10520 as shown below the other way to use the conv convolution_op api is to directly call the convolution_op method from the call method of a convolution layer subclass a comparable class implemented using this approach is shown below example usage both of these layers work as drop in replacements for conv2d the following demonstration performs classification on the mnist dataset model data parameters the data split between train and test sets scale images to the 0 1 range make sure images have shape 28 28 1 convert class vectors to binary class matrices conclusion the conv convolution_op api provides an easy and readable way to implement custom convolution layers a standardizedconvolution implementation using the api is quite terse consisting of only four lines of code
import tensorflow as tf
import keras
from keras import layers
import numpy as np


class StandardizedConv2DWithOverride(layers.Conv2D):
    def convolution_op(self, inputs, kernel):
        mean, var = tf.nn.moments(kernel, axes=[0, 1, 2], keepdims=True)
        return tf.nn.conv2d(
            inputs,
            (kernel - mean) / tf.sqrt(var + 1e-10),
            padding="valid",
            strides=list(self.strides),
            name=self.__class__.__name__,
        )





class StandardizedConv2DWithCall(layers.Conv2D):
    def call(self, inputs):
        mean, var = tf.nn.moments(self.kernel, axes=[0, 1, 2], keepdims=True)
        result = self.convolution_op(
            inputs, (self.kernel - mean) / tf.sqrt(var + 1e-10)
        )
        if self.use_bias:
            result = result + self.bias
        return result





num_classes = 10
input_shape = (28, 28, 1)


(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()


x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255

x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")


y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

model = keras.Sequential(
    [
        keras.layers.Input(shape=input_shape),
        StandardizedConv2DWithCall(32, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        StandardizedConv2DWithOverride(
            64, kernel_size=(3, 3), activation="relu"
        ),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation="softmax"),
    ]
)

model.summary()

batch_size = 128
epochs = 5

model.compile(
    loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"]
)

model.fit(
    x_train, y_train, batch_size=batch_size, epochs=5, validation_split=0.1
)