Implementing SKFF with Tensorflow

[IwakuraRein]

Hi. I tried to implement SKFF with tf 1.15:

def SKFF(self, inputs:list, reduction=8, name='SKFF'):
    with tf.variable_scope(name):
        ch_n=inputs[0].shape[3]
        num=len(inputs)
        d=max(ch_n//reduction, 4)
        
        inputs=tf.stack(inputs, 0)
        fea=tf.reduce_sum(inputs, 0)
        fea=tf.reduce_mean(fea, [1, 2], keep_dims=True)
        fea=self.conv_layer(fea, d, 1, name='du')
        fea=tf.keras.layers.PReLU()(fea)
        
        vecs=[self.conv_layer(fea, ch_n, 1, name=str(no)) for no in range(num)]
        vec=tf.concat(vecs,axis=1)
        weight=tf.nn.softmax(vec, axis=1)
        weight=tf.transpose(weight, (1, 0, 2, 3))
        weight=tf.expand_dims(weight, 2)
        out = inputs*weight
        out=tf.reduce_sum(out, 0)
        return out

Then I used timeline to profile my network. I noticed that there were lots of transpose operations (i.e., convert data from NHWC to NCHW) so the inference speed was actually slower than direct contact different scales.

Is there any way I can optimize the TensorFlow codes? Thanks.