Оптимизация пользовательского слоя conv2d в tf.keras

Question

Я пытаюсь эффективно реализовать следующий вид слоя conv2d.Текущая реализация, на мой взгляд, работает, но очень неэффективна.

Входной тензор размера

(размер партии x Ш x В x C_in)

Тензор вывода

(batch_size x W x H x C_out)

Слой принимает два параметра: количество единиц (C_u) и список из K ядер ядер (известных заранее).Каждое ядро ​​конвоя имеет размер (W, H, 1, N), где N - количество выходных каналов (в каналах 1).Обратите внимание, что разные ядра в одном и том же подобии имеют разные N!

Сначала мы применяем плотно связанный слой (обучаемый), преобразуя входную форму в

(пакетный размер x W x H x C_u)

Затем я хочу применить каждое из сверточных ядер к каждому из каналов.

В результате получается C_u x K x (размер пакета x x x x N)

I тогдахочу взять максимум вдоль N (поэтому я получаю (размер партии x Ш x В x 1)) и объединить все, чтобы получить

(размер партии x Ш x В x (C_u x K))

(т. е. C_out = C_u x K)

Вот один из способов реализовать это, но время тренировки очень медленное, и это не очень хорошо при установке на GPU:

import tensorflow as tf
from tensorflow.keras import layers

class fixedConvLayer(layers.Dense):

    def __init__(self, units, conv_kernels, **params):
        params['units']=units
        self.conv_kernels_numpy = conv_kernels
        super().__init__(**params)
        return

    def build(self, input_shape):

        super().build(input_shape)
        self.conv_kernels = [tf.convert_to_tensor(np.reshape(kernels,[3,3,1,-1]))                                                           
                                  for kernels in self.conv_kernels_numpy]

        return 

    def comp_filters(self,channel):

        return tf.concat([
                    tf.math.reduce_max(tf.nn.conv2d(channel,
                                      filter=kernel,
                                      strides=1,
                                      padding='SAME'),axis=3,keepdims=True)

                        for kernel in self.conv_kernels],axis=3)


    def call(self,inputs):

        #take from Dense definition and slightly modify
        inputs = tf.convert_to_tensor(inputs)
        rank = tf.rank(inputs)
        if rank != 4:
            assert 'Rank expected to be 4'


        # Broadcasting is required for the inputs.
        outputs = tf.tensordot(inputs, self.kernel, [[3], [0]])

        # Reshape the output back to the original ndim of the input.
        shape = inputs.shape.as_list()
        output_shape = shape[:-1] + [self.units]
        outputs.set_shape(output_shape)

        if self.use_bias:
            outputs = tf.nn.bias_add(outputs, self.bias)
        if self.activation is not None:
            outputs = self.activation(outputs)  

        #apply the conv filters
        channel_list = tf.split(outputs,num_or_size_splits= self.units,axis = -1)
        max_layers = tf.concat([self.comp_filters(channel) for channel in channel_list],axis=3)

        return max_layers