MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / bindPlugin

Function bindPlugin

python/src/infer/pyPlugin.cpp:157–277  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

155} // namespace lambdas
156
157void bindPlugin(py::module& m)
158{
159 py::class_<IPluginV2>(m, "IPluginV2", IPluginV2Doc::descr, py::module_local())
160 .def_property_readonly("num_outputs", &IPluginV2::getNbOutputs)
161 .def_property_readonly("tensorrt_version", &IPluginV2::getTensorRTVersion)
162 .def_property_readonly("plugin_type", &IPluginV2::getPluginType)
163 .def_property_readonly("plugin_version", &IPluginV2::getPluginVersion)
164 .def_property("plugin_namespace", &IPluginV2::getPluginNamespace,
165 py::cpp_function(&IPluginV2::setPluginNamespace, py::keep_alive<1, 2>{}))
166 .def("get_output_shape", lambdas::IPluginV2_get_output_shape, "index"_a, "input_shapes"_a,
167 IPluginV2Doc::get_output_shape)
168 .def("supports_format", &IPluginV2::supportsFormat, "dtype"_a, "format"_a, IPluginV2Doc::supports_format)
169 .def("configure_with_format", lambdas::IPluginV2_configure_with_format, "input_shapes"_a, "output_shapes"_a,
170 "dtype"_a, "format"_a, "max_batch_size"_a, IPluginV2Doc::configure_with_format)
171 .def("initialize", &IPluginV2::initialize, IPluginV2Doc::initialize)
172 .def("terminate", &IPluginV2::terminate, IPluginV2Doc::terminate)
173 .def("get_workspace_size", &IPluginV2::getWorkspaceSize, "max_batch_size"_a, IPluginV2Doc::get_workspace_size)
174 .def("execute_async", lambdas::IPluginV2_execute_async, "batch_size"_a, "inputs"_a, "outputs"_a, "workspace"_a,
175 "stream_handle"_a, IPluginV2Doc::execute_async)
176 .def_property_readonly("serialization_size", &IPluginV2::getSerializationSize)
177 .def(
178 "serialize", lambdas::IPluginV2_serialize, IPluginV2Doc::serialize, py::return_value_policy::take_ownership)
179 .def("destroy", &IPluginV2::destroy, IPluginV2Doc::destroy)
180 .def("clone", &IPluginV2::clone, IPluginV2Doc::clone);
181
182 py::class_<IPluginV2Ext, IPluginV2>(m, "IPluginV2Ext", IPluginV2ExtDoc::descr, py::module_local())
183 .def("get_output_data_type", lambdas::get_output_data_type, "index"_a, "input_types"_a,
184 IPluginV2ExtDoc::get_output_data_type)
185 .def("configure_plugin", lambdas::configure_plugin, "input_shapes"_a, "output_shapes"_a, "input_types"_a,
186 "output_types"_a, "input_is_broadcasted"_a, "output_is_broacasted"_a, "format"_a, "max_batch_size"_a,
187 IPluginV2ExtDoc::configure_plugin)
188 .def("attach_to_context", lambdas::attach_to_context, "cudnn"_a, "cublas"_a, "allocator"_a,
189 IPluginV2ExtDoc::attach_to_context)
190 .def("detach_from_context", &IPluginV2Ext::detachFromContext, IPluginV2ExtDoc::detach_from_context)
191 .def("clone", &IPluginV2Ext::clone, IPluginV2ExtDoc::clone);
192 ;
193
194 py::enum_<PluginFieldType>(m, "PluginFieldType", PluginFieldTypeDoc::descr, py::module_local())
195 .value("FLOAT16", PluginFieldType::kFLOAT16)
196 .value("FLOAT32", PluginFieldType::kFLOAT32)
197 .value("FLOAT64", PluginFieldType::kFLOAT64)
198 .value("INT8", PluginFieldType::kINT8)
199 .value("INT16", PluginFieldType::kINT16)
200 .value("INT32", PluginFieldType::kINT32)
201 .value("CHAR", PluginFieldType::kCHAR)
202 .value("DIMS", PluginFieldType::kDIMS)
203 .value("UNKNOWN", PluginFieldType::kUNKNOWN);
204
205 py::class_<PluginField>(m, "PluginField", PluginFieldDoc::descr, py::module_local())
206 .def(py::init(lambdas::plugin_field_default_constructor), "name"_a = "", py::keep_alive<1, 2>{})
207 .def(py::init(lambdas::plugin_field_constructor), "name"_a, "data"_a,
208 "type"_a = nvinfer1::PluginFieldType::kUNKNOWN, py::keep_alive<1, 2>{}, py::keep_alive<1, 3>{})
209 .def_property(
210 "name", [](PluginField& self) { return self.name; },
211 py::cpp_function(
212 [](PluginField& self, FallbackString& name) { self.name = name.c_str(); }, py::keep_alive<1, 2>{}))
213 .def_property(
214 "data", [](PluginField& self) { return self.data; },

Callers 1

PYBIND11_MODULEFunction · 0.85

Calls 1

c_strMethod · 0.80

Tested by

no test coverage detected