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 61 kB

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