Question

我正在尝试有效地实现以下类型的conv2d层。我认为当前的实现可行，但效率很低。

输入张量大小

（批量大小x宽x高x C_in）

输出张量

（批量大小x宽x高x C_out）

该层采用两个参数，单位数（C_u）和K个conv内核的列表（提前知道）。每个转换内核的大小为（W，H，1，N），其中N是输出通道数（输入通道为1）。注意，相同的不同内核具有不同的Ns！

首先，我们将紧密连接的图层（可训练）应用于将输入形状转换为

（批量大小x宽x高x C_u）

然后，我想将每个卷积内核应用于每个通道。

这将导致C_u x K x（批量大小x W x H x N）

然后我想沿N取一个最大值（所以我得到了（batch_size x W x H 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

在tf.keras中优化自定义conv2d层

0 个答案: