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

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