diff --git a/python/src/main.cpp b/python/src/main.cpp index 4dbaf0038..fdd026ed5 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -760,17 +760,20 @@ PYBIND11_MODULE(ncnn, m) .value("PIXEL_BGRA2RGBA", ncnn::Mat::PixelType::PIXEL_BGRA2RGBA); py::class_(m, "Extractor") + .def("__enter__", [](Extractor& ex) -> Extractor& { return ex; }) + .def("__exit__", [](Extractor& ex, pybind11::args) { + ex.clear(); + }) + .def("clear", &Extractor::clear) .def("set_light_mode", &Extractor::set_light_mode, py::arg("enable")) .def("set_num_threads", &Extractor::set_num_threads, py::arg("num_threads")) .def("set_blob_allocator", &Extractor::set_blob_allocator, py::arg("allocator")) .def("set_workspace_allocator", &Extractor::set_workspace_allocator, py::arg("allocator")) #if NCNN_STRING .def("input", (int (Extractor::*)(const char*, const Mat&)) & Extractor::input, py::arg("blob_name"), py::arg("in")) - .def("extract", (int (Extractor::*)(const char*, Mat&, int)) & Extractor::extract, - py::arg("blob_name"), py::arg("feat"), py::arg("type") = 0) + .def("extract", (int (Extractor::*)(const char*, Mat&, int)) & Extractor::extract, py::arg("blob_name"), py::arg("feat"), py::arg("type") = 0) .def( - "extract", - [](Extractor& ex, const char* blob_name, int type) { + "extract", [](Extractor& ex, const char* blob_name, int type) { ncnn::Mat feat; int ret = ex.extract(blob_name, feat, type); return py::make_tuple(ret, feat); @@ -778,11 +781,9 @@ PYBIND11_MODULE(ncnn, m) py::arg("blob_name"), py::arg("type") = 0) #endif .def("input", (int (Extractor::*)(int, const Mat&)) & Extractor::input) - .def("extract", (int (Extractor::*)(int, Mat&, int)) & Extractor::extract, - py::arg("blob_index"), py::arg("feat"), py::arg("type") = 0) + .def("extract", (int (Extractor::*)(int, Mat&, int)) & Extractor::extract, py::arg("blob_index"), py::arg("feat"), py::arg("type") = 0) .def( - "extract", - [](Extractor& ex, int blob_index, int type) { + "extract", [](Extractor& ex, int blob_index, int type) { ncnn::Mat feat; int ret = ex.extract(blob_index, feat, type); return py::make_tuple(ret, feat); @@ -825,6 +826,10 @@ PYBIND11_MODULE(ncnn, m) py::class_(m, "Net") .def(py::init<>()) .def_readwrite("opt", &Net::opt) + .def("__enter__", [](Net& net) -> Net& { return net; }) + .def("__exit__", [](Net& net, pybind11::args) { + net.clear(); + }) #if NCNN_VULKAN .def("set_vulkan_device", (void (Net::*)(int)) & Net::set_vulkan_device, py::arg("device_index")) @@ -834,10 +839,7 @@ PYBIND11_MODULE(ncnn, m) #if NCNN_STRING .def( - "register_custom_layer", - [](Net& net, const char* type, - const std::function& creator, - const std::function& destroyer) { + "register_custom_layer", [](Net& net, const char* type, const std::function& creator, const std::function& destroyer) { if (g_layer_factroy_index == g_layer_factroys.size()) { std::stringstream ss; @@ -853,10 +855,7 @@ PYBIND11_MODULE(ncnn, m) py::arg("type"), py::arg("creator"), py::arg("destroyer")) #endif //NCNN_STRING .def( - "register_custom_layer", - [](Net& net, int index, - const std::function& creator, - const std::function& destroyer) { + "register_custom_layer", [](Net& net, int index, const std::function& creator, const std::function& destroyer) { if (g_layer_factroy_index == g_layer_factroys.size()) { std::stringstream ss; @@ -1114,8 +1113,8 @@ PYBIND11_MODULE(ncnn, m) m.def("destroy_gpu_instance", &destroy_gpu_instance); m.def("get_gpu_count", &get_gpu_count); m.def("get_default_gpu_index", &get_default_gpu_index); - m.def("get_gpu_info", &get_gpu_info, py::arg("device_index") = 0); - m.def("get_gpu_device", &get_gpu_device, py::arg("device_index") = 0); + m.def("get_gpu_info", &get_gpu_info, py::arg("device_index") = 0, py::return_value_policy::reference_internal); + m.def("get_gpu_device", &get_gpu_device, py::arg("device_index") = 0, py::return_value_policy::reference_internal); py::class_ >(m, "VkAllocator") .def_readwrite("vkdev", &VkAllocator::vkdev) diff --git a/python/tests/benchmark.py b/python/tests/benchmark.py index 15955297f..21e5804ba 100644 --- a/python/tests/benchmark.py +++ b/python/tests/benchmark.py @@ -56,9 +56,9 @@ def benchmark(comment, _in, opt): # warm up for i in range(g_warmup_loop_count): - ex = net.create_extractor() - ex.input("data", _in) - ex.extract("output") + with net.create_extractor() as ex: + ex.input("data", _in) + ex.extract("output") time_min = sys.float_info.max time_max = -sys.float_info.max @@ -67,9 +67,9 @@ def benchmark(comment, _in, opt): for i in range(g_loop_count): start = time.time() - ex = net.create_extractor() - ex.input("data", _in) - ex.extract("output") + with net.create_extractor() as ex: + ex.input("data", _in) + ex.extract("output") end = time.time() @@ -79,10 +79,6 @@ def benchmark(comment, _in, opt): time_max = timespan if timespan > time_max else time_max time_avg += timespan - # extractor need relese manually when build ncnn with vuklan, - # due to python relese ex after net, but in extractor.destruction use net - ex = None - time_avg /= g_loop_count print( diff --git a/python/tests/test_extractor.py b/python/tests/test_extractor.py index 610829282..5fb5fa010 100644 --- a/python/tests/test_extractor.py +++ b/python/tests/test_extractor.py @@ -31,25 +31,25 @@ def test_extractor(): net.load_model(dr) in_mat = ncnn.Mat((227, 227, 3)) - ex = net.create_extractor() - ex.set_light_mode(True) - ex.set_num_threads(2) + with net.create_extractor() as ex: + ex.set_light_mode(True) + ex.set_num_threads(2) - ex.set_blob_allocator(alloctor) - ex.set_workspace_allocator(alloctor) + ex.set_blob_allocator(alloctor) + ex.set_workspace_allocator(alloctor) - ex.input("data", in_mat) - ret, out_mat = ex.extract("conv0_fwd") - assert ( - ret == 0 - and out_mat.dims == 3 - and out_mat.w == 225 - and out_mat.h == 225 - and out_mat.c == 3 - ) + ex.input("data", in_mat) + ret, out_mat = ex.extract("conv0_fwd") + assert ( + ret == 0 + and out_mat.dims == 3 + and out_mat.w == 225 + and out_mat.h == 225 + and out_mat.c == 3 + ) - ret, out_mat = ex.extract("output") - assert ret == 0 and out_mat.dims == 1 and out_mat.w == 1 + ret, out_mat = ex.extract("output") + assert ret == 0 and out_mat.dims == 1 and out_mat.w == 1 def test_extractor_index(): @@ -83,3 +83,6 @@ def test_extractor_index(): ret, out_mat = ex.extract(2) assert ret == 0 and out_mat.dims == 1 and out_mat.w == 1 + + # not use with sentence, call clear manually to ensure ex destruct before net + ex.clear() diff --git a/python/tests/test_net.py b/python/tests/test_net.py index e312902c1..2a96ece81 100644 --- a/python/tests/test_net.py +++ b/python/tests/test_net.py @@ -21,21 +21,21 @@ import ncnn def test_net(): dr = ncnn.DataReaderFromEmpty() - net = ncnn.Net() - ret = net.load_param("tests/test.param") - net.load_model(dr) - assert ret == 0 and len(net.blobs()) == 3 and len(net.layers()) == 3 + with ncnn.Net() as net: + ret = net.load_param("tests/test.param") + net.load_model(dr) + assert ret == 0 and len(net.blobs()) == 3 and len(net.layers()) == 3 - in_mat = ncnn.Mat((227, 227, 3)) + in_mat = ncnn.Mat((227, 227, 3)) - ex = net.create_extractor() - ex.input("data", in_mat) - ret, out_mat = ex.extract("output") + with net.create_extractor() as ex: + ex.input("data", in_mat) + ret, out_mat = ex.extract("output") - assert ret == 0 and out_mat.dims == 1 and out_mat.w == 1 + assert ret == 0 and out_mat.dims == 1 and out_mat.w == 1 - net.clear() - assert len(net.blobs()) == 0 and len(net.layers()) == 0 + net.clear() + assert len(net.blobs()) == 0 and len(net.layers()) == 0 def test_custom_layer(): @@ -85,5 +85,7 @@ def test_custom_layer(): ret, out_mat = ex.extract("output") assert ret == 0 and out_mat.dims == 1 and out_mat.w == 1 and out_mat[0] == 2.0 + ex.clear() + net.clear() assert len(net.blobs()) == 0 and len(net.layers()) == 0 diff --git a/src/net.cpp b/src/net.cpp index f9f6900e2..404eaca03 100644 --- a/src/net.cpp +++ b/src/net.cpp @@ -2264,24 +2264,7 @@ Extractor::Extractor(const Net* _net, size_t blob_count) Extractor::~Extractor() { - d->blob_mats.clear(); - -#if NCNN_VULKAN - if (d->net->opt.use_vulkan_compute) - { - d->blob_mats_gpu.clear(); - d->blob_mats_gpu_image.clear(); - - if (d->local_blob_vkallocator) - { - d->net->vulkan_device()->reclaim_blob_allocator(d->local_blob_vkallocator); - } - if (d->local_staging_vkallocator) - { - d->net->vulkan_device()->reclaim_staging_allocator(d->local_staging_vkallocator); - } - } -#endif // NCNN_VULKAN + clear(); delete d; } @@ -2322,6 +2305,28 @@ Extractor& Extractor::operator=(const Extractor& rhs) return *this; } +void Extractor::clear() +{ + d->blob_mats.clear(); + +#if NCNN_VULKAN + if (d->opt.use_vulkan_compute) + { + d->blob_mats_gpu.clear(); + d->blob_mats_gpu_image.clear(); + + if (d->local_blob_vkallocator) + { + d->net->vulkan_device()->reclaim_blob_allocator(d->local_blob_vkallocator); + } + if (d->local_staging_vkallocator) + { + d->net->vulkan_device()->reclaim_staging_allocator(d->local_staging_vkallocator); + } + } +#endif // NCNN_VULKAN +} + void Extractor::set_light_mode(bool enable) { d->opt.lightmode = enable; diff --git a/src/net.h b/src/net.h index 26aee43f1..262a4ba04 100644 --- a/src/net.h +++ b/src/net.h @@ -164,6 +164,9 @@ public: // assign Extractor& operator=(const Extractor&); + // clear blob mats and alloctors + void clear(); + // enable light mode // intermediate blob will be recycled when enabled // enabled by default