Creating python binding for a simple IPlugin layer

AI & Data Science Deep Learning (Training & Inference) TensorRT
1

Hello,

I have followed through the IPlugin python example (FCPlugin) which basically implements a layer in C++ and then generate python bindings for it and then use those bindings to generate plugin using caffeparser.

class definitions:

class FCPlugin: public nvinfer1::IPluginExt

class FCPluginFactory : public nvcaffeparser1::IPluginFactoryExt, public nvinfer1::IPluginFactory

pybind11 binding:

py::class_<FCPlugin, nvinfer1::IPluginExt, std::unique_ptr<FCPlugin, py::nodelete>>(m, "FCPlugin")
        // Bind the normal constructor as well as the one which deserializes the plugin
        .def(py::init<const nvinfer1::Weights*, int>())
        .def(py::init<const void*, size_t>())
    ;

    // Since the createPlugin function overrides IPluginFactory functionality, we do not need to explicitly bind it here.
    // We specify py::multiple_inheritance because we have not explicitly specified nvinfer1::IPluginFactory as a base class.
    py::class_<FCPluginFactory, nvcaffeparser1::IPluginFactoryExt>(m, "FCPluginFactory", py::multiple_inheritance())
        // Bind the default constructor.
        .def(py::init<>())
        // The destroy_plugin function does not override the base class, so we must bind it explicitly.
        .def("destroy_plugin", &FCPluginFactory::destroyPlugin)
    ;

Call from python as in tutorial example:

... trt.CaffeParser() as parser:
        fc_factory = fcplugin.FCPluginFactory()
        parser.plugin_factory_ext = fc_factory

I intend to bypass caffeparser and instead directly define the network in python and call my custom layer as plugin. My code structure is as follows:

Class definitions:

class Toy : public nvinfer1::IPlugin

class ToyFactory : public nvinfer1::IPluginFactory

pybind11 code segment

py::class_<Toy, nvinfer1::IPlugin, std::unique_ptr<Toy, py::nodelete>>(m, "Toy")
		.def(py::init<>())
	;

	py::class_<ToyFactory, nvinfer1::IPluginFactory, std::unique_ptr<ToyFactory, py::nodelete>>(m, "ToyFactory", py::multiple_inheritance())
		.def(py::init<>())
	;

I am able to build the python module. But I am stuck at how to register/call the plugin. In caffe example you are registering the pluginfactory with parser directly

... trt.CaffeParser() as parser:
        fc_factory = fcplugin.FCPluginFactory()
        parser.plugin_factory_ext = fc_factory

I intend to register the plugin directly with network as is done in LeakyReLU example

trt.register_all_tensorrt_plugins()
PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list

def get_trt_plugin(plugin_name):
        plugin = None
        for plugin_creator in PLUGIN_CREATORS:
            if plugin_creator.name == plugin_name:
                lrelu_slope_field = trt.PluginField("neg_slope", np.array([0.1], dtype=np.float32), trt.PluginFieldType.FLOAT32)
                field_collection = trt.PluginFieldCollection([lrelu_slope_field])
                plugin = plugin_creator.create_plugin(name=plugin_name, field_collection=field_collection)
        return plugin

lrelu = network.add_plugin_ext(inputs=[input_layer], plugin=get_trt_plugin("LReLU_TRT"))
        lrelu.get_output(0).name = "outputs"

But I cannot figure out how to call/register the plugin correctly. Can you guys help?