How to addin new Tensorflow layers to TensorRT engine?

Question

I have tensorflow freezed model from which TensorRT engine is produced.

I can't retrain model since I don't have all those required images.

But Tensorflow process has some post processing layers and I like to add into TensorRT engine.

What would be the best approach?

Can I create plugin layer using TensorRT layers?

Those Tensorflow layers are mostly available in TensorRT as follows.

self.tensor_heatMat_up = tf.image.resize_area(self.tensor_output[:, :, :, :19], self.upsample_size,
                                                      align_corners=False, name='upsample_heatmat')
        self.tensor_pafMat_up = tf.image.resize_area(self.tensor_output[:, :, :, 19:], self.upsample_size,
                                                     align_corners=False, name='upsample_pafmat')
        if trt_bool is True:
            smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0, 19)
        else:
            smoother = Smoother({'data': self.tensor_heatMat_up}, 25, 3.0)
        gaussian_heatMat = smoother.get_output()

        max_pooled_in_tensor = tf.nn.pool(gaussian_heatMat, window_shape=(3, 3), pooling_type='MAX', padding='SAME')
        self.tensor_peaks = tf.where(tf.equal(gaussian_heatMat, max_pooled_in_tensor), gaussian_heatMat,
tf.zeros_like(gaussian_heatMat))

TensorRT has scale for resize_area, conv for Smoother. Not sure tf.equal in TensorRT.

How to addin those layers to TensorRT? Possible to use graphsurgeon or UFF model?

Answer 1

The following steps add a custom plugin layer in C++ for TensorFlow networks:

1. Implement the IPluginV2 and IPluginCreator classes as shown in: Adding A Custom Layer Using C++ For Caffe.
2. Map the TensorFlow operation to the plugin operation. You can use GraphSurgeon for this.

3. Call the UFF converter with the preprocess -p flag set. This will generate a UFF file with the TensorFlow operations replaced by TensorRT plugin nodes. convert-to-uff frozen_inference_graph.pb -p config.py -t

4. Run the pre-processed and converted UFF file with TensorRT using the UFF parser. For details, see Using Custom Layers When Importing A Model From A Framework. The Object Detection With A TensorFlow SSD Network sample illustrates how to add a custom layer that is not supported in UFF using C++. See config.py in the sample folder for a demonstration of how to pre-process the graph.

Although the C++ API is the preferred language to implement custom layers; due to easily accessing libraries like CUDA and cuDNN, you can also work with custom layers in Python applications. You can follow Adding Custom Layers Using The Python API guide.