| 62 | } // namespace lambdas |
| 63 | |
| 64 | void bindCaffe(py::module& m) |
| 65 | { |
| 66 | py::class_<IBlobNameToTensor, std::unique_ptr<IBlobNameToTensor, py::nodelete>>( |
| 67 | m, "IBlobNameToTensor", IBlobNameToTensorDoc::descr, py::module_local()) |
| 68 | .def("find", &IBlobNameToTensor::find, "name"_a, IBlobNameToTensorDoc::find); |
| 69 | |
| 70 | py::class_<IPluginFactoryV2>(m, "ICaffePluginFactoryV2", ICaffePluginFactoryV2Doc::descr, py::module_local()) |
| 71 | .def("is_plugin_v2", &IPluginFactoryV2::isPluginV2, "layer_name"_a, ICaffePluginFactoryV2Doc::is_plugin_v2) |
| 72 | .def("create_plugin", lambdas::PluginV2_create_plugin, "layer_name"_a, "weights"_a, py::keep_alive<1, 3>{}, |
| 73 | ICaffePluginFactoryV2Doc::create_plugin); |
| 74 | |
| 75 | py::class_<ICaffeParser, std::unique_ptr<ICaffeParser, py::nodelete>>( |
| 76 | m, "CaffeParser", ICaffeParserDoc::descr, py::module_local()) |
| 77 | .def(py::init(&nvcaffeparser1::createCaffeParser)) |
| 78 | .def_property("protobuf_buffer_size", nullptr, &ICaffeParser::setProtobufBufferSize) |
| 79 | .def_property( |
| 80 | "plugin_factory_v2", nullptr, py::cpp_function(&ICaffeParser::setPluginFactoryV2, py::keep_alive<1, 2>{})) |
| 81 | .def_property( |
| 82 | "plugin_namespace", nullptr, py::cpp_function(&ICaffeParser::setPluginNamespace, py::keep_alive<1, 2>{})) |
| 83 | .def("parse", &ICaffeParser::parse, "deploy"_a, "model"_a, "network"_a, "dtype"_a, ICaffeParserDoc::parse, |
| 84 | py::keep_alive<4, 1>{}) |
| 85 | .def("parse_buffer", lambdas::parse_buffer, "deploy_buffer"_a, "model_buffer"_a, "network"_a, "dtype"_a, |
| 86 | ICaffeParserDoc::parse_buffer, py::keep_alive<4, 1>{}) |
| 87 | .def("parse_binary_proto", lambdas::parse_binary_proto, "filename"_a, ICaffeParserDoc::parse_binary_proto) |
| 88 | .def_property("error_recorder", &ICaffeParser::getErrorRecorder, |
| 89 | py::cpp_function(&ICaffeParser::setErrorRecorder, py::keep_alive<1, 2>{})) |
| 90 | .def("__del__", &utils::doNothingDel<ICaffeParser>); |
| 91 | |
| 92 | m.def("shutdown_protobuf_library", &nvcaffeparser1::shutdownProtobufLibrary, |
| 93 | FreeFunctionsDoc::shutdown_protobuf_library); |
| 94 | } |
| 95 | } // namespace tensorrt |