You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

main.cpp 60 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314
  1. // Copyright 2020 Tencent
  2. // SPDX-License-Identifier: BSD-3-Clause
  3. #include <pybind11/pybind11.h>
  4. #include <pybind11/stl.h>
  5. #include <pybind11/numpy.h>
  6. #include <pybind11/functional.h>
  7. #include <cpu.h>
  8. #include <gpu.h>
  9. #include <net.h>
  10. #include <option.h>
  11. #include <blob.h>
  12. #include <paramdict.h>
  13. #include "pybind11_mat.h"
  14. #include "pybind11_datareader.h"
  15. #include "pybind11_allocator.h"
  16. #include "pybind11_modelbin.h"
  17. #include "pybind11_layer.h"
  18. using namespace ncnn;
  19. namespace py = pybind11;
  20. class DataReaderFromMemoryCopy : public DataReaderFromMemory
  21. {
  22. public:
  23. explicit DataReaderFromMemoryCopy(const unsigned char*& mem)
  24. : DataReaderFromMemory(mem)
  25. {
  26. }
  27. virtual size_t reference(size_t size, const void** buf) const
  28. {
  29. return 0;
  30. }
  31. };
  32. struct LayerFactory
  33. {
  34. std::string name;
  35. int index;
  36. std::function<Layer*()> creator;
  37. std::function<void(Layer*)> destroyer;
  38. layer_creator_func creator_func;
  39. layer_destroyer_func destroyer_func;
  40. };
  41. #define LayerFactoryDeclear(n) \
  42. static ncnn::Layer* LayerCreator##n(void*); \
  43. static void LayerDestroyer##n(ncnn::Layer*, void*);
  44. LayerFactoryDeclear(0);
  45. LayerFactoryDeclear(1);
  46. LayerFactoryDeclear(2);
  47. LayerFactoryDeclear(3);
  48. LayerFactoryDeclear(4);
  49. LayerFactoryDeclear(5);
  50. LayerFactoryDeclear(6);
  51. LayerFactoryDeclear(7);
  52. LayerFactoryDeclear(8);
  53. LayerFactoryDeclear(9);
  54. std::vector<LayerFactory> g_layer_factroys = {
  55. {"", -1, nullptr, nullptr, LayerCreator0, LayerDestroyer0},
  56. {"", -1, nullptr, nullptr, LayerCreator1, LayerDestroyer1},
  57. {"", -1, nullptr, nullptr, LayerCreator2, LayerDestroyer2},
  58. {"", -1, nullptr, nullptr, LayerCreator3, LayerDestroyer3},
  59. {"", -1, nullptr, nullptr, LayerCreator4, LayerDestroyer4},
  60. {"", -1, nullptr, nullptr, LayerCreator5, LayerDestroyer5},
  61. {"", -1, nullptr, nullptr, LayerCreator6, LayerDestroyer6},
  62. {"", -1, nullptr, nullptr, LayerCreator7, LayerDestroyer7},
  63. {"", -1, nullptr, nullptr, LayerCreator8, LayerDestroyer8},
  64. {"", -1, nullptr, nullptr, LayerCreator9, LayerDestroyer9},
  65. };
  66. int g_layer_factroy_index = 0;
  67. #define LayerFactoryDefine(n) \
  68. static ncnn::Layer* LayerCreator##n(void* p) \
  69. { \
  70. if (g_layer_factroys[n].creator != nullptr) \
  71. { \
  72. return g_layer_factroys[n].creator(); \
  73. } \
  74. return nullptr; \
  75. } \
  76. static void LayerDestroyer##n(ncnn::Layer* layer, void* p) \
  77. { \
  78. if (g_layer_factroys[n].destroyer) \
  79. { \
  80. g_layer_factroys[n].destroyer(layer); \
  81. } \
  82. }
  83. LayerFactoryDefine(0);
  84. LayerFactoryDefine(1);
  85. LayerFactoryDefine(2);
  86. LayerFactoryDefine(3);
  87. LayerFactoryDefine(4);
  88. LayerFactoryDefine(5);
  89. LayerFactoryDefine(6);
  90. LayerFactoryDefine(7);
  91. LayerFactoryDefine(8);
  92. LayerFactoryDefine(9);
  93. PYBIND11_MODULE(ncnn, m)
  94. {
  95. auto atexit = py::module_::import("atexit");
  96. atexit.attr("register")(py::cpp_function([]() {
  97. for (int i = 0; i < g_layer_factroys.size(); i++)
  98. {
  99. g_layer_factroys[i].creator = nullptr;
  100. g_layer_factroys[i].destroyer = nullptr;
  101. }
  102. }));
  103. py::class_<Allocator, PyAllocator<> >(m, "Allocator");
  104. py::class_<PoolAllocator, Allocator, PyAllocatorOther<PoolAllocator> >(m, "PoolAllocator")
  105. .def(py::init<>())
  106. .def("set_size_compare_ratio", &PoolAllocator::set_size_compare_ratio, py::arg("src"))
  107. .def("clear", &PoolAllocator::clear)
  108. .def("fastMalloc", &PoolAllocator::fastMalloc, py::arg("size"))
  109. .def("fastFree", &PoolAllocator::fastFree, py::arg("ptr"));
  110. py::class_<UnlockedPoolAllocator, Allocator, PyAllocatorOther<UnlockedPoolAllocator> >(m, "UnlockedPoolAllocator")
  111. .def(py::init<>())
  112. .def("set_size_compare_ratio", &UnlockedPoolAllocator::set_size_compare_ratio, py::arg("src"))
  113. .def("clear", &UnlockedPoolAllocator::clear)
  114. .def("fastMalloc", &UnlockedPoolAllocator::fastMalloc, py::arg("size"))
  115. .def("fastFree", &UnlockedPoolAllocator::fastFree, py::arg("ptr"));
  116. py::class_<DataReader, PyDataReader<> >(m, "DataReader")
  117. .def(py::init<>())
  118. #if NCNN_STRING
  119. .def("scan", &DataReader::scan, py::arg("format"), py::arg("p"))
  120. #endif // NCNN_STRING
  121. .def("read", &DataReader::read, py::arg("buf"), py::arg("size"));
  122. py::class_<DataReaderFromEmpty, DataReader, PyDataReaderOther<DataReaderFromEmpty> >(m, "DataReaderFromEmpty")
  123. .def(py::init<>())
  124. #if NCNN_STRING
  125. .def("scan", &DataReaderFromEmpty::scan, py::arg("format"), py::arg("p"))
  126. #endif // NCNN_STRING
  127. .def("read", &DataReaderFromEmpty::read, py::arg("buf"), py::arg("size"));
  128. py::class_<Blob>(m, "Blob")
  129. .def(py::init<>())
  130. #if NCNN_STRING
  131. .def_readwrite("name", &Blob::name)
  132. #endif // NCNN_STRING
  133. .def_readwrite("producer", &Blob::producer)
  134. .def_readwrite("consumer", &Blob::consumer)
  135. .def_readwrite("shape", &Blob::shape);
  136. py::class_<ModelBin, PyModelBin<> >(m, "ModelBin")
  137. .def(py::init<>())
  138. .def("load", (Mat(ModelBin::*)(int, int) const) & ModelBin::load, py::arg("w"), py::arg("type"))
  139. .def("load", (Mat(ModelBin::*)(int, int, int) const) & ModelBin::load, py::arg("w"), py::arg("h"), py::arg("type"))
  140. .def("load", (Mat(ModelBin::*)(int, int, int, int) const) & ModelBin::load, py::arg("w"), py::arg("h"), py::arg("c"), py::arg("type"))
  141. .def("load", (Mat(ModelBin::*)(int, int, int, int, int) const) & ModelBin::load, py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::arg("type"));
  142. py::class_<ModelBinFromDataReader, ModelBin, PyModelBinOther<ModelBinFromDataReader> >(m, "ModelBinFromDataReader")
  143. .def(py::init<const DataReader&>(), py::arg("dr"))
  144. .def("load", &ModelBinFromDataReader::load, py::arg("w"), py::arg("type"));
  145. py::class_<ModelBinFromMatArray, ModelBin, PyModelBinOther<ModelBinFromMatArray> >(m, "ModelBinFromMatArray")
  146. .def(py::init<const Mat*>(), py::arg("weights"))
  147. .def("load", &ModelBinFromMatArray::load, py::arg("w"), py::arg("type"));
  148. py::class_<ParamDict>(m, "ParamDict")
  149. .def(py::init<>())
  150. .def("type", &ParamDict::type, py::arg("id"))
  151. .def("get", (int (ParamDict::*)(int, int) const) & ParamDict::get, py::arg("id"), py::arg("def"))
  152. .def("get", (float (ParamDict::*)(int, float) const) & ParamDict::get, py::arg("id"), py::arg("def"))
  153. .def("get", (Mat(ParamDict::*)(int, const Mat&) const) & ParamDict::get, py::arg("id"), py::arg("def"))
  154. .def("set", (void (ParamDict::*)(int, int)) & ParamDict::set, py::arg("id"), py::arg("i"))
  155. .def("set", (void (ParamDict::*)(int, float)) & ParamDict::set, py::arg("id"), py::arg("f"))
  156. .def("set", (void (ParamDict::*)(int, const Mat&)) & ParamDict::set, py::arg("id"), py::arg("v"));
  157. py::class_<Option>(m, "Option")
  158. .def(py::init<>())
  159. .def_readwrite("lightmode", &Option::lightmode)
  160. .def_readwrite("num_threads", &Option::num_threads)
  161. .def_readwrite("blob_allocator", &Option::blob_allocator)
  162. .def_readwrite("workspace_allocator", &Option::workspace_allocator)
  163. #if NCNN_VULKAN
  164. .def_readwrite("blob_vkallocator", &Option::blob_vkallocator)
  165. .def_readwrite("workspace_vkallocator", &Option::workspace_vkallocator)
  166. .def_readwrite("staging_vkallocator", &Option::staging_vkallocator)
  167. //.def_readwrite("pipeline_cache", &Option::pipeline_cache)
  168. #endif // NCNN_VULKAN
  169. .def_readwrite("openmp_blocktime", &Option::openmp_blocktime)
  170. .def_readwrite("use_winograd_convolution", &Option::use_winograd_convolution)
  171. .def_readwrite("use_winograd23_convolution", &Option::use_winograd23_convolution)
  172. .def_readwrite("use_winograd43_convolution", &Option::use_winograd43_convolution)
  173. .def_readwrite("use_winograd63_convolution", &Option::use_winograd63_convolution)
  174. .def_readwrite("use_sgemm_convolution", &Option::use_sgemm_convolution)
  175. .def_readwrite("use_int8_inference", &Option::use_int8_inference)
  176. .def_readwrite("use_vulkan_compute", &Option::use_vulkan_compute)
  177. .def_readwrite("use_bf16_storage", &Option::use_bf16_storage)
  178. .def_readwrite("use_fp16_packed", &Option::use_fp16_packed)
  179. .def_readwrite("use_fp16_storage", &Option::use_fp16_storage)
  180. .def_readwrite("use_fp16_arithmetic", &Option::use_fp16_arithmetic)
  181. .def_readwrite("use_int8_packed", &Option::use_int8_packed)
  182. .def_readwrite("use_int8_storage", &Option::use_int8_storage)
  183. .def_readwrite("use_int8_arithmetic", &Option::use_int8_arithmetic)
  184. .def_readwrite("use_packing_layout", &Option::use_packing_layout)
  185. .def_readwrite("use_shader_pack8", &Option::use_shader_pack8)
  186. .def_readwrite("use_subgroup_ops", &Option::use_subgroup_ops)
  187. .def_readwrite("use_tensor_storage", &Option::use_tensor_storage);
  188. py::class_<Mat> mat(m, "Mat", py::buffer_protocol());
  189. mat.def(py::init<>())
  190. .def(py::init(
  191. [](py::tuple shape, size_t elemsize, int elempack, Allocator* allocator) {
  192. Mat* mat = nullptr;
  193. switch (shape.size())
  194. {
  195. case 1:
  196. mat = new Mat(shape[0].cast<int>(), elemsize, elempack, allocator);
  197. break;
  198. case 2:
  199. mat = new Mat(shape[0].cast<int>(), shape[1].cast<int>(), elemsize, elempack, allocator);
  200. break;
  201. case 3:
  202. mat = new Mat(shape[0].cast<int>(), shape[1].cast<int>(), shape[2].cast<int>(), elemsize, elempack, allocator);
  203. break;
  204. case 4:
  205. mat = new Mat(shape[0].cast<int>(), shape[1].cast<int>(), shape[2].cast<int>(), shape[3].cast<int>(), elemsize, elempack, allocator);
  206. break;
  207. default:
  208. std::stringstream ss;
  209. ss << "shape must be 1, 2, 3 or 4 dims, not " << shape.size();
  210. pybind11::pybind11_fail(ss.str());
  211. }
  212. return mat;
  213. }),
  214. py::arg("shape"), py::kw_only(),
  215. py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  216. .def(py::init<int, size_t, int, Allocator*>(),
  217. py::arg("w"), py::kw_only(),
  218. py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  219. .def(py::init<int, int, size_t, int, Allocator*>(),
  220. py::arg("w"), py::arg("h"), py::kw_only(),
  221. py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  222. .def(py::init<int, int, int, size_t, int, Allocator*>(),
  223. py::arg("w"), py::arg("h"), py::arg("c"), py::kw_only(),
  224. py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  225. .def(py::init<int, int, int, int, size_t, int, Allocator*>(),
  226. py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::kw_only(),
  227. py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  228. .def(py::init<const Mat&>(), py::arg("m"))
  229. .def(py::init([](py::buffer const b) {
  230. py::buffer_info info = b.request();
  231. if (info.ndim > 4)
  232. {
  233. std::stringstream ss;
  234. ss << "convert numpy.ndarray to ncnn.Mat only dims <=4 support now, but given " << info.ndim;
  235. pybind11::pybind11_fail(ss.str());
  236. }
  237. size_t elemsize = info.itemsize;
  238. Mat* v = nullptr;
  239. if (info.ndim == 1)
  240. {
  241. v = new Mat((int)info.shape[0], info.ptr, elemsize);
  242. }
  243. else if (info.ndim == 2)
  244. {
  245. v = new Mat((int)info.shape[1], (int)info.shape[0], info.ptr, elemsize);
  246. }
  247. else if (info.ndim == 3)
  248. {
  249. v = new Mat((int)info.shape[2], (int)info.shape[1], (int)info.shape[0], info.ptr, elemsize);
  250. // in ncnn, buffer to construct ncnn::Mat need align to ncnn::alignSize
  251. // with (w * h * elemsize, 16) / elemsize, but the buffer from numpy not
  252. // so we set the cstep as numpy's cstep
  253. v->cstep = (int)info.shape[2] * (int)info.shape[1];
  254. }
  255. else if (info.ndim == 4)
  256. {
  257. v = new Mat((int)info.shape[3], (int)info.shape[2], (int)info.shape[1], (int)info.shape[0], info.ptr, elemsize);
  258. // in ncnn, buffer to construct ncnn::Mat need align to ncnn::alignSize
  259. // with (w * h * d elemsize, 16) / elemsize, but the buffer from numpy not
  260. // so we set the cstep as numpy's cstep
  261. v->cstep = (int)info.shape[3] * (int)info.shape[2] * (int)info.shape[1];
  262. }
  263. return std::unique_ptr<Mat>(v);
  264. }),
  265. py::arg("array"))
  266. .def_buffer([](Mat& m) -> py::buffer_info {
  267. return to_buffer_info(m);
  268. })
  269. .def(
  270. "numpy", [](py::object obj, const std::string& format = "") -> py::array {
  271. auto* m = obj.cast<Mat*>();
  272. return py::array(to_buffer_info(*m, format), obj);
  273. },
  274. py::arg("format") = "", "i for int32, f for float32, d for double")
  275. //.def("fill", (void (Mat::*)(int))(&Mat::fill), py::arg("v"))
  276. .def("fill", (void (Mat::*)(float))(&Mat::fill), py::arg("v"))
  277. .def("clone", &Mat::clone, py::arg("allocator") = nullptr)
  278. .def("clone_from", &Mat::clone_from, py::arg("mat"), py::arg("allocator") = nullptr)
  279. .def(
  280. "reshape", [](Mat& mat, py::tuple shape, Allocator* allocator) {
  281. switch (shape.size())
  282. {
  283. case 1:
  284. return mat.reshape(shape[0].cast<int>(), allocator);
  285. case 2:
  286. return mat.reshape(shape[0].cast<int>(), shape[1].cast<int>(), allocator);
  287. case 3:
  288. return mat.reshape(shape[0].cast<int>(), shape[1].cast<int>(), shape[2].cast<int>(), allocator);
  289. case 4:
  290. return mat.reshape(shape[0].cast<int>(), shape[1].cast<int>(), shape[2].cast<int>(), shape[3].cast<int>(), allocator);
  291. default:
  292. std::stringstream ss;
  293. ss << "shape must be 1, 2, 3 or 4 dims, not " << shape.size();
  294. pybind11::pybind11_fail(ss.str());
  295. }
  296. return Mat();
  297. },
  298. py::arg("shape") = py::tuple(1), py::arg("allocator") = nullptr)
  299. .def("reshape", (Mat(Mat::*)(int, Allocator*) const) & Mat::reshape, py::arg("w"), py::kw_only(), py::arg("allocator") = nullptr)
  300. .def("reshape", (Mat(Mat::*)(int, int, Allocator*) const) & Mat::reshape, py::arg("w"), py::arg("h"), py::kw_only(), py::arg("allocator") = nullptr)
  301. .def("reshape", (Mat(Mat::*)(int, int, int, Allocator*) const) & Mat::reshape, py::arg("w"), py::arg("h"), py::arg("c"), py::kw_only(), py::arg("allocator") = nullptr)
  302. .def("reshape", (Mat(Mat::*)(int, int, int, int, Allocator*) const) & Mat::reshape, py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::kw_only(), py::arg("allocator") = nullptr)
  303. .def(
  304. "create", [](Mat& mat, py::tuple shape, size_t elemsize, int elempack, Allocator* allocator) {
  305. switch (shape.size())
  306. {
  307. case 1:
  308. return mat.create(shape[0].cast<int>(), elemsize, elempack, allocator);
  309. case 2:
  310. return mat.create(shape[0].cast<int>(), shape[1].cast<int>(), elemsize, elempack, allocator);
  311. case 3:
  312. return mat.create(shape[0].cast<int>(), shape[1].cast<int>(), shape[2].cast<int>(), elemsize, elempack, allocator);
  313. case 4:
  314. return mat.create(shape[0].cast<int>(), shape[1].cast<int>(), shape[2].cast<int>(), shape[3].cast<int>(), elemsize, elempack, allocator);
  315. default:
  316. std::stringstream ss;
  317. ss << "shape must be 1, 2, 3 or 4 dims, not " << shape.size();
  318. pybind11::pybind11_fail(ss.str());
  319. }
  320. return;
  321. },
  322. py::arg("shape"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  323. .def("create", (void (Mat::*)(int, size_t, int, Allocator*)) & Mat::create, py::arg("w"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  324. .def("create", (void (Mat::*)(int, int, size_t, int, Allocator*)) & Mat::create, py::arg("w"), py::arg("h"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  325. .def("create", (void (Mat::*)(int, int, int, size_t, int, Allocator*)) & Mat::create, py::arg("w"), py::arg("h"), py::arg("c"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  326. .def("create", (void (Mat::*)(int, int, int, int, size_t, int, Allocator*)) & Mat::create, py::arg("w"), py::arg("h"), py::arg("d"), py::arg("c"), py::kw_only(), py::arg("elemsize") = 4, py::arg("elempack") = 1, py::arg("allocator") = nullptr)
  327. .def("create_like", (void (Mat::*)(const Mat&, Allocator*)) & Mat::create_like, py::arg("m"), py::arg("allocator") = nullptr)
  328. .def("addref", &Mat::addref)
  329. .def("release", &Mat::release)
  330. .def("empty", &Mat::empty)
  331. .def("total", &Mat::total)
  332. .def("elembits", &Mat::elembits)
  333. .def("shape", &Mat::shape)
  334. .def("channel", (Mat(Mat::*)(int)) & Mat::channel, py::arg("c"))
  335. //.def("channel", (const Mat (Mat::*)(int) const) & Mat::channel, py::arg("c"))
  336. .def("depth", (Mat(Mat::*)(int)) & Mat::depth, py::arg("z"))
  337. //.def("depth", (const Mat (Mat::*)(int) const) & Mat::depth, py::arg("z"))
  338. .def(
  339. "row", [](Mat& m, int y) {
  340. if (m.elempack != 1)
  341. {
  342. std::stringstream ss;
  343. ss << "get ncnn.Mat row only elempack 1 support now, but given " << m.elempack;
  344. pybind11::pybind11_fail(ss.str());
  345. }
  346. switch (m.elemsize)
  347. {
  348. case 1:
  349. return py::memoryview::from_buffer(m.row<int8_t>(y), {m.w}, {sizeof(int8_t)});
  350. //case 2:
  351. // return py::memoryview::from_buffer(m.row<short>(y), {m.w}, {sizeof(short)});
  352. case 4:
  353. return py::memoryview::from_buffer(m.row<float>(y), {m.w}, {sizeof(float)});
  354. default:
  355. std::stringstream ss;
  356. ss << "ncnn.Mat row elemsize " << m.elemsize << "not support now";
  357. pybind11::pybind11_fail(ss.str());
  358. }
  359. return py::memoryview::from_buffer(m.row<float>(y), {m.w}, {sizeof(float)});
  360. },
  361. py::arg("y"))
  362. .def("channel_range", (Mat(Mat::*)(int, int)) & Mat::channel_range, py::arg("c"), py::arg("channels"))
  363. //.def("channel_range", (const Mat (Mat::*)(int, int) const) & Mat::channel_range, py::arg("c"), py::arg("channels"))
  364. .def("depth_range", (Mat(Mat::*)(int, int)) & Mat::depth_range, py::arg("z"), py::arg("depths"))
  365. //.def("depth_range", (const Mat (Mat::*)(int, int) const) & Mat::depth_range, py::arg("z"), py::arg("depths"))
  366. .def("row_range", (Mat(Mat::*)(int, int)) & Mat::row_range, py::arg("y"), py::arg("rows"))
  367. //.def("row_range", (const Mat (Mat::*)(int, int) const) & Mat::row_range, py::arg("y"), py::arg("rows"))
  368. .def("range", (Mat(Mat::*)(int, int)) & Mat::range, py::arg("x"), py::arg("n"))
  369. //.def("range", (const Mat (Mat::*)(int, int) const) & Mat::range, py::arg("x"), py::arg("n"))
  370. .def(
  371. "__getitem__", [](const Mat& m, size_t i) {
  372. return m[i];
  373. },
  374. py::arg("i"))
  375. .def(
  376. "__setitem__", [](Mat& m, size_t i, float v) {
  377. m[i] = v;
  378. },
  379. py::arg("i"), py::arg("v"))
  380. .def("__len__", [](Mat& m) {
  381. return m.w;
  382. })
  383. //convenient construct from pixel data
  384. .def_static(
  385. "from_pixels", [](py::buffer const b, int type, int w, int h, Allocator* allocator) {
  386. return Mat::from_pixels((const unsigned char*)b.request().ptr, type, w, h, allocator);
  387. },
  388. py::arg("array"), py::arg("type"), py::arg("w"), py::arg("h"), py::arg("allocator") = nullptr)
  389. .def_static(
  390. "from_pixels", [](py::buffer const b, int type, int w, int h, int stride, Allocator* allocator) {
  391. return Mat::from_pixels((const unsigned char*)b.request().ptr, type, w, h, stride, allocator);
  392. },
  393. py::arg("array"), py::arg("type"), py::arg("w"), py::arg("h"), py::arg("stride"), py::arg("allocator") = nullptr)
  394. .def_static(
  395. "from_pixels_resize", [](py::buffer const b, int type, int w, int h, int target_width, int target_height, Allocator* allocator) {
  396. return Mat::from_pixels_resize((const unsigned char*)b.request().ptr,
  397. type, w, h, target_width, target_height, allocator);
  398. },
  399. py::arg("array"), py::arg("type"), py::arg("w"), py::arg("h"), py::arg("target_width"), py::arg("target_height"), py::arg("allocator") = nullptr)
  400. .def_static(
  401. "from_pixels_resize", [](py::buffer const b, int type, int w, int h, int stride, int target_width, int target_height, Allocator* allocator) {
  402. return Mat::from_pixels_resize((const unsigned char*)b.request().ptr,
  403. type, w, h, stride, target_width, target_height, allocator);
  404. },
  405. py::arg("array"), py::arg("type"), py::arg("w"), py::arg("h"), py::arg("stride"), py::arg("target_width"), py::arg("target_height"), py::arg("allocator") = nullptr)
  406. .def_static(
  407. "from_pixels_roi", [](py::buffer const b, int type, int w, int h, int roix, int roiy, int roiw, int roih, Allocator* allocator) {
  408. return Mat::from_pixels_roi((const unsigned char*)b.request().ptr,
  409. type, w, h, roix, roiy, roiw, roih, allocator);
  410. },
  411. py::arg("array"), py::arg("type"), py::arg("w"), py::arg("h"), py::arg("roix"), py::arg("roiy"), py::arg("roiw"), py::arg("roih"), py::arg("allocator") = nullptr)
  412. .def_static(
  413. "from_pixels_roi", [](py::buffer const b, int type, int w, int h, int stride, int roix, int roiy, int roiw, int roih, Allocator* allocator) {
  414. return Mat::from_pixels_roi((const unsigned char*)b.request().ptr,
  415. type, w, h, stride, roix, roiy, roiw, roih, allocator);
  416. },
  417. py::arg("array"), py::arg("type"), py::arg("w"), py::arg("h"), py::arg("stride"), py::arg("roix"), py::arg("roiy"), py::arg("roiw"), py::arg("roih"), py::arg("allocator") = nullptr)
  418. .def_static(
  419. "from_pixels_roi_resize", [](py::buffer const b, int type, int w, int h, int roix, int roiy, int roiw, int roih, int target_width, int target_height, Allocator* allocator) {
  420. return Mat::from_pixels_roi_resize((const unsigned char*)b.request().ptr,
  421. type, w, h, roix, roiy, roiw, roih, target_width, target_height, allocator);
  422. },
  423. py::arg("array"), py::arg("type"), py::arg("w"), py::arg("h"), py::arg("roix"), py::arg("roiy"), py::arg("roiw"), py::arg("roih"), py::arg("target_width"), py::arg("target_height"), py::arg("allocator") = nullptr)
  424. .def_static(
  425. "from_pixels_roi_resize", [](py::buffer const b, int type, int w, int h, int stride, int roix, int roiy, int roiw, int roih, int target_width, int target_height, Allocator* allocator) {
  426. return Mat::from_pixels_roi_resize((const unsigned char*)b.request().ptr,
  427. type, w, h, stride, roix, roiy, roiw, roih, target_width, target_height, allocator);
  428. },
  429. py::arg("array"), py::arg("type"), py::arg("w"), py::arg("h"), py::arg("stride"), py::arg("roix"), py::arg("roiy"), py::arg("roiw"), py::arg("roih"), py::arg("target_width"), py::arg("target_height"), py::arg("allocator") = nullptr)
  430. .def(
  431. "substract_mean_normalize", [](Mat& mat, std::vector<float>& mean, std::vector<float>& norm) {
  432. return mat.substract_mean_normalize(mean.size() > 0 ? &mean[0] : 0, norm.size() > 0 ? &norm[0] : 0);
  433. },
  434. py::arg("mean"), py::arg("norm"))
  435. .def_readwrite("refcount", &Mat::refcount)
  436. .def_readwrite("elemsize", &Mat::elemsize)
  437. .def_readwrite("elempack", &Mat::elempack)
  438. .def_readwrite("allocator", &Mat::allocator)
  439. .def_readwrite("dims", &Mat::dims)
  440. .def_readwrite("w", &Mat::w)
  441. .def_readwrite("h", &Mat::h)
  442. .def_readwrite("d", &Mat::d)
  443. .def_readwrite("c", &Mat::c)
  444. .def_readwrite("cstep", &Mat::cstep)
  445. .def("__repr__", [](const Mat& m) {
  446. std::stringstream ss;
  447. ss << "<ncnn.Mat w=" << m.w << " h=" << m.h << " d=" << m.d << " c=" << m.c << " dims=" << m.dims
  448. << " cstep=" << m.cstep << " elemsize=" << m.elemsize << " elempack=" << m.elempack << "\n\t"
  449. << "refcount=" << (m.refcount ? *m.refcount : 0) << " data=0x" << static_cast<const void*>(m.data)
  450. << " allocator=0x" << static_cast<const void*>(m.allocator) << ">\n";
  451. const int max_count = m.dims == 1 ? 10 : 6;
  452. if (m.dims == 1)
  453. {
  454. ss << "[";
  455. bool dot_printed_w = false;
  456. if (m.elemsize == 1)
  457. {
  458. const int8_t* row = m.row<int8_t>(0);
  459. for (int i = 0; i < m.w; i++)
  460. {
  461. if (i < max_count / 2 || i >= m.w - max_count / 2)
  462. {
  463. if (i > 0)
  464. {
  465. ss << ", ";
  466. }
  467. ss << static_cast<int>(row[i]);
  468. }
  469. else if (!dot_printed_w)
  470. {
  471. dot_printed_w = true;
  472. ss << ", ...";
  473. }
  474. }
  475. }
  476. if (m.elemsize == 4)
  477. {
  478. const float* row = m.row<float>(0);
  479. for (int i = 0; i < m.w; i++)
  480. {
  481. if (i < max_count / 2 || i >= m.w - max_count / 2)
  482. {
  483. if (i > 0)
  484. {
  485. ss << ", ";
  486. }
  487. ss << row[i];
  488. }
  489. else if (!dot_printed_w)
  490. {
  491. dot_printed_w = true;
  492. ss << ", ...";
  493. }
  494. }
  495. }
  496. ss << "]";
  497. }
  498. else if (m.dims == 2)
  499. {
  500. bool dot_printed_h = false;
  501. ss << "[";
  502. for (int j = 0; j < m.h; j++)
  503. {
  504. bool dot_printed_w = false;
  505. if (j < max_count / 2 || j >= m.h - max_count / 2)
  506. {
  507. ss << "[";
  508. if (m.elemsize == 1)
  509. {
  510. const int8_t* row = m.row<int8_t>(j);
  511. for (int i = 0; i < m.w; i++)
  512. {
  513. if (i < max_count / 2 || i >= m.w - max_count / 2)
  514. {
  515. if (i > 0)
  516. {
  517. ss << ", ";
  518. }
  519. ss << static_cast<int>(row[i]);
  520. }
  521. else if (!dot_printed_w)
  522. {
  523. dot_printed_w = true;
  524. ss << ", ...";
  525. }
  526. }
  527. }
  528. if (m.elemsize == 4)
  529. {
  530. const float* row = m.row<float>(j);
  531. for (int i = 0; i < m.w; i++)
  532. {
  533. if (i < max_count / 2 || i >= m.w - max_count / 2)
  534. {
  535. if (i > 0)
  536. {
  537. ss << ", ";
  538. }
  539. ss << row[i];
  540. }
  541. else if (!dot_printed_w)
  542. {
  543. dot_printed_w = true;
  544. ss << ", ...";
  545. }
  546. }
  547. }
  548. ss << "]";
  549. if (j < m.h - 1)
  550. {
  551. ss << "\n";
  552. }
  553. }
  554. else if (!dot_printed_h)
  555. {
  556. dot_printed_h = true;
  557. ss << "...\n";
  558. }
  559. }
  560. ss << "]\n";
  561. }
  562. else if (m.dims == 3)
  563. {
  564. bool dot_printed_c = false;
  565. ss << "[";
  566. for (int k = 0; k < m.c; k++)
  567. {
  568. bool dot_printed_h = false;
  569. if (k < max_count / 2 || k >= m.c - max_count / 2)
  570. {
  571. Mat channel = m.channel(k);
  572. if (k > 0)
  573. {
  574. ss << " ";
  575. }
  576. ss << "[";
  577. for (int j = 0; j < channel.h; j++)
  578. {
  579. bool dot_printed_w = false;
  580. if (j < max_count / 2 || j >= channel.h - max_count / 2)
  581. {
  582. if (j > 0)
  583. {
  584. ss << " ";
  585. }
  586. ss << "[";
  587. if (m.elemsize == 1)
  588. {
  589. const int8_t* row = channel.row<int8_t>(j);
  590. for (int i = 0; i < channel.w; i++)
  591. {
  592. if (i < max_count / 2 || i >= channel.w - max_count / 2)
  593. {
  594. if (i > 0)
  595. {
  596. ss << ", ";
  597. }
  598. ss << static_cast<int>(row[i]);
  599. }
  600. else if (!dot_printed_w)
  601. {
  602. dot_printed_w = true;
  603. ss << ", ...";
  604. }
  605. }
  606. }
  607. if (m.elemsize == 4)
  608. {
  609. const float* row = channel.row<float>(j);
  610. for (int i = 0; i < m.w; i++)
  611. {
  612. if (i < max_count / 2 || i >= m.w - max_count / 2)
  613. {
  614. if (i > 0)
  615. {
  616. ss << ", ";
  617. }
  618. ss << row[i];
  619. }
  620. else if (!dot_printed_w)
  621. {
  622. dot_printed_w = true;
  623. ss << ", ...";
  624. }
  625. }
  626. }
  627. ss << "]";
  628. if (j < channel.h - 1)
  629. {
  630. ss << "\n";
  631. }
  632. }
  633. else if (!dot_printed_h)
  634. {
  635. dot_printed_h = true;
  636. ss << " ...\n";
  637. }
  638. } // for j
  639. ss << "]";
  640. if (k < m.c - 1)
  641. {
  642. ss << "\n\n";
  643. }
  644. }
  645. else if (!dot_printed_c)
  646. {
  647. dot_printed_c = true;
  648. ss << " ...\n";
  649. }
  650. } // for k
  651. ss << "]\n";
  652. }
  653. else if (m.dims == 4)
  654. {
  655. bool dot_printed_c = false;
  656. ss << "[";
  657. for (int k = 0; k < m.c; k++)
  658. {
  659. bool dot_printed_d = false;
  660. if (k < max_count / 2 || k >= m.c - max_count / 2)
  661. {
  662. Mat channel = m.channel(k);
  663. if (k > 0)
  664. {
  665. ss << " ";
  666. }
  667. ss << "[";
  668. for (int z = 0; z < channel.d; z++)
  669. {
  670. bool dot_printed_h = false;
  671. if (z < max_count / 2 || z >= channel.d - max_count / 2)
  672. {
  673. if (z > 0)
  674. {
  675. ss << " ";
  676. }
  677. ss << "[";
  678. for (int j = 0; j < channel.h; j++)
  679. {
  680. bool dot_printed_w = false;
  681. if (j < max_count / 2 || j >= channel.h - max_count / 2)
  682. {
  683. if (j > 0)
  684. {
  685. ss << " ";
  686. }
  687. ss << "[";
  688. if (m.elemsize == 1)
  689. {
  690. const int8_t* row = channel.depth(z).row<int8_t>(j);
  691. for (int i = 0; i < channel.w; i++)
  692. {
  693. if (i < max_count / 2 || i >= channel.w - max_count / 2)
  694. {
  695. if (i > 0)
  696. {
  697. ss << ", ";
  698. }
  699. ss << static_cast<int>(row[i]);
  700. }
  701. else if (!dot_printed_w)
  702. {
  703. dot_printed_w = true;
  704. ss << ", ...";
  705. }
  706. }
  707. }
  708. if (m.elemsize == 4)
  709. {
  710. const float* row = channel.depth(z).row<float>(j);
  711. for (int i = 0; i < m.w; i++)
  712. {
  713. if (i < max_count / 2 || i >= m.w - max_count / 2)
  714. {
  715. if (i > 0)
  716. {
  717. ss << ", ";
  718. }
  719. ss << row[i];
  720. }
  721. else if (!dot_printed_w)
  722. {
  723. dot_printed_w = true;
  724. ss << ", ...";
  725. }
  726. }
  727. }
  728. ss << "]";
  729. if (j < channel.h - 1)
  730. {
  731. ss << "\n";
  732. }
  733. }
  734. else if (!dot_printed_h)
  735. {
  736. dot_printed_h = true;
  737. ss << " ...\n";
  738. }
  739. } // for j
  740. ss << "]";
  741. if (z < channel.d - 1)
  742. {
  743. ss << "\n";
  744. }
  745. }
  746. else if (!dot_printed_d)
  747. {
  748. dot_printed_d = true;
  749. ss << " ...\n";
  750. }
  751. } // for z
  752. ss << "]";
  753. if (k < m.c - 1)
  754. {
  755. ss << "\n\n";
  756. }
  757. }
  758. else if (!dot_printed_c)
  759. {
  760. dot_printed_c = true;
  761. ss << " ...\n";
  762. }
  763. } // for k
  764. ss << "]\n";
  765. }
  766. return ss.str();
  767. });
  768. py::enum_<ncnn::Mat::PixelType>(mat, "PixelType")
  769. .value("PIXEL_CONVERT_SHIFT", ncnn::Mat::PixelType::PIXEL_CONVERT_SHIFT)
  770. .value("PIXEL_FORMAT_MASK", ncnn::Mat::PixelType::PIXEL_FORMAT_MASK)
  771. .value("PIXEL_CONVERT_MASK", ncnn::Mat::PixelType::PIXEL_CONVERT_MASK)
  772. .value("PIXEL_RGB", ncnn::Mat::PixelType::PIXEL_RGB)
  773. .value("PIXEL_BGR", ncnn::Mat::PixelType::PIXEL_BGR)
  774. .value("PIXEL_GRAY", ncnn::Mat::PixelType::PIXEL_GRAY)
  775. .value("PIXEL_RGBA", ncnn::Mat::PixelType::PIXEL_RGBA)
  776. .value("PIXEL_BGRA", ncnn::Mat::PixelType::PIXEL_BGRA)
  777. .value("PIXEL_RGB2BGR", ncnn::Mat::PixelType::PIXEL_RGB2BGR)
  778. .value("PIXEL_RGB2GRAY", ncnn::Mat::PixelType::PIXEL_RGB2GRAY)
  779. .value("PIXEL_RGB2RGBA", ncnn::Mat::PixelType::PIXEL_RGB2RGBA)
  780. .value("PIXEL_RGB2BGRA", ncnn::Mat::PixelType::PIXEL_RGB2BGRA)
  781. .value("PIXEL_BGR2RGB", ncnn::Mat::PixelType::PIXEL_BGR2RGB)
  782. .value("PIXEL_BGR2GRAY", ncnn::Mat::PixelType::PIXEL_BGR2GRAY)
  783. .value("PIXEL_BGR2RGBA", ncnn::Mat::PixelType::PIXEL_BGR2RGBA)
  784. .value("PIXEL_BGR2BGRA", ncnn::Mat::PixelType::PIXEL_BGR2BGRA)
  785. .value("PIXEL_GRAY2RGB", ncnn::Mat::PixelType::PIXEL_GRAY2RGB)
  786. .value("PIXEL_GRAY2BGR", ncnn::Mat::PixelType::PIXEL_GRAY2BGR)
  787. .value("PIXEL_GRAY2RGBA", ncnn::Mat::PixelType::PIXEL_GRAY2RGBA)
  788. .value("PIXEL_GRAY2BGRA", ncnn::Mat::PixelType::PIXEL_GRAY2BGRA)
  789. .value("PIXEL_RGBA2RGB", ncnn::Mat::PixelType::PIXEL_RGBA2RGB)
  790. .value("PIXEL_RGBA2BGR", ncnn::Mat::PixelType::PIXEL_RGBA2BGR)
  791. .value("PIXEL_RGBA2GRAY", ncnn::Mat::PixelType::PIXEL_RGBA2GRAY)
  792. .value("PIXEL_RGBA2BGRA", ncnn::Mat::PixelType::PIXEL_RGBA2BGRA)
  793. .value("PIXEL_BGRA2RGB", ncnn::Mat::PixelType::PIXEL_BGRA2RGB)
  794. .value("PIXEL_BGRA2BGR", ncnn::Mat::PixelType::PIXEL_BGRA2BGR)
  795. .value("PIXEL_BGRA2GRAY", ncnn::Mat::PixelType::PIXEL_BGRA2GRAY)
  796. .value("PIXEL_BGRA2RGBA", ncnn::Mat::PixelType::PIXEL_BGRA2RGBA);
  797. py::class_<Extractor>(m, "Extractor")
  798. .def("__enter__", [](Extractor& ex) -> Extractor& { return ex; })
  799. .def("__exit__", [](Extractor& ex, pybind11::args) {
  800. ex.clear();
  801. })
  802. .def("clear", &Extractor::clear)
  803. .def("set_light_mode", &Extractor::set_light_mode, py::arg("enable"))
  804. .def("set_num_threads", &Extractor::set_num_threads, py::arg("num_threads"))
  805. .def("set_blob_allocator", &Extractor::set_blob_allocator, py::arg("allocator"))
  806. .def("set_workspace_allocator", &Extractor::set_workspace_allocator, py::arg("allocator"))
  807. #if NCNN_STRING
  808. .def("input", (int (Extractor::*)(const char*, const Mat&)) & Extractor::input, py::arg("blob_name"), py::arg("in"))
  809. .def("extract", (int (Extractor::*)(const char*, Mat&, int)) & Extractor::extract, py::arg("blob_name"), py::arg("feat"), py::arg("type") = 0)
  810. .def(
  811. "extract", [](Extractor& ex, const char* blob_name, int type) {
  812. ncnn::Mat feat;
  813. int ret = ex.extract(blob_name, feat, type);
  814. return py::make_tuple(ret, feat.clone());
  815. },
  816. py::arg("blob_name"), py::arg("type") = 0)
  817. #endif
  818. .def("input", (int (Extractor::*)(int, const Mat&)) & Extractor::input)
  819. .def("extract", (int (Extractor::*)(int, Mat&, int)) & Extractor::extract, py::arg("blob_index"), py::arg("feat"), py::arg("type") = 0)
  820. .def(
  821. "extract", [](Extractor& ex, int blob_index, int type) {
  822. ncnn::Mat feat;
  823. int ret = ex.extract(blob_index, feat, type);
  824. return py::make_tuple(ret, feat.clone());
  825. },
  826. py::arg("blob_index"), py::arg("type") = 0);
  827. py::class_<Layer, PyLayer>(m, "Layer")
  828. .def(py::init<>())
  829. .def("load_param", &Layer::load_param, py::arg("pd"))
  830. .def("load_model", &Layer::load_model, py::arg("mb"))
  831. .def("create_pipeline", &Layer::create_pipeline, py::arg("opt"))
  832. .def("destroy_pipeline", &Layer::destroy_pipeline, py::arg("opt"))
  833. .def_readwrite("one_blob_only", &Layer::one_blob_only)
  834. .def_readwrite("support_inplace", &Layer::support_inplace)
  835. .def_readwrite("support_vulkan", &Layer::support_vulkan)
  836. .def_readwrite("support_packing", &Layer::support_packing)
  837. .def_readwrite("support_bf16_storage", &Layer::support_bf16_storage)
  838. .def_readwrite("support_fp16_storage", &Layer::support_fp16_storage)
  839. .def("forward", (int (Layer::*)(const std::vector<Mat>&, std::vector<Mat>&, const Option&) const) & Layer::forward,
  840. py::arg("bottom_blobs"), py::arg("top_blobs"), py::arg("opt"))
  841. .def("forward", (int (Layer::*)(const Mat&, Mat&, const Option&) const) & Layer::forward,
  842. py::arg("bottom_blob"), py::arg("top_blob"), py::arg("opt"))
  843. .def("forward_inplace", (int (Layer::*)(std::vector<Mat>&, const Option&) const) & Layer::forward_inplace,
  844. py::arg("bottom_top_blobs"), py::arg("opt"))
  845. .def("forward_inplace", (int (Layer::*)(Mat&, const Option&) const) & Layer::forward_inplace,
  846. py::arg("bottom_top_blob"), py::arg("opt"))
  847. .def_readwrite("typeindex", &Layer::typeindex)
  848. #if NCNN_STRING
  849. .def_readwrite("type", &Layer::type)
  850. .def_readwrite("name", &Layer::name)
  851. #endif // NCNN_STRING
  852. .def_readwrite("bottoms", &Layer::bottoms)
  853. .def_readwrite("tops", &Layer::tops)
  854. .def_readwrite("bottom_shapes", &Layer::bottom_shapes)
  855. .def_readwrite("top_shapes", &Layer::top_shapes);
  856. py::class_<Net>(m, "Net")
  857. .def(py::init<>())
  858. .def_readwrite("opt", &Net::opt)
  859. .def("__enter__", [](Net& net) -> Net& { return net; })
  860. .def("__exit__", [](Net& net, pybind11::args) {
  861. net.clear();
  862. })
  863. #if NCNN_VULKAN
  864. .def("set_vulkan_device", (void (Net::*)(int)) & Net::set_vulkan_device, py::arg("device_index"))
  865. .def("set_vulkan_device", (void (Net::*)(const VulkanDevice*)) & Net::set_vulkan_device, py::arg("vkdev"))
  866. .def("vulkan_device", &Net::vulkan_device, py::return_value_policy::reference_internal)
  867. #endif // NCNN_VULKAN
  868. #if NCNN_STRING
  869. .def(
  870. "register_custom_layer", [](Net& net, const char* type, const std::function<ncnn::Layer*()>& creator, const std::function<void(ncnn::Layer*)>& destroyer) {
  871. if (g_layer_factroy_index == g_layer_factroys.size())
  872. {
  873. std::stringstream ss;
  874. ss << "python version only support " << g_layer_factroys.size() << " custom layers now";
  875. pybind11::pybind11_fail(ss.str());
  876. }
  877. LayerFactory& lf = g_layer_factroys[g_layer_factroy_index++];
  878. lf.name = type;
  879. lf.creator = creator;
  880. lf.destroyer = destroyer;
  881. return net.register_custom_layer(lf.name.c_str(), lf.creator_func, lf.destroyer_func);
  882. },
  883. py::arg("type"), py::arg("creator"), py::arg("destroyer"))
  884. #endif //NCNN_STRING
  885. .def(
  886. "register_custom_layer", [](Net& net, int index, const std::function<ncnn::Layer*()>& creator, const std::function<void(ncnn::Layer*)>& destroyer) {
  887. if (g_layer_factroy_index == g_layer_factroys.size())
  888. {
  889. std::stringstream ss;
  890. ss << "python version only support " << g_layer_factroys.size() << " custom layers now";
  891. pybind11::pybind11_fail(ss.str());
  892. }
  893. LayerFactory& lf = g_layer_factroys[g_layer_factroy_index++];
  894. lf.index = index;
  895. lf.creator = creator;
  896. lf.destroyer = destroyer;
  897. return net.register_custom_layer(index, lf.creator_func, lf.destroyer_func);
  898. },
  899. py::arg("index"), py::arg("creator"), py::arg("destroyer"))
  900. #if NCNN_STRING
  901. .def("load_param", (int (Net::*)(const DataReader&)) & Net::load_param, py::arg("dr"))
  902. #endif // NCNN_STRING
  903. .def("load_param_bin", (int (Net::*)(const DataReader&)) & Net::load_param_bin, py::arg("dr"))
  904. .def("load_model", (int (Net::*)(const DataReader&)) & Net::load_model, py::arg("dr"))
  905. #if NCNN_STDIO
  906. #if NCNN_STRING
  907. .def("load_param", (int (Net::*)(const char*)) & Net::load_param, py::arg("protopath"))
  908. .def("load_param_mem", (int (Net::*)(const char*)) & Net::load_param_mem, py::arg("mem"))
  909. #endif // NCNN_STRING
  910. .def("load_param_bin", (int (Net::*)(const char*)) & Net::load_param_bin, py::arg("protopath"))
  911. .def("load_model", (int (Net::*)(const char*)) & Net::load_model, py::arg("modelpath"))
  912. .def(
  913. "load_model_mem", [](Net& net, const char* mem) {
  914. const unsigned char* _mem = (const unsigned char*)mem;
  915. DataReaderFromMemoryCopy dr(_mem);
  916. net.load_model(dr);
  917. },
  918. py::arg("mem"))
  919. #endif // NCNN_STDIO
  920. .def("clear", &Net::clear)
  921. .def("create_extractor", &Net::create_extractor, py::keep_alive<0, 1>()) //net should be kept alive until retuned ex is freed by gc
  922. .def("input_indexes", &Net::input_indexes, py::return_value_policy::reference)
  923. .def("output_indexes", &Net::output_indexes, py::return_value_policy::reference)
  924. #if NCNN_STRING
  925. .def("input_names", &Net::input_names, py::return_value_policy::reference)
  926. .def("output_names", &Net::output_names, py::return_value_policy::reference)
  927. #endif // NCNN_STRING
  928. .def("blobs", &Net::blobs, py::return_value_policy::reference_internal)
  929. .def("layers", &Net::layers, py::return_value_policy::reference_internal);
  930. py::enum_<ncnn::BorderType>(m, "BorderType")
  931. .value("BORDER_CONSTANT", ncnn::BorderType::BORDER_CONSTANT)
  932. .value("BORDER_REPLICATE", ncnn::BorderType::BORDER_REPLICATE);
  933. m.def("cpu_support_arm_neon", &cpu_support_arm_neon);
  934. m.def("cpu_support_arm_vfpv4", &cpu_support_arm_vfpv4);
  935. m.def("cpu_support_arm_asimdhp", &cpu_support_arm_asimdhp);
  936. m.def("cpu_support_x86_avx2", &cpu_support_x86_avx2);
  937. m.def("cpu_support_x86_avx", &cpu_support_x86_avx);
  938. m.def("get_cpu_count", &get_cpu_count);
  939. m.def("get_little_cpu_count", &get_little_cpu_count);
  940. m.def("get_big_cpu_count", &get_big_cpu_count);
  941. m.def("get_physical_cpu_count", &get_physical_cpu_count);
  942. m.def("get_physical_little_cpu_count", &get_physical_little_cpu_count);
  943. m.def("get_physical_big_cpu_count", &get_physical_big_cpu_count);
  944. m.def("get_cpu_powersave", &get_cpu_powersave);
  945. m.def("set_cpu_powersave", &set_cpu_powersave, py::arg("powersave"));
  946. m.def("get_omp_num_threads", &get_omp_num_threads);
  947. m.def("set_omp_num_threads", &set_omp_num_threads, py::arg("num_threads"));
  948. m.def("get_omp_dynamic", &get_omp_dynamic);
  949. m.def("set_omp_dynamic", &set_omp_dynamic, py::arg("dynamic"));
  950. m.def("get_omp_thread_num", &get_omp_thread_num);
  951. m.def("get_kmp_blocktime", &get_kmp_blocktime);
  952. m.def("set_kmp_blocktime", &set_kmp_blocktime, py::arg("time_ms"));
  953. m.def("copy_make_border", &copy_make_border,
  954. py::arg("src"), py::arg("dst"),
  955. py::arg("top"), py::arg("bottom"), py::arg("left"), py::arg("right"),
  956. py::arg("type"), py::arg("v"), py::arg("opt") = Option());
  957. m.def(
  958. "copy_make_border",
  959. [](const Mat& src, int top, int bottom, int left, int right, int type, float v, const Option& opt) {
  960. Mat dst;
  961. copy_make_border(src, dst, top, bottom, left, right, type, v, opt);
  962. return dst;
  963. },
  964. py::arg("src"),
  965. py::arg("top"), py::arg("bottom"), py::arg("left"), py::arg("right"),
  966. py::arg("type"), py::arg("v"), py::arg("opt") = Option());
  967. m.def("copy_make_border_3d", &copy_make_border_3d,
  968. py::arg("src"), py::arg("dst"),
  969. py::arg("top"), py::arg("bottom"), py::arg("left"), py::arg("right"), py::arg("front"), py::arg("behind"),
  970. py::arg("type"), py::arg("v"), py::arg("opt") = Option());
  971. m.def(
  972. "copy_make_border_3d",
  973. [](const Mat& src, int top, int bottom, int left, int right, int front, int behind, int type, float v, const Option& opt) {
  974. Mat dst;
  975. copy_make_border_3d(src, dst, top, bottom, left, right, front, behind, type, v, opt);
  976. return dst;
  977. },
  978. py::arg("src"),
  979. py::arg("top"), py::arg("bottom"), py::arg("left"), py::arg("right"), py::arg("front"), py::arg("behind"),
  980. py::arg("type"), py::arg("v"), py::arg("opt") = Option());
  981. m.def("copy_cut_border", &copy_cut_border,
  982. py::arg("src"), py::arg("dst"),
  983. py::arg("top"), py::arg("bottom"), py::arg("left"), py::arg("right"),
  984. py::arg("opt") = Option());
  985. m.def(
  986. "copy_cut_border",
  987. [](const Mat& src, int top, int bottom, int left, int right, const Option& opt) {
  988. Mat dst;
  989. copy_cut_border(src, dst, top, bottom, left, right, opt);
  990. return dst;
  991. },
  992. py::arg("src"),
  993. py::arg("top"), py::arg("bottom"), py::arg("left"), py::arg("right"),
  994. py::arg("opt") = Option());
  995. m.def("resize_nearest", &resize_nearest,
  996. py::arg("src"), py::arg("dst"),
  997. py::arg("w"), py::arg("h"),
  998. py::arg("opt") = Option());
  999. m.def(
  1000. "resize_nearest",
  1001. [](const Mat& src, int w, int h, const Option& opt) {
  1002. Mat dst;
  1003. resize_nearest(src, dst, w, h);
  1004. return dst;
  1005. },
  1006. py::arg("src"),
  1007. py::arg("w"), py::arg("h"),
  1008. py::arg("opt") = Option());
  1009. m.def("resize_bilinear", &resize_bilinear,
  1010. py::arg("src"), py::arg("dst"),
  1011. py::arg("w"), py::arg("h"),
  1012. py::arg("opt") = Option());
  1013. m.def(
  1014. "resize_bilinear",
  1015. [](const Mat& src, int w, int h, const Option& opt) {
  1016. Mat dst;
  1017. resize_bilinear(src, dst, w, h, opt);
  1018. return dst;
  1019. },
  1020. py::arg("src"),
  1021. py::arg("w"), py::arg("h"),
  1022. py::arg("opt") = Option());
  1023. m.def("resize_bicubic", &resize_bicubic,
  1024. py::arg("src"), py::arg("dst"),
  1025. py::arg("w"), py::arg("h"),
  1026. py::arg("opt") = Option());
  1027. m.def(
  1028. "resize_bicubic",
  1029. [](const Mat& src, int w, int h, const Option& opt) {
  1030. Mat dst;
  1031. resize_bicubic(src, dst, w, h, opt);
  1032. return dst;
  1033. },
  1034. py::arg("src"),
  1035. py::arg("w"), py::arg("h"),
  1036. py::arg("opt") = Option());
  1037. m.def("convert_packing", &convert_packing,
  1038. py::arg("src"), py::arg("dst"),
  1039. py::arg("elempack"),
  1040. py::arg("opt") = Option());
  1041. m.def(
  1042. "convert_packing",
  1043. [](const Mat& src, int elempack, const Option& opt) {
  1044. Mat dst;
  1045. convert_packing(src, dst, elempack, opt);
  1046. return dst;
  1047. },
  1048. py::arg("src"),
  1049. py::arg("elempack"),
  1050. py::arg("opt") = Option());
  1051. m.def("flatten", &flatten,
  1052. py::arg("src"), py::arg("dst"),
  1053. py::arg("opt") = Option());
  1054. m.def(
  1055. "flatten",
  1056. [](const Mat& src, const Option& opt) {
  1057. Mat dst;
  1058. flatten(src, dst, opt);
  1059. return dst;
  1060. },
  1061. py::arg("src"),
  1062. py::arg("opt") = Option());
  1063. m.def("cast_float32_to_float16", &cast_float32_to_float16,
  1064. py::arg("src"), py::arg("dst"),
  1065. py::arg("opt") = Option());
  1066. m.def(
  1067. "cast_float32_to_float16",
  1068. [](const Mat& src, const Option& opt) {
  1069. Mat dst;
  1070. cast_float32_to_float16(src, dst, opt);
  1071. return dst;
  1072. },
  1073. py::arg("src"),
  1074. py::arg("opt") = Option());
  1075. m.def("cast_float16_to_float32", &cast_float16_to_float32,
  1076. py::arg("src"), py::arg("dst"),
  1077. py::arg("opt") = Option());
  1078. m.def(
  1079. "cast_float16_to_float32",
  1080. [](const Mat& src, const Option& opt) {
  1081. Mat dst;
  1082. cast_float16_to_float32(src, dst, opt);
  1083. return dst;
  1084. },
  1085. py::arg("src"),
  1086. py::arg("opt") = Option());
  1087. m.def("cast_int8_to_float32", &cast_int8_to_float32,
  1088. py::arg("src"), py::arg("dst"),
  1089. py::arg("opt") = Option());
  1090. m.def(
  1091. "cast_int8_to_float32",
  1092. [](const Mat& src, const Option& opt) {
  1093. Mat dst;
  1094. cast_int8_to_float32(src, dst, opt);
  1095. return dst;
  1096. },
  1097. py::arg("src"),
  1098. py::arg("opt") = Option());
  1099. m.def("cast_float32_to_bfloat16", &cast_float32_to_bfloat16,
  1100. py::arg("src"), py::arg("dst"),
  1101. py::arg("opt") = Option());
  1102. m.def(
  1103. "cast_float32_to_bfloat16",
  1104. [](const Mat& src, const Option& opt) {
  1105. Mat dst;
  1106. cast_float32_to_bfloat16(src, dst, opt);
  1107. return dst;
  1108. },
  1109. py::arg("src"),
  1110. py::arg("opt") = Option());
  1111. m.def("cast_bfloat16_to_float32", &cast_bfloat16_to_float32,
  1112. py::arg("src"), py::arg("dst"),
  1113. py::arg("opt") = Option());
  1114. m.def(
  1115. "cast_bfloat16_to_float32",
  1116. [](const Mat& src, const Option& opt) {
  1117. Mat dst;
  1118. cast_bfloat16_to_float32(src, dst, opt);
  1119. return dst;
  1120. },
  1121. py::arg("src"),
  1122. py::arg("opt") = Option());
  1123. m.def("quantize_to_int8", &quantize_to_int8,
  1124. py::arg("src"), py::arg("dst"),
  1125. py::arg("scale_data"),
  1126. py::arg("opt") = Option());
  1127. m.def(
  1128. "quantize_to_int8",
  1129. [](const Mat& src, const Mat& scale_data, const Option& opt) {
  1130. Mat dst;
  1131. quantize_to_int8(src, dst, scale_data, opt);
  1132. return dst;
  1133. },
  1134. py::arg("src"),
  1135. py::arg("scale_data"),
  1136. py::arg("opt") = Option());
  1137. #if NCNN_STRING
  1138. m.def("layer_to_index", &layer_to_index, py::arg("type"));
  1139. m.def(
  1140. "create_layer",
  1141. [](const char* type) {
  1142. return static_cast<Layer*>(create_layer(type));
  1143. },
  1144. py::arg("type"));
  1145. m.def(
  1146. "create_layer",
  1147. [](int index) {
  1148. return static_cast<Layer*>(create_layer(index));
  1149. },
  1150. py::arg("index"));
  1151. #endif //NCNN_STRING
  1152. #if NCNN_VULKAN
  1153. m.def("create_gpu_instance", &create_gpu_instance, py::arg("driver_path") = ((const char*)0));
  1154. m.def("destroy_gpu_instance", &destroy_gpu_instance);
  1155. m.def("get_gpu_count", &get_gpu_count);
  1156. m.def("get_default_gpu_index", &get_default_gpu_index);
  1157. m.def("get_gpu_info", &get_gpu_info, py::arg("device_index") = 0, py::return_value_policy::reference);
  1158. m.def("get_gpu_device", &get_gpu_device, py::arg("device_index") = 0, py::return_value_policy::reference);
  1159. py::class_<VkBufferMemory>(m, "VkBufferMemory")
  1160. .def_readwrite("offset", &VkBufferMemory::offset)
  1161. .def_readwrite("capacity", &VkBufferMemory::capacity)
  1162. .def_readwrite("refcount", &VkBufferMemory::refcount);
  1163. py::class_<VkImageMemory>(m, "VkImageMemory")
  1164. .def_readwrite("width", &VkImageMemory::width)
  1165. .def_readwrite("height", &VkImageMemory::height)
  1166. .def_readwrite("depth", &VkImageMemory::depth)
  1167. .def_readwrite("refcount", &VkImageMemory::refcount);
  1168. py::class_<VkAllocator, PyVkAllocator<> >(m, "VkAllocator")
  1169. .def_readonly("vkdev", &VkAllocator::vkdev)
  1170. .def_readwrite("buffer_memory_type_index", &VkAllocator::buffer_memory_type_index)
  1171. .def_readwrite("image_memory_type_index", &VkAllocator::image_memory_type_index)
  1172. .def_readwrite("mappable", &VkAllocator::mappable)
  1173. .def_readwrite("coherent", &VkAllocator::coherent);
  1174. py::class_<VkBlobAllocator, VkAllocator, PyVkAllocatorOther<VkBlobAllocator> >(m, "VkBlobAllocator")
  1175. .def(py::init<const VulkanDevice*>())
  1176. .def("clear", &VkBlobAllocator::clear)
  1177. .def("fastMalloc", (VkBufferMemory * (VkBlobAllocator::*)(size_t size)) & VkBlobAllocator::fastMalloc, py::return_value_policy::reference_internal)
  1178. .def("fastFree", (void (VkBlobAllocator::*)(VkBufferMemory * ptr)) & VkBlobAllocator::fastFree)
  1179. .def("fastMalloc", (VkImageMemory * (VkBlobAllocator::*)(int, int, int, size_t, int)) & VkBlobAllocator::fastMalloc, py::return_value_policy::reference_internal)
  1180. .def("fastFree", (void (VkBlobAllocator::*)(VkImageMemory * ptr)) & VkBlobAllocator::fastFree);
  1181. py::class_<VkWeightAllocator, VkAllocator, PyVkAllocatorOther<VkWeightAllocator> >(m, "VkWeightAllocator")
  1182. .def(py::init<const VulkanDevice*>())
  1183. .def("clear", &VkWeightAllocator::clear)
  1184. .def("fastMalloc", (VkBufferMemory * (VkWeightAllocator::*)(size_t size)) & VkWeightAllocator::fastMalloc, py::return_value_policy::reference_internal)
  1185. .def("fastFree", (void (VkWeightAllocator::*)(VkBufferMemory * ptr)) & VkWeightAllocator::fastFree)
  1186. .def("fastMalloc", (VkImageMemory * (VkWeightAllocator::*)(int, int, int, size_t, int)) & VkWeightAllocator::fastMalloc, py::return_value_policy::reference_internal)
  1187. .def("fastFree", (void (VkWeightAllocator::*)(VkImageMemory * ptr)) & VkWeightAllocator::fastFree);
  1188. py::class_<VkStagingAllocator, VkAllocator, PyVkAllocatorOther<VkStagingAllocator> >(m, "VkStagingAllocator")
  1189. .def(py::init<const VulkanDevice*>())
  1190. .def("set_size_compare_ratio", &VkStagingAllocator::set_size_compare_ratio)
  1191. .def("clear", &VkStagingAllocator::clear)
  1192. .def("fastMalloc", (VkBufferMemory * (VkStagingAllocator::*)(size_t size)) & VkStagingAllocator::fastMalloc, py::return_value_policy::reference_internal)
  1193. .def("fastFree", (void (VkStagingAllocator::*)(VkBufferMemory * ptr)) & VkStagingAllocator::fastFree)
  1194. .def("fastMalloc", (VkImageMemory * (VkStagingAllocator::*)(int, int, int, size_t, int)) & VkStagingAllocator::fastMalloc, py::return_value_policy::reference_internal)
  1195. .def("fastFree", (void (VkStagingAllocator::*)(VkImageMemory * ptr)) & VkStagingAllocator::fastFree);
  1196. py::class_<VkWeightStagingAllocator, VkAllocator, PyVkAllocatorOther<VkWeightStagingAllocator> >(m, "VkWeightStagingAllocator")
  1197. .def(py::init<const VulkanDevice*>())
  1198. .def("fastMalloc", (VkBufferMemory * (VkWeightStagingAllocator::*)(size_t size)) & VkWeightStagingAllocator::fastMalloc, py::return_value_policy::reference_internal)
  1199. .def("fastFree", (void (VkWeightStagingAllocator::*)(VkBufferMemory * ptr)) & VkWeightStagingAllocator::fastFree)
  1200. .def("fastMalloc", (VkImageMemory * (VkWeightStagingAllocator::*)(int, int, int, size_t, int)) & VkWeightStagingAllocator::fastMalloc, py::return_value_policy::reference_internal)
  1201. .def("fastFree", (void (VkWeightStagingAllocator::*)(VkImageMemory * ptr)) & VkWeightStagingAllocator::fastFree);
  1202. py::class_<GpuInfo>(m, "GpuInfo")
  1203. .def(py::init<>())
  1204. .def("api_version", &GpuInfo::api_version)
  1205. .def("driver_version", &GpuInfo::driver_version)
  1206. .def("vendor_id", &GpuInfo::vendor_id)
  1207. .def("device_id", &GpuInfo::device_id)
  1208. .def("pipeline_cache_uuid", [](GpuInfo& gpuinfo) {
  1209. return py::memoryview::from_buffer(gpuinfo.pipeline_cache_uuid(), {VK_UUID_SIZE}, {sizeof(uint8_t) * VK_UUID_SIZE});
  1210. })
  1211. .def("type", &GpuInfo::type)
  1212. .def("device_name", &GpuInfo::device_name);
  1213. py::class_<VulkanDevice>(m, "VulkanDevice")
  1214. .def(py::init<int>(), py::arg("device_index") = 0)
  1215. .def(
  1216. "info", [](VulkanDevice& dev) {
  1217. return &dev.info;
  1218. },
  1219. py::return_value_policy::reference_internal)
  1220. .def("acquire_blob_allocator", &VulkanDevice::acquire_blob_allocator)
  1221. .def("reclaim_blob_allocator", &VulkanDevice::reclaim_blob_allocator, py::arg("vkallocator"))
  1222. .def("acquire_staging_allocator", &VulkanDevice::acquire_staging_allocator)
  1223. .def("reclaim_staging_allocator", &VulkanDevice::reclaim_staging_allocator, py::arg("vkallocator"))
  1224. .def("get_heap_budget", &VulkanDevice::get_heap_budget);
  1225. #endif // NCNN_VULKAN
  1226. m.doc() = R"pbdoc(
  1227. ncnn python wrapper
  1228. -----------------------
  1229. .. currentmodule:: pyncnn
  1230. .. autosummary::
  1231. :toctree: _generate
  1232. )pbdoc";
  1233. #ifdef VERSION_INFO
  1234. m.attr("__version__") = VERSION_INFO;
  1235. #else
  1236. m.attr("__version__") = "dev";
  1237. #endif
  1238. }