tnn2pnnx (#5898)

Tencent · Feb 12, 2025 · 5ea7681 · 5ea7681
1 parent 2389090
commit 5ea7681
Show file tree

Hide file tree

Showing 50 changed files with 3,195 additions and 3 deletions.
diff --git a/tools/pnnx/CMakeLists.txt b/tools/pnnx/CMakeLists.txt
@@ -123,6 +123,8 @@ else()
     set(onnxruntime_FOUND FALSE)
 endif()
 
+option(PNNX_TNN2PNNX "build tnn2pnnx" ON)
+
 add_subdirectory(src)
 
 enable_testing()

diff --git a/tools/pnnx/src/CMakeLists.txt b/tools/pnnx/src/CMakeLists.txt
@@ -715,6 +715,29 @@ else()
     message(STATUS "Building without onnx2pnnx")
 endif()
 
+if(PNNX_TNN2PNNX)
+    set(pnnx_pass_tnn_SRCS
+        pass_tnn/fuse_shape_size.cpp
+        pass_tnn/fuse_shape_list_construct.cpp
+        pass_tnn/lower_concat.cpp
+        pass_tnn/lower_convolution_activation.cpp
+        pass_tnn/lower_power.cpp
+    )
+
+    set(tnn2pnnx_SRCS
+        ${pnnx_pass_tnn_SRCS}
+        load_tnn.cpp
+    )
+
+    add_library(tnn2pnnx OBJECT ${tnn2pnnx_SRCS})
+    target_compile_definitions(tnn2pnnx PRIVATE BUILD_TNN2PNNX)
+    target_compile_options(tnn2pnnx PUBLIC "${TORCH_CXX_FLAGS}")
+
+    message(STATUS "Building with tnn2pnnx")
+else()
+    message(STATUS "Building without tnn2pnnx")
+endif()
+
 if(NOT MSVC)
     add_definitions(-Wall -Wextra)
 endif()
@@ -765,6 +788,11 @@ if(onnxruntime_FOUND)
     target_link_libraries(pnnx PRIVATE onnx2pnnx)
 endif()
 
+if(PNNX_TNN2PNNX)
+    set_property(SOURCE main.cpp APPEND PROPERTY COMPILE_DEFINITIONS BUILD_TNN2PNNX)
+    target_link_libraries(pnnx PRIVATE tnn2pnnx)
+endif()
+
 if(PNNX_COVERAGE)
     target_compile_options(pnnx PUBLIC -coverage -fprofile-arcs -ftest-coverage)
     target_link_libraries(pnnx PUBLIC -coverage -lgcov)

diff --git a/tools/pnnx/src/ir.cpp b/tools/pnnx/src/ir.cpp
@@ -1030,6 +1030,7 @@ static std::string expand_expression(const Operator* op)
                  || t == "ceil"
                  || t == "cos"
                  || t == "cosh"
+                 || t == "erf"
                  || t == "exp"
                  || t == "floor"
                  || t == "log"
@@ -1062,6 +1063,7 @@ static std::string expand_expression(const Operator* op)
             if (t == "ceil") unaryop = "torch.ceil";
             if (t == "cos") unaryop = "torch.cos";
             if (t == "cosh") unaryop = "torch.cosh";
+            if (t == "erf") unaryop = "torch.erf";
             if (t == "exp") unaryop = "torch.exp";
             if (t == "floor") unaryop = "torch.floor";
             if (t == "log") unaryop = "torch.log";
@@ -2253,11 +2255,17 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath)
                     {
                         if (op->type == "Tensor.index_put" && it.first == "values")
                         {
-                            fprintf(pyfp, "torch.tensor(%f)", param.f);
+                            if (param.f == (int)param.f)
+                                fprintf(pyfp, "torch.tensor(%.1f)", param.f);
+                            else
+                                fprintf(pyfp, "torch.tensor(%g)", param.f);
                         }
                         else
                         {
-                            fprintf(pyfp, "%f", param.f);
+                            if (param.f == (int)param.f)
+                                fprintf(pyfp, "%.1f", param.f);
+                            else
+                                fprintf(pyfp, "%g", param.f);
                         }
                     }
                     if (param.type == 4)
@@ -2316,7 +2324,10 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath)
                         fprintf(pyfp, "(");
                         for (size_t i = 0; i < param.af.size(); i++)
                         {
-                            fprintf(pyfp, "%f", param.af[i]);
+                            if (param.af[i] == (int)param.af[i])
+                                fprintf(pyfp, "%.1f", param.af[i]);
+                            else
+                                fprintf(pyfp, "%g", param.af[i]);
                             if (i + 1 != param.af.size() || param.af.size() == 1)
                                 fprintf(pyfp, ",");
                         }

diff --git a/tools/pnnx/src/ir.h b/tools/pnnx/src/ir.h
@@ -234,6 +234,9 @@ class Attribute
 #if BUILD_ONNX2PNNX
     Attribute(const onnx::TensorProto& t);
 #endif
+#if BUILD_TNN2PNNX
+    Attribute(FILE* bp);
+#endif
 
     Attribute(const std::initializer_list<int>& shape, const std::vector<float>& t);