How to initialize a Conv2D layer with predetermined list of kernels in tensorflow/keras?

Question

I would like to use Conv2D layer in order to stride an input image and run three 2x2 kernels.

This is not the purpose of tensorflow, but I really want to use tensorflow as the backend engine to run kernels efficiently and to distribute the work load between different devices GPUs and/or CPUs.

I tried something like the following code. But It doesn't seem to work very well.

import tensorflow as tf

class InitConvKernels(tf.keras.initializers.Initializer):

  def __init__(self, num_kernels, kernel_tensor):
    self.kernel_list= kernel_tensor
    self.index = -1
    self.num_kernels = num_kernels

  def __call__(self, shape, dtype=None):
    index += 1 
    assert(self.index <= self.num_kernels) # doesn't affect anything
    tf.print(shape) # doesn't work
    return self.kernel_list[index]

  def get_config(self):
    return {'kernel_list': self.kernel_list, 'num_kernels': self.num_kernels}

I am calling the custom initializer, but the returned layer is empty:

kernel_list = tf.constant([[[-1, -1],  [-1, -1]], [[1, 1],   [1, 1]],  [[-1, 1],  [1, -1]],])
layer = layers.Conv2D(
    filters=3,
    kernel_size=2,
    kernel_initializer=InitConvKernels(3,kernel_list),
    bias_initializer=initializers.Zeros()
)

layer.variables is empty ([]) layer.layer.get_weights() is also empty ([])

My goal is to evaluate the convolution of the three kernels in kernel_list on an input image and aggregate all the results.

Answer 1

from PIL import Image
import requests
from io import BytesIO
import numpy as np
from keras.models import Sequential
from keras.layers import Conv2D


response = requests.get('https://upload.wikimedia.org/wikipedia/commons/thumb/0/02/Stack_Overflow_logo.svg/1280px-Stack_Overflow_logo.svg.png')
image = Image.open(BytesIO(response.content))

Loading an image from url.

Build a model to run a kernel (to run more kernels make kernel_init a generator and readily adjust the number of filters when initializing Conv2D)

def kernel_init(shape, dtype=None, partition_info=None):
    kernel = np.zeros(shape)
    kernel[:,:,0,0] = np.array([[1,0,1],[-1,0,-1],[1,0,1]])
    return kernel

#Build Keras model
model = Sequential()
model.add(Conv2D(1, [3,3], kernel_initializer=kernel_init, 
                 input_shape=(251,1280,4), padding="valid"))
model.build()

# To apply existing filter, we use predict with no training
out = model.predict(image)

And visualizing the output:

import matplotlib.pyplot as plt
plt.matshow(out[0,:,:,0])

Edit: It's worth mentioning OpenAI's Triton which can help using higher level language and framework such as pytorch to run efficient GPU code:

Python-like programming language which enables researchers with no CUDA experience to write highly efficient GPU code—most of the time on par with what an expert would be able to produce.