diff --git a/cmake/ncnn_add_layer.cmake b/cmake/ncnn_add_layer.cmake index cddde852b..48c5638b7 100644 --- a/cmake/ncnn_add_layer.cmake +++ b/cmake/ncnn_add_layer.cmake @@ -76,9 +76,9 @@ macro(ncnn_add_layer class) endif() if(WITH_LAYER_${name}) - set(layer_registry "${layer_registry}#if NCNN_STRING\n{\"${class}\",${class}_final_layer_creator},\n#else\n{${class}_final_layer_creator},\n#endif\n") + set(layer_registry "${layer_registry}#if NCNN_STRING\n{\"${class}\", ${class}_final_layer_creator, 0},\n#else\n{${class}_final_layer_creator, 0},\n#endif\n") else() - set(layer_registry "${layer_registry}#if NCNN_STRING\n{\"${class}\",0},\n#else\n{0},\n#endif\n") + set(layer_registry "${layer_registry}#if NCNN_STRING\n{\"${class}\", 0, 0},\n#else\n{0, 0},\n#endif\n") endif() @@ -143,13 +143,13 @@ macro(ncnn_add_layer class) set(layer_declaration "${layer_declaration}};\n") set(layer_declaration "${layer_declaration}DEFINE_LAYER_CREATOR(${class}_final_avx2)\n} // namespace ncnn\n\n") - set(layer_registry_avx2 "${layer_registry_avx2}#if NCNN_STRING\n{\"${class}\",${class}_final_avx2_layer_creator},\n#else\n{${class}_final_avx2_layer_creator},\n#endif\n") + set(layer_registry_avx2 "${layer_registry_avx2}#if NCNN_STRING\n{\"${class}\", ${class}_final_avx2_layer_creator, 0},\n#else\n{${class}_final_avx2_layer_creator, 0},\n#endif\n") else() # no arm optimized version if(WITH_LAYER_${name}) - set(layer_registry_avx2 "${layer_registry_avx2}#if NCNN_STRING\n{\"${class}\",${class}_final_layer_creator},\n#else\n{${class}_final_layer_creator},\n#endif\n") + set(layer_registry_avx2 "${layer_registry_avx2}#if NCNN_STRING\n{\"${class}\", ${class}_final_layer_creator, 0},\n#else\n{${class}_final_layer_creator, 0},\n#endif\n") else() - set(layer_registry_avx2 "${layer_registry_avx2}#if NCNN_STRING\n{\"${class}\",0},\n#else\n{0},\n#endif\n") + set(layer_registry_avx2 "${layer_registry_avx2}#if NCNN_STRING\n{\"${class}\", 0, 0},\n#else\n{0, 0},\n#endif\n") endif() endif() endif() @@ -211,13 +211,13 @@ macro(ncnn_add_layer class) set(layer_declaration "${layer_declaration}};\n") set(layer_declaration "${layer_declaration}DEFINE_LAYER_CREATOR(${class}_final_arm82)\n} // namespace ncnn\n\n") - set(layer_registry_arm82 "${layer_registry_arm82}#if NCNN_STRING\n{\"${class}\",${class}_final_arm82_layer_creator},\n#else\n{${class}_final_arm82_layer_creator},\n#endif\n") + set(layer_registry_arm82 "${layer_registry_arm82}#if NCNN_STRING\n{\"${class}\", ${class}_final_arm82_layer_creator, 0},\n#else\n{${class}_final_arm82_layer_creator, 0},\n#endif\n") else() # no arm optimized version if(WITH_LAYER_${name}) - set(layer_registry_arm82 "${layer_registry_arm82}#if NCNN_STRING\n{\"${class}\",${class}_final_layer_creator},\n#else\n{${class}_final_layer_creator},\n#endif\n") + set(layer_registry_arm82 "${layer_registry_arm82}#if NCNN_STRING\n{\"${class}\", ${class}_final_layer_creator, 0},\n#else\n{${class}_final_layer_creator, 0},\n#endif\n") else() - set(layer_registry_arm82 "${layer_registry_arm82}#if NCNN_STRING\n{\"${class}\",0},\n#else\n{0},\n#endif\n") + set(layer_registry_arm82 "${layer_registry_arm82}#if NCNN_STRING\n{\"${class}\", 0, 0},\n#else\n{0, 0},\n#endif\n") endif() endif() endif() diff --git a/src/layer.cpp b/src/layer.cpp index cb600e1ab..67a3a4e52 100644 --- a/src/layer.cpp +++ b/src/layer.cpp @@ -44,6 +44,8 @@ Layer::Layer() use_int8_inference = false; support_weight_fp16_storage = false; + typeindex = -1; + #if NCNN_VULKAN vkdev = 0; #endif // NCNN_VULKAN diff --git a/src/layer.h b/src/layer.h index 373f41725..554f7a8eb 100644 --- a/src/layer.h +++ b/src/layer.h @@ -144,6 +144,7 @@ public: // layer factory function typedef Layer* (*layer_creator_func)(); +typedef void (*layer_destroyer_func)(Layer*); struct layer_registry_entry { @@ -153,6 +154,7 @@ struct layer_registry_entry #endif // NCNN_STRING // layer factory entry layer_creator_func creator; + layer_destroyer_func destroyer; }; #if NCNN_STRING @@ -170,6 +172,12 @@ Layer* create_layer(int index); return new name; \ } +#define DEFINE_LAYER_DESTROYER(name) \ + void name##_layer_destroyer(::ncnn::Layer* layer) \ + { \ + delete layer; \ + } + } // namespace ncnn #endif // NCNN_LAYER_H diff --git a/src/net.cpp b/src/net.cpp index fe3ad50ea..0e48caed6 100644 --- a/src/net.cpp +++ b/src/net.cpp @@ -57,7 +57,7 @@ Net::~Net() } #if NCNN_STRING -int Net::register_custom_layer(const char* type, layer_creator_func creator) +int Net::register_custom_layer(const char* type, layer_creator_func creator, layer_destroyer_func destroyer) { int typeindex = layer_to_index(type); if (typeindex != -1) @@ -69,7 +69,7 @@ int Net::register_custom_layer(const char* type, layer_creator_func creator) int custom_index = custom_layer_to_index(type); if (custom_index == -1) { - struct layer_registry_entry entry = {type, creator}; + struct layer_registry_entry entry = {type, creator, destroyer}; custom_layer_registry.push_back(entry); } else @@ -77,13 +77,14 @@ int Net::register_custom_layer(const char* type, layer_creator_func creator) NCNN_LOGE("overwrite existing custom layer type %s", type); custom_layer_registry[custom_index].name = type; custom_layer_registry[custom_index].creator = creator; + custom_layer_registry[custom_index].destroyer = destroyer; } return 0; } #endif // NCNN_STRING -int Net::register_custom_layer(int index, layer_creator_func creator) +int Net::register_custom_layer(int index, layer_creator_func creator, layer_destroyer_func destroyer) { int custom_index = index & ~LayerType::CustomBit; if (index == custom_index) @@ -108,6 +109,7 @@ int Net::register_custom_layer(int index, layer_creator_func creator) } custom_layer_registry[custom_index].creator = creator; + custom_layer_registry[custom_index].destroyer = destroyer; return 0; } @@ -948,7 +950,22 @@ void Net::clear() // ignore anyway } - delete layer; + if (layer->typeindex & ncnn::LayerType::CustomBit) + { + int custom_index = layer->typeindex & ~ncnn::LayerType::CustomBit; + if (custom_layer_registry[custom_index].destroyer) + { + custom_layer_registry[custom_index].destroyer(layer); + } + else + { + delete layer; + } + } + else + { + delete layer; + } } layers.clear(); @@ -1104,7 +1121,9 @@ Layer* Net::create_custom_layer(int index) if (!layer_creator) return 0; - return layer_creator(); + Layer* layer = layer_creator(); + layer->typeindex = ncnn::LayerType::CustomBit | index; + return layer; } int Net::forward_layer(int layer_index, std::vector& blob_mats, const Option& opt) const diff --git a/src/net.h b/src/net.h index a7c05bd0f..b78bd05d7 100644 --- a/src/net.h +++ b/src/net.h @@ -58,11 +58,11 @@ public: #if NCNN_STRING // register custom layer by layer type name // return 0 if success - int register_custom_layer(const char* type, layer_creator_func creator); + int register_custom_layer(const char* type, layer_creator_func creator, layer_destroyer_func destroyer = NULL); #endif // NCNN_STRING // register custom layer by layer type // return 0 if success - int register_custom_layer(int index, layer_creator_func creator); + int register_custom_layer(int index, layer_creator_func creator, layer_destroyer_func destroyer = NULL); #if NCNN_STRING int load_param(const DataReader& dr);