Ncnn load net refactor

backend/src/nodes/ncnn_nodes.py

@@ -4,26 +4,27 @@
 from __future__ import annotations
 
 import os
+from contextlib import contextmanager
 from typing import Tuple
 
 import numpy as np
 from ncnn_vulkan import ncnn
 from sanic.log import logger
-from contextlib import contextmanager
 
 from .categories import NCNNCategory
+
+# NCNN Save Model node
+# pylint: disable=unused-import
+from .model_save_nodes import NcnnSaveNode
 from .node_base import NodeBase
 from .node_factory import NodeFactory
 from .properties.inputs import *
 from .properties.outputs import *
+from .utils.exec_options import get_execution_options
 from .utils.ncnn_auto_split import ncnn_auto_split_process
 from .utils.ncnn_model import NcnnModel
-from .utils.utils import get_h_w_c, convenient_upscale
-from .utils.exec_options import get_execution_options
-
-# NCNN Save Model node
-# pylint: disable=unused-import
-from .model_save_nodes import NcnnSaveNode
+from .utils.ncnn_session import get_ncnn_net
+from .utils.utils import convenient_upscale, get_h_w_c
 
 
 @contextmanager
@@ -127,19 +128,10 @@ def upscale(
 
     def run(self, model: NcnnModel, img: np.ndarray, tile_mode: int) -> np.ndarray:
         exec_options = get_execution_options()
+        net = get_ncnn_net(model, exec_options)
 
         model_c = model.get_model_in_nc()
 
-        net = ncnn.Net()
-
-        # Use vulkan compute
-        net.opt.use_vulkan_compute = True
-        net.set_vulkan_device(exec_options.ncnn_gpu_index)
-
-        # Load model param and bin
-        net.load_param_mem(model.write_param())
-        net.load_model_mem(model.weights_bin)
-
         def upscale(i: np.ndarray) -> np.ndarray:
             ic = get_h_w_c(i)[2]
             if ic == 3:

backend/src/nodes/utils/ncnn_auto_split.py

@@ -32,7 +32,7 @@ def ncnn_auto_split_process(
     overlap: int = 16,
     max_depth: Union[int, None] = None,
     current_depth: int = 1,
-    input_name: str = "data",
+    input_name: str = "input",
     output_name: str = "output",
     blob_vkallocator=None,
     staging_vkallocator=None,
@@ -129,6 +129,8 @@ def ncnn_auto_split_process(
         overlap=overlap,
         max_depth=max_depth,
         current_depth=current_depth + 1,
+        input_name=input_name,
+        output_name=output_name,
         blob_vkallocator=blob_vkallocator,
         staging_vkallocator=staging_vkallocator,
     )
@@ -138,6 +140,8 @@ def ncnn_auto_split_process(
         overlap=overlap,
         max_depth=depth,
         current_depth=current_depth + 1,
+        input_name=input_name,
+        output_name=output_name,
         blob_vkallocator=blob_vkallocator,
         staging_vkallocator=staging_vkallocator,
     )
@@ -147,6 +151,8 @@ def ncnn_auto_split_process(
         overlap=overlap,
         max_depth=depth,
         current_depth=current_depth + 1,
+        input_name=input_name,
+        output_name=output_name,
         blob_vkallocator=blob_vkallocator,
         staging_vkallocator=staging_vkallocator,
     )
@@ -156,6 +162,8 @@ def ncnn_auto_split_process(
         overlap=overlap,
         max_depth=depth,
         current_depth=current_depth + 1,
+        input_name=input_name,
+        output_name=output_name,
         blob_vkallocator=blob_vkallocator,
         staging_vkallocator=staging_vkallocator,
     )

backend/src/nodes/utils/ncnn_session.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+from ncnn_vulkan import ncnn
+from weakref import WeakKeyDictionary
+
+from .exec_options import ExecutionOptions
+from .ncnn_model import NcnnModel
+
+
+def create_ncnn_net(model: NcnnModel, exec_options: ExecutionOptions) -> ncnn.Net:
+    net = ncnn.Net()
+
+    # Use vulkan compute
+    net.opt.use_vulkan_compute = True
+    net.set_vulkan_device(exec_options.ncnn_gpu_index)
+
+    # Load model param and bin
+    net.load_param_mem(model.write_param())
+    net.load_model_mem(model.weights_bin)
+
+    # Attribute not needed anymore, save memory
+    model.weights_bin = b""
+
+    return net
+
+
+__session_cache: WeakKeyDictionary[NcnnModel, ncnn.Net] = WeakKeyDictionary()
+
+
+def get_ncnn_net(model: NcnnModel, exec_options: ExecutionOptions) -> ncnn.Net:
+    cached = __session_cache.get(model)
+    if cached is None:
+        cached = create_ncnn_net(model, exec_options)
+        __session_cache[model] = cached
+    return cached