add custom layer destroyer (#2481)

* add custom layer destroyer

* set default layer destroyer with 0

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()

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

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

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<Mat>& blob_mats, const Option& opt) const

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);