title involutional neural networks aritra roy gosthipatyhttps twitter comarig23498 date created 20210725 last modified 20210725 description deep dive into locationspecific and channelagnostic involution kernels accelerator gpu introduction convolution has been the basis of most modern neural networks for computer vision a convolution kernel is spatialagnostic and channelspecific because of this it isn t able to adapt to different visual patterns with respect to different spatial locations along with locationrelated problems the receptive field of convolution creates challenges with regard to capturing longrange spatial interactions to address the above issues li et al rethink the properties of convolution in involution inverting the inherence of convolution for visualrecognitionhttps arxiv orgabs2103 06255 the s propose the involution kernel that is locationspecific and channelagnostic due to the locationspecific nature of the operation the s say that selfattention falls under the design paradigm of involution this example describes the involution kernel compares two image classification models one with convolution and the other with involution and also tries drawing a parallel with the selfattention layer setup set seed for reproducibility convolution convolution remains the mainstay of deep neural networks for computer vision to understand involution it is necessary to talk about the convolution operation imgurhttps i imgur commsklsm5 png consider an input tensor x with dimensions h w and cin we take a collection of cout convolution kernels each of shape k k cin with the multiplyadd operation between the input tensor and the kernels we obtain an output tensor y with dimensions h w cout in the diagram above cout3 this makes the output tensor of shape h w and 3 one can notice that the convoltuion kernel does not depend on the spatial position of the input tensor which makes it locationagnostic on the other hand each channel in the output tensor is based on a specific convolution filter which makes is channelspecific involution the idea is to have an operation that is both locationspecific and channelagnostic trying to implement these specific properties poses a challenge with a fixed number of involution kernels for each spatial position we will not be able to process variableresolution input tensors to solve this problem the s have considered generating each kernel conditioned on specific spatial positions with this method we should be able to process variableresolution input tensors with ease the diagram below provides an intuition on this kernel generation method imgurhttps i imgur comjtrggqg png initialize the parameters get the shape of the input scale the height and width with respect to the strides define a layer that average pools the input tensor if stride is more than 1 define the kernel generation layer define reshape layers generate the kernel with respect to the input tensor b h w kkg reshape the kerenl b h w kk 1 g extract input patches b h w kkc reshape the input patches to align with later operations b h w kk cg g compute the multiplyadd operation of kernels and patches b h w kk cg g b h w cg g reshape the output kernel b h w c return the output tensor and the kernel testing the involution layer define the input tensor compute involution with stride 1 compute involution with stride 2 compute involution with stride 1 channel 16 and reduction ratio 2 image classification in this section we will build an imageclassifier model there will be two models one with convolutions and the other with involutions the imageclassification model is heavily inspired by this convolutional neural network cnnhttps www tensorflow orgtutorialsimagescnn tutorial from google get the cifar10 dataset load the cifar10 dataset normalize pixel values to be between 0 and 1 shuffle and batch the dataset visualise the data convolutional neural network build the conv model compile the mode with the necessary loss function and optimizer train the model involutional neural network build the involution model compile the mode with the necessary loss function and optimizer train the model comparisons in this section we will be looking at both the models and compare a few pointers parameters one can see that with a similar architecture the parameters in a cnn is much larger than that of an inn involutional neural network loss and accuracy plots here the loss and the accuracy plots demonstrate that inns are slow learners with lower parameters visualizing involution kernels to visualize the kernels we take the sum of kk values from each involution kernel all the representatives at different spatial locations frame the corresponding heat map the s mention our proposed involution is reminiscent of selfattention and essentially could become a generalized version of it with the visualization of the kernel we can indeed obtain an attention map of the image the learned involution kernels provides attention to individual spatial positions of the input tensor the locationspecific property makes involution a generic space of models in which selfattention belongs the kernel is the second output of the layer conclusions in this example the main focus was to build an involution layer which can be easily reused while our comparisons were based on a specific task feel free to use the layer for different tasks and report your results according to me the key takeaway of involution is its relationship with selfattention the intuition behind locationspecific and channelspefic processing makes sense in a lot of tasks moving forward one can look at yannick s videohttps youtu beph2jzun8moy on involution for a better understanding experiment with the various hyperparameters of the involution layer build different models with the involution layer try building a different kernel generation method altogether you can use the trained model hosted on hugging face hubhttps huggingface cokerasioinvolution and try the demo on hugging face spaceshttps huggingface cospaceskerasioinvolution introduction convolution has been the basis of most modern neural networks for computer vision a convolution kernel is spatial agnostic and channel specific because of this it isn t able to adapt to different visual patterns with respect to different spatial locations along with location related problems the receptive field of convolution creates challenges with regard to capturing long range spatial interactions to address the above issues li et al rethink the properties of convolution in involution inverting the inherence of convolution for visualrecognition https arxiv org abs 2103 06255 the s propose the involution kernel that is location specific and channel agnostic due to the location specific nature of the operation the s say that self attention falls under the design paradigm of involution this example describes the involution kernel compares two image classification models one with convolution and the other with involution and also tries drawing a parallel with the self attention layer setup set seed for reproducibility convolution convolution remains the mainstay of deep neural networks for computer vision to understand involution it is necessary to talk about the convolution operation imgur https i imgur com msklsm5 png consider an input tensor x with dimensions h w and c_in we take a collection of c_out convolution kernels each of shape k k c_in with the multiply add operation between the input tensor and the kernels we obtain an output tensor y with dimensions h w c_out in the diagram above c_out 3 this makes the output tensor of shape h w and 3 one can notice that the convoltuion kernel does not depend on the spatial position of the input tensor which makes it location agnostic on the other hand each channel in the output tensor is based on a specific convolution filter which makes is channel specific involution the idea is to have an operation that is both location specific and channel agnostic trying to implement these specific properties poses a challenge with a fixed number of involution kernels for each spatial position we will not be able to process variable resolution input tensors to solve this problem the s have considered generating each kernel conditioned on specific spatial positions with this method we should be able to process variable resolution input tensors with ease the diagram below provides an intuition on this kernel generation method imgur https i imgur com jtrggqg png initialize the parameters get the shape of the input scale the height and width with respect to the strides define a layer that average pools the input tensor if stride is more than 1 define the kernel generation layer define reshape layers generate the kernel with respect to the input tensor b h w k k g reshape the kerenl b h w k k 1 g extract input patches b h w k k c reshape the input patches to align with later operations b h w k k c g g compute the multiply add operation of kernels and patches b h w k k c g g b h w c g g reshape the output kernel b h w c return the output tensor and the kernel testing the involution layer define the input tensor compute involution with stride 1 compute involution with stride 2 compute involution with stride 1 channel 16 and reduction ratio 2 image classification in this section we will build an image classifier model there will be two models one with convolutions and the other with involutions the image classification model is heavily inspired by this convolutional neural network cnn https www tensorflow org tutorials images cnn tutorial from google get the cifar10 dataset load the cifar10 dataset normalize pixel values to be between 0 and 1 shuffle and batch the dataset visualise the data convolutional neural network build the conv model compile the mode with the necessary loss function and optimizer train the model involutional neural network build the involution model compile the mode with the necessary loss function and optimizer train the model comparisons in this section we will be looking at both the models and compare a few pointers parameters one can see that with a similar architecture the parameters in a cnn is much larger than that of an inn involutional neural network loss and accuracy plots here the loss and the accuracy plots demonstrate that inns are slow learners with lower parameters visualizing involution kernels to visualize the kernels we take the sum of k k values from each involution kernel all the representatives at different spatial locations frame the corresponding heat map the s mention our proposed involution is reminiscent of self attention and essentially could become a generalized version of it with the visualization of the kernel we can indeed obtain an attention map of the image the learned involution kernels provides attention to individual spatial positions of the input tensor the location specific property makes involution a generic space of models in which self attention belongs the kernel is the second output of the layer conclusions in this example the main focus was to build an involution layer which can be easily reused while our comparisons were based on a specific task feel free to use the layer for different tasks and report your results according to me the key take away of involution is its relationship with self attention the intuition behind location specific and channel spefic processing makes sense in a lot of tasks moving forward one can look at yannick s video https youtu be ph2jzun8moy on involution for a better understanding experiment with the various hyperparameters of the involution layer build different models with the involution layer try building a different kernel generation method altogether you can use the trained model hosted on hugging face hub https huggingface co keras io involution and try the demo on hugging face spaces https huggingface co spaces keras io involution
import tensorflow as tf
import keras
import matplotlib.pyplot as plt


tf.random.set_seed(42)






class Involution(keras.layers.Layer):
    def __init__(
        self, channel, group_number, kernel_size, stride, reduction_ratio, name
    ):
        super().__init__(name=name)

        
        self.channel = channel
        self.group_number = group_number
        self.kernel_size = kernel_size
        self.stride = stride
        self.reduction_ratio = reduction_ratio

    def build(self, input_shape):
        
        (_, height, width, num_channels) = input_shape

        
        height = height // self.stride
        width = width // self.stride

        
        
        self.stride_layer = (
            keras.layers.AveragePooling2D(
                pool_size=self.stride, strides=self.stride, padding="same"
            )
            if self.stride > 1
            else tf.identity
        )
        
        self.kernel_gen = keras.Sequential(
            [
                keras.layers.Conv2D(
                    filters=self.channel // self.reduction_ratio, kernel_size=1
                ),
                keras.layers.BatchNormalization(),
                keras.layers.ReLU(),
                keras.layers.Conv2D(
                    filters=self.kernel_size
                    * self.kernel_size
                    * self.group_number,
                    kernel_size=1,
                ),
            ]
        )
        
        self.kernel_reshape = keras.layers.Reshape(
            target_shape=(
                height,
                width,
                self.kernel_size * self.kernel_size,
                1,
                self.group_number,
            )
        )
        self.input_patches_reshape = keras.layers.Reshape(
            target_shape=(
                height,
                width,
                self.kernel_size * self.kernel_size,
                num_channels // self.group_number,
                self.group_number,
            )
        )
        self.output_reshape = keras.layers.Reshape(
            target_shape=(height, width, num_channels)
        )

    def call(self, x):
        
        
        kernel_input = self.stride_layer(x)
        kernel = self.kernel_gen(kernel_input)

        
        
        kernel = self.kernel_reshape(kernel)

        
        
        input_patches = tf.image.extract_patches(
            images=x,
            sizes=[1, self.kernel_size, self.kernel_size, 1],
            strides=[1, self.stride, self.stride, 1],
            rates=[1, 1, 1, 1],
            padding="SAME",
        )

        
        
        input_patches = self.input_patches_reshape(input_patches)

        
        
        output = tf.multiply(kernel, input_patches)
        
        output = tf.reduce_sum(output, axis=3)

        
        
        output = self.output_reshape(output)

        
        return output, kernel





input_tensor = tf.random.normal((32, 256, 256, 3))


output_tensor, _ = Involution(
    channel=3,
    group_number=1,
    kernel_size=5,
    stride=1,
    reduction_ratio=1,
    name="inv_1",
)(input_tensor)
print(f"with stride 1 ouput shape: {output_tensor.shape}")


output_tensor, _ = Involution(
    channel=3,
    group_number=1,
    kernel_size=5,
    stride=2,
    reduction_ratio=1,
    name="inv_2",
)(input_tensor)
print(f"with stride 2 ouput shape: {output_tensor.shape}")


output_tensor, _ = Involution(
    channel=16,
    group_number=1,
    kernel_size=5,
    stride=1,
    reduction_ratio=2,
    name="inv_3",
)(input_tensor)
print(
    "with channel 16 and reduction ratio 2 ouput shape: {}".format(
        output_tensor.shape
    )
)






print("loading the CIFAR10 dataset...")
(
    (train_images, train_labels),
    (
        test_images,
        test_labels,
    ),
) = keras.datasets.cifar10.load_data()


(train_images, test_images) = (train_images / 255.0, test_images / 255.0)


train_ds = (
    tf.data.Dataset.from_tensor_slices((train_images, train_labels))
    .shuffle(256)
    .batch(256)
)
test_ds = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(
    256
)



class_names = [
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
]

plt.figure(figsize=(10, 10))
for i in range(25):
    plt.subplot(5, 5, i + 1)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)
    plt.imshow(train_images[i])
    plt.xlabel(class_names[train_labels[i][0]])
plt.show()




print("building the convolution model...")
conv_model = keras.Sequential(
    [
        keras.layers.Conv2D(
            32, (3, 3), input_shape=(32, 32, 3), padding="same"
        ),
        keras.layers.ReLU(name="relu1"),
        keras.layers.MaxPooling2D((2, 2)),
        keras.layers.Conv2D(64, (3, 3), padding="same"),
        keras.layers.ReLU(name="relu2"),
        keras.layers.MaxPooling2D((2, 2)),
        keras.layers.Conv2D(64, (3, 3), padding="same"),
        keras.layers.ReLU(name="relu3"),
        keras.layers.Flatten(),
        keras.layers.Dense(64, activation="relu"),
        keras.layers.Dense(10),
    ]
)


print("compiling the convolution model...")
conv_model.compile(
    optimizer="adam",
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
    jit_compile=False,
)


print("conv model training...")
conv_hist = conv_model.fit(train_ds, epochs=20, validation_data=test_ds)




print("building the involution model...")

inputs = keras.Input(shape=(32, 32, 3))
x, _ = Involution(
    channel=3,
    group_number=1,
    kernel_size=3,
    stride=1,
    reduction_ratio=2,
    name="inv_1",
)(inputs)
x = keras.layers.ReLU()(x)
x = keras.layers.MaxPooling2D((2, 2))(x)
x, _ = Involution(
    channel=3,
    group_number=1,
    kernel_size=3,
    stride=1,
    reduction_ratio=2,
    name="inv_2",
)(x)
x = keras.layers.ReLU()(x)
x = keras.layers.MaxPooling2D((2, 2))(x)
x, _ = Involution(
    channel=3,
    group_number=1,
    kernel_size=3,
    stride=1,
    reduction_ratio=2,
    name="inv_3",
)(x)
x = keras.layers.ReLU()(x)
x = keras.layers.Flatten()(x)
x = keras.layers.Dense(64, activation="relu")(x)
outputs = keras.layers.Dense(10)(x)

inv_model = keras.Model(inputs=[inputs], outputs=[outputs], name="inv_model")


print("compiling the involution model...")
inv_model.compile(
    optimizer="adam",
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
    jit_compile=False,
)


print("inv model training...")
inv_hist = inv_model.fit(train_ds, epochs=20, validation_data=test_ds)





conv_model.summary()

inv_model.summary()



plt.figure(figsize=(20, 5))

plt.subplot(1, 2, 1)
plt.title("Convolution Loss")
plt.plot(conv_hist.history["loss"], label="loss")
plt.plot(conv_hist.history["val_loss"], label="val_loss")
plt.legend()

plt.subplot(1, 2, 2)
plt.title("Involution Loss")
plt.plot(inv_hist.history["loss"], label="loss")
plt.plot(inv_hist.history["val_loss"], label="val_loss")
plt.legend()

plt.show()

plt.figure(figsize=(20, 5))

plt.subplot(1, 2, 1)
plt.title("Convolution Accuracy")
plt.plot(conv_hist.history["accuracy"], label="accuracy")
plt.plot(conv_hist.history["val_accuracy"], label="val_accuracy")
plt.legend()

plt.subplot(1, 2, 2)
plt.title("Involution Accuracy")
plt.plot(inv_hist.history["accuracy"], label="accuracy")
plt.plot(inv_hist.history["val_accuracy"], label="val_accuracy")
plt.legend()

plt.show()



layer_names = ["inv_1", "inv_2", "inv_3"]

outputs = [inv_model.get_layer(name).output[1] for name in layer_names]
vis_model = keras.Model(inv_model.input, outputs)

fig, axes = plt.subplots(nrows=10, ncols=4, figsize=(10, 30))

for ax, test_image in zip(axes, test_images[:10]):
    inv_out = vis_model.predict(test_image[None, ...])
    inv1_kernel, inv2_kernel, inv3_kernel = inv_out

    inv1_kernel = tf.reduce_sum(inv1_kernel, axis=[-1, -2, -3])
    inv2_kernel = tf.reduce_sum(inv2_kernel, axis=[-1, -2, -3])
    inv3_kernel = tf.reduce_sum(inv3_kernel, axis=[-1, -2, -3])

    ax[0].imshow(keras.utils.array_to_img(test_image))
    ax[0].set_title("Input Image")

    ax[1].imshow(keras.utils.array_to_img(inv1_kernel[0, ..., None]))
    ax[1].set_title("Involution Kernel 1")

    ax[2].imshow(keras.utils.array_to_img(inv2_kernel[0, ..., None]))
    ax[2].set_title("Involution Kernel 2")

    ax[3].imshow(keras.utils.array_to_img(inv3_kernel[0, ..., None]))
    ax[3].set_title("Involution Kernel 3")