Swish Plugin Layer on TensorRT


I am trying to implement a plugin layer for swish activation function in TensorRT. The model was initially trained on keras and was converted to UFF format using uff converter in python. A custom config.py was used in the conversion process. Kindly note that the network only has a single unsupported node which is swish activation (https://tensorlayer.readthedocs.io/en/latest/modules/activation.html).

import graphsurgeon as gs

import tensorflow as tf

Input = gs.create_node("Input",
    shape=[1, 3, 224, 224])

swish_fp32_list = []
for i in range(49):
    layer_name = "swish_{}/swish_f32".format(i+1)
    swish_fp32 = gs.create_plugin_node(name=layer_name, op="Swish_TRT",dtype=tf.float32)

namespace_plugin_map = {

for i in range(49):
    layer_name = "swish_{}/swish_f32".format(i+1)
    namespace_plugin_map[layer_name] = swish_fp32_list[i]

def preprocess(dynamic_graph):
    #dynamic_graph.remove(dynamic_graph.find_nodes_by_path(namespace_remove), remove_exclusive_dependencies=True)
    #dynamic_graph.remove(dynamic_graph.graph_outputs, remove_exclusive_dependencies=False)

While trying to parse the above-converted uff file, I am coming across the following error :

ERROR: multiply_1/mul: elementwise inputs must have same dimensions or follow broadcast rules (input dimensions were [32,1,1] and [112,112,32])
ERROR: conv2d_4/convolution: at least three non-batch dimensions are required for input
ERROR: UFFParser: Parser error: batch_normalization_3/FusedBatchNorm_1: The input to the Scale Layer is required to have a minimum of 3 dimensions.

The input and output of the swish layer are of shape 32 x 112 x112. When it is passed to the next layer, which is Reduce mean layer, TensorRT shuffles the node to reshape it to NHWC format (112x112x32)

This shuffling is not needed for the depthwise conv layer which uses swish activation as input and hence the above error is generated.

My question is why does UFF converter insert a shuffle layer into the node to convert it to NHWC format. How can I reshuffle it manually? Any help would be deeply appreciated.
Thanks in advance.

Linux distro and version - Ubuntu 18.04
GPU type - Titan Xp
nvidia driver version - 418.39
CUDA version - 10.0
Python version - 3.6
Tensorflow version - 1.13
TensorRT version - 5.0
Device - Desktop

Met the same question, you can print all layers and you’ll find that there are some unnamed shuffle layer, I guess it is because TensorRT prefer NHWC in mean-reduce but also send shuffled data to other path by mistake.
You can reshuffle it by adding a new custom_node, than create a shuffled_data-> custom_node->multiply_1/mul path and reshuffle it in custom_node, you can also adding a custom_node to make sure that the uff parsing sucess and then create a path from unshuffled_data -> mul layer