pnnx: support npy input tensors

Co-authored-by: lovedream-ms <209071654+lovedream-ms@users.noreply.github.com>
Co-authored-by: lancerstadium <57707690+lancerstadium@users.noreply.github.com>
Co-authored-by: AtomAlpaca <66774326+AtomAlpaca@users.noreply.github.com>

Author: 4 people

tools/pnnx/src/ir.cpp

@@ -107,6 +107,51 @@ static size_t type_to_elemsize(int type)
     return 0; // null
 }
 
+static std::string python_string_literal(const std::string& s)
+{
+    std::string r = "'";
+    for (char ch : s)
+    {
+        if (ch == '\\' || ch == '\'')
+        {
+            r += '\\';
+            r += ch;
+        }
+        else if (ch == '\n')
+        {
+            r += "\\n";
+        }
+        else if (ch == '\r')
+        {
+            r += "\\r";
+        }
+        else if (ch == '\t')
+        {
+            r += "\\t";
+        }
+        else
+        {
+            r += ch;
+        }
+    }
+    r += "'";
+    return r;
+}
+
+static void write_python_load_input_helper(FILE* pyfp)
+{
+    fprintf(pyfp, "\n");
+    fprintf(pyfp, "def _pnnx_load_input(path):\n");
+    fprintf(pyfp, "    arr = np.load(path)\n");
+    fprintf(pyfp, "    if arr.dtype.byteorder not in ('=', '|'):\n");
+    fprintf(pyfp, "        native = '<' if np.little_endian else '>'\n");
+    fprintf(pyfp, "        if arr.dtype.byteorder != native:\n");
+    fprintf(pyfp, "            arr = arr.byteswap().view(arr.dtype.newbyteorder('='))\n");
+    fprintf(pyfp, "    if not arr.flags.c_contiguous:\n");
+    fprintf(pyfp, "        arr = np.ascontiguousarray(arr)\n");
+    fprintf(pyfp, "    return torch.from_numpy(arr)\n");
+}
+
 static int string_to_type(const char* s)
 {
     if (strcmp(s, "f32") == 0) return 1;
@@ -1456,7 +1501,7 @@ static std::string make_index_expression(const Operator* op)
     return index_expr;
 }
 
-int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, const std::vector<std::vector<int64_t> >& input_shapes)
+int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, const std::vector<std::vector<int64_t> >& input_shapes, const std::vector<std::string>& input_npy_paths)
 {
     FILE* pyfp = fopen(pypath.c_str(), "wb");
     if (!pyfp)
@@ -1477,6 +1522,9 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
     fprintf(pyfp, "except:\n");
     fprintf(pyfp, "    pass\n");
 
+    if (!input_npy_paths.empty())
+        write_python_load_input_helper(pyfp);
+
     fprintf(pyfp, "\n");
 
     fprintf(pyfp, "class Model(nn.Module):\n");
@@ -2771,6 +2819,7 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
         fprintf(pyfp, "    torch.manual_seed(0)\n");
 
         int input_shapes_i = 0;
+        int input_npy_i = 0;
 
         std::vector<std::string> input_names;
         for (const Operator* op : ops)
@@ -2779,6 +2828,18 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
                 continue;
 
             const Operand* r = op->outputs[0];
+            std::string input_name = std::string("v_") + sanitize_identifier(r->name);
+
+            if (input_npy_i < (int)input_npy_paths.size())
+            {
+                std::string path_literal = python_string_literal(input_npy_paths[input_npy_i]);
+                fprintf(pyfp, "    %s = _pnnx_load_input(%s)\n", input_name.c_str(), path_literal.c_str());
+                input_names.push_back(input_name);
+                input_npy_i++;
+                if (!input_shapes.empty())
+                    input_shapes_i++;
+                continue;
+            }
 
             std::vector<int> input_shape;
             if (input_shapes.empty())
@@ -2794,7 +2855,6 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
                 }
             }
 
-            std::string input_name = std::string("v_") + sanitize_identifier(r->name);
             if (type_is_integer(r->type))
             {
                 fprintf(pyfp, "    %s = torch.randint(10, (", input_name.c_str());

tools/pnnx/src/ir.h

@@ -324,7 +324,7 @@ class Graph
     int load(const std::string& parampath, const std::string& binpath);
     int save(const std::string& parampath, const std::string& binpath);
 
-    int python(const std::string& pypath, const std::string& binpath, const std::vector<std::vector<int64_t> >& input_shapes);
+    int python(const std::string& pypath, const std::string& binpath, const std::vector<std::vector<int64_t> >& input_shapes, const std::vector<std::string>& input_npy_paths);
 
     int parse(const std::string& param);
 

tools/pnnx/src/load_tnn.cpp

@@ -141,6 +141,26 @@ static int get_tnn_tensor_type(int dt)
     return 0; // unknown type
 }
 
+static int input_type_to_pnnx_type(const std::string& t)
+{
+    if (t == "f32") return 1;
+    if (t == "f64") return 2;
+    if (t == "f16") return 3;
+    if (t == "i32") return 4;
+    if (t == "i64") return 5;
+    if (t == "i16") return 6;
+    if (t == "i8") return 7;
+    if (t == "u8") return 8;
+    if (t == "bool") return 9;
+    if (t == "c64") return 10;
+    if (t == "c128") return 11;
+    if (t == "c32") return 12;
+    if (t == "bf16") return 13;
+
+    fprintf(stderr, "unsupported input type %s\n", t.c_str());
+    return 0;
+}
+
 Attribute::Attribute(FILE* bp)
 {
     unsigned int magic;
@@ -178,7 +198,9 @@ Attribute::Attribute(FILE* bp)
     fread((void*)data.data(), 1, length, bp);
 }
 
-int load_tnn(const std::string& tnnpath, Graph& pnnx_graph)
+int load_tnn(const std::string& tnnpath, Graph& pnnx_graph,
+             const std::vector<std::vector<int64_t> >& input_shapes,
+             const std::vector<std::string>& input_types)
 {
     fprintf(stderr, "############# pass_level0 tnn\n");
 
@@ -244,14 +266,48 @@ int load_tnn(const std::string& tnnpath, Graph& pnnx_graph)
         int datatype = 0;
         sscanf(pline, "%d%n", &datatype, &ncomsumed);
 
+        int tensor_type = get_tnn_tensor_type(datatype);
+
+        if (!input_shapes.empty())
+        {
+            if (input_shapes.size() != 1)
+            {
+                fprintf(stderr, "tnn input expect 1 tensor but got %d\n", (int)input_shapes.size());
+                return -1;
+            }
+
+            shape.clear();
+            const std::vector<int64_t>& input_shape = input_shapes[0];
+            for (size_t i = 0; i < input_shape.size(); i++)
+            {
+                if (input_shape[i] < 0 || input_shape[i] > INT_MAX)
+                {
+                    fprintf(stderr, "invalid tnn input shape dimension %ld\n", input_shape[i]);
+                    return -1;
+                }
+                shape.push_back((int)input_shape[i]);
+            }
+
+            const int input_tensor_type = input_type_to_pnnx_type(input_types[0]);
+            if (input_tensor_type == 0)
+                return -1;
+            if (tensor_type != 0 && tensor_type != input_tensor_type)
+            {
+                fprintf(stderr, "tnn input type mismatch, model type %d but got %s\n", tensor_type, input_types[0].c_str());
+                return -1;
+            }
+
+            tensor_type = input_tensor_type;
+        }
+
         Operator* op = pnnx_graph.new_operator("pnnx.Input", "input0");
 
         Operand* r = pnnx_graph.new_operand(blob_name);
 
         r->producer = op;
 
         r->shape = shape;
-        r->type = get_tnn_tensor_type(datatype);
+        r->type = tensor_type;
 
         op->outputs.push_back(r);
     }

tools/pnnx/src/load_tnn.h

@@ -8,7 +8,9 @@
 
 namespace pnnx {
 
-int load_tnn(const std::string& tnnpath, Graph& g);
+int load_tnn(const std::string& tnnpath, Graph& g,
+             const std::vector<std::vector<int64_t> >& input_shapes,
+             const std::vector<std::string>& input_types);
 
 } // namespace pnnx