Python 如何在Tensorflow 2中使自定义非线性卷积层快速？

我需要在tensorflow2中创建类似卷积的层（然后将图形导出到神经加速器）。 滤波器采用7x7面积，产生一个像素，但它是非线性的，使用abs函数，因此不可能直接利用优化的nn.conv2d函数

我试着让自定义层，代码可以工作，但是对于过滤器来说，它的代码非常慢（一个图像行10-20秒！），应该只比卷积运算慢一点。当我试图以同样的方式制作卷积层时，它也太慢了。下面代码中实现过滤器的最佳方法是什么

这是非线性且缓慢的过滤器的最小示例：

import tensorflow as tf
import numpy as np

class FilterLayer (tf.keras.layers.Layer) :
def __init__(self) :
    super(FilterLayer, self).__init__()
    self.offset = 4
    
def build(self, input_shape) :
    pass


def call(self, input): 
    output = tf.Variable(np.zeros((input.shape[0] - self.offset * 2, input.shape[1] - self.offset * 2)), dtype=tf.uint8)
    
    a = tf.Variable(1, dtype=tf.float32)
    s = tf.Variable(2, dtype=tf.float32)
    d = tf.Variable(3, dtype=tf.float32)
    f = tf.Variable(4, dtype=tf.float32)
    g = tf.Variable(5, dtype=tf.float32)
    for y in range(self.offset, input.shape[0] - self.offset) :
        for x in range(self.offset, input.shape[1] - self.offset) :

            q = tf.cast(input[y][x], dtype=tf.float32)
            w = tf.cast(input[y+1][x+1], dtype=tf.float32)
            e = tf.cast(input[y-1][x+1], dtype=tf.float32)
            r = tf.cast(input[y+1][x-1], dtype=tf.float32)
            t = tf.cast(input[y-1][x-1], dtype=tf.float32)
            summ = a*q+tf.abs(s*w-d*e)*tf.abs(f*r-g*t)
            tf.tensor_scatter_nd_update(output, [(y-4, x-4)], [summ])
        print(y)
    return output
q = FilterLayer()
q(tf.Variable(np.ones((3000, 3000)), dtype=tf.uint8))

相比之下，普通卷积几乎立即执行（示例取自文档，它不同于此过滤器）

另外，我不需要对这个滤波器使用反向传播或训练，所有系数都是已知的

x_in = np.ones((3000, 3000), dtype=np.uint8)
print(x_in.shape)
x_in = x_in[np.newaxis][..., np.newaxis]
print(x_in.shape)

kernel_in = np.array([
 [ [[2, 0.1]], [[3, 0.2]] ],
 [ [[0, 0.3]],[[1, 0.4]] ], ])

x = tf.constant(x_in, dtype=tf.float32)
kernel = tf.constant(kernel_in, dtype=tf.float32)
out_res = tf.nn.conv2d(x, kernel, strides=[1, 1, 1, 1], padding='VALID')
print(out_res.shape)