pnnx user provided parameter (#3417)

tools/pnnx/README.md

@@ -549,9 +549,9 @@ TORCH_LIBRARY(upfirdn2d_op, m) {
 |F.celu                     | :heavy_check_mark: |
 |F.conv1d                   | :heavy_check_mark: | :heavy_check_mark: |
 |F.conv2d                   | :heavy_check_mark: | :heavy_check_mark: |
-|F.conv3d                   | :heavy_check_mark: |
+|F.conv3d                   | :heavy_check_mark: | :heavy_check_mark: |
 |F.conv_transpose1d         | :heavy_check_mark: |
-|F.conv_transpose2d         | :heavy_check_mark: |
+|F.conv_transpose2d         | :heavy_check_mark: | :heavy_check_mark: |
 |F.conv_transpose3d         | :heavy_check_mark: |
 |F.cosine_similarity        |  |
 |F.dropout                  |  |

tools/pnnx/src/CMakeLists.txt

@@ -109,9 +109,7 @@ set(pnnx_pass_level2_SRCS
     pass_level2/F_conv1d.cpp
     pass_level2/F_conv2d.cpp
     pass_level2/F_conv3d.cpp
-    pass_level2/F_conv_transpose1d.cpp
-    pass_level2/F_conv_transpose2d.cpp
-    pass_level2/F_conv_transpose3d.cpp
+    pass_level2/F_conv_transpose123d.cpp
     pass_level2/F_elu.cpp
     pass_level2/F_gelu.cpp
     pass_level2/F_grid_sample.cpp
@@ -186,6 +184,7 @@ set(pnnx_pass_level3_SRCS
     pass_level3/fuse_chunk_split_unpack.cpp
     pass_level3/fuse_expression.cpp
     pass_level3/fuse_rnn_unpack.cpp
+    pass_level3/rename_F_conv_transposend.cpp
     pass_level3/rename_F_convmode.cpp
 )
 
@@ -242,8 +241,10 @@ set(pnnx_pass_ncnn_SRCS
     pass_ncnn/F_avg_pool2d.cpp
     pass_ncnn/F_avg_pool3d.cpp
     pass_ncnn/F_batch_norm.cpp
+    pass_ncnn/F_conv_transpose2d.cpp
     pass_ncnn/F_conv1d.cpp
     pass_ncnn/F_conv2d.cpp
+    pass_ncnn/F_conv3d.cpp
     pass_ncnn/F_elu.cpp
     pass_ncnn/F_gelu.cpp
     pass_ncnn/F_group_norm.cpp

tools/pnnx/src/ir.cpp

@@ -1254,6 +1254,49 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath)
             }
         }
 
+        for (const Operator* op : ops)
+        {
+            if (op->type != "pnnx.Attribute")
+                continue;
+
+            const std::string& key = op->attrs.begin()->first;
+            const Attribute& attr = op->attrs.begin()->second;
+
+            bool is_running_mean_var = false;
+            {
+                const Operand* r = op->outputs[0];
+                if (r->consumers.size() == 1)
+                {
+                    const Operator* op2 = r->consumers[0];
+                    if (op2->type == "F.batch_norm" || op2->type == "F.instance_norm")
+                    {
+                        if (r == op2->inputs[1] || r == op2->inputs[2])
+                        {
+                            is_running_mean_var = true;
+                        }
+                    }
+                }
+            }
+
+            if (is_running_mean_var)
+            {
+                fprintf(pyfp, "        self.%s = self.load_pnnx_bin_as_tensor(archive, '%s.%s', (", key.c_str(), sanitize_identifier(op->name).c_str(), key.c_str());
+            }
+            else
+            {
+                fprintf(pyfp, "        self.%s = self.load_pnnx_bin_as_parameter(archive, '%s.%s', (", key.c_str(), sanitize_identifier(op->name).c_str(), key.c_str());
+            }
+
+            for (size_t i = 0; i < attr.shape.size(); i++)
+            {
+                fprintf(pyfp, "%d", attr.shape[i]);
+                if (i + 1 != attr.shape.size())
+                    fprintf(pyfp, ",");
+            }
+
+            fprintf(pyfp, "), '%s')\n", type_to_numpy_string(attr.type));
+        }
+
         fprintf(pyfp, "        archive.close()\n");
     }
 
@@ -1313,6 +1356,11 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath)
                 std::string expanded_expr = expand_expression(op);
                 fprintf(pyfp, " = %s\n", expanded_expr.c_str());
             }
+            else if (op->type == "pnnx.Attribute")
+            {
+                const std::string& key = op->attrs.begin()->first;
+                fprintf(pyfp, "v_%s = self.%s\n", sanitize_identifier(op->outputs[0]->name).c_str(), key.c_str());
+            }
             else if (op->type == "Tensor.slice")
             {
                 // slice expr

tools/pnnx/src/pass_level1.cpp

@@ -118,11 +118,17 @@ void pass_level1(const torch::jit::Module& mod, const std::shared_ptr<torch::jit
                     sub_mod = sub_mod.attr(module_name).toModule();
                 }
 
+                if (wrapped_name.empty())
+                {
+                    // top-level module
+                    wrapped_name = name;
+                }
+
                 op->name = wrapped_name;
 
-                //                 op->params["this"] = n->input(i)
+                // op->params["this"] = n->input(i)
 
-                //                 sub_mod.dump(true, true, true);
+                // sub_mod.dump(true, true, true);
 
                 op->attrs[name] = sub_mod.attr(name).toTensor();
             }

tools/pnnx/src/pass_level2/F_conv_transpose1d.cpp → tools/pnnx/src/pass_level2/F_conv_transpose123d.cpp

@@ -16,7 +16,7 @@
 
 namespace pnnx {
 
-class F_conv_transpose1d : public GraphRewriterPass
+class F_conv_transposend : public GraphRewriterPass
 {
 public:
     const char* match_pattern_graph() const
@@ -43,10 +43,10 @@ pnnx.Output             output      1 0 out
 
     const char* type_str() const
     {
-        return "F.conv_transpose1d";
+        return "F.conv_transposend";
     }
 };
 
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_conv_transpose1d, 10)
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_conv_transposend, 10)
 
 } // namespace pnnx

tools/pnnx/src/pass_level2/F_conv_transpose2d.cpp

@@ -1,52 +0,0 @@
-// Tencent is pleased to support the open source community by making ncnn available.
-//
-// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
-//
-// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
-// in compliance with the License. You may obtain a copy of the License at
-//
-// https://opensource.org/licenses/BSD-3-Clause
-//
-// Unless required by applicable law or agreed to in writing, software distributed
-// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
-// CONDITIONS OF ANY KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations under the License.
-
-#include "pass_level2.h"
-
-namespace pnnx {
-
-class F_conv_transpose2d : public GraphRewriterPass
-{
-public:
-    const char* match_pattern_graph() const
-    {
-        return R"PNNXIR(7767517
-15 14
-pnnx.Input              input_0     0 1 input
-pnnx.Input              input_1     0 1 weight
-pnnx.Input              input_2     0 1 bias
-pnnx.Input              input_3     0 1 stride
-pnnx.Input              input_4     0 1 padding
-pnnx.Input              input_5     0 1 dilation
-pnnx.Input              input_6     0 1 output_padding
-pnnx.Input              input_7     0 1 groups
-prim::Constant          op_0        0 1 transposed value=True
-prim::Constant          op_1        0 1 benchmark value=*
-prim::Constant          op_2        0 1 deterministic value=*
-prim::Constant          op_3        0 1 cudnn_enabled value=*
-prim::Constant          op_4        0 1 allow_tf32 value=*
-aten::_convolution      op_5        13 1 input weight bias stride padding dilation transposed output_padding groups benchmark deterministic cudnn_enabled allow_tf32 out
-pnnx.Output             output      1 0 out
-)PNNXIR";
-    }
-
-    const char* type_str() const
-    {
-        return "F.conv_transpose2d";
-    }
-};
-
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_conv_transpose2d, 10)
-
-} // namespace pnnx

tools/pnnx/src/pass_level2/F_conv_transpose3d.cpp

@@ -1,52 +0,0 @@
-// Tencent is pleased to support the open source community by making ncnn available.
-//
-// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
-//
-// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
-// in compliance with the License. You may obtain a copy of the License at
-//
-// https://opensource.org/licenses/BSD-3-Clause
-//
-// Unless required by applicable law or agreed to in writing, software distributed
-// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
-// CONDITIONS OF ANY KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations under the License.
-
-#include "pass_level2.h"
-
-namespace pnnx {
-
-class F_conv_transpose3d : public GraphRewriterPass
-{
-public:
-    const char* match_pattern_graph() const
-    {
-        return R"PNNXIR(7767517
-15 14
-pnnx.Input              input_0     0 1 input
-pnnx.Input              input_1     0 1 weight
-pnnx.Input              input_2     0 1 bias
-pnnx.Input              input_3     0 1 stride
-pnnx.Input              input_4     0 1 padding
-pnnx.Input              input_5     0 1 dilation
-pnnx.Input              input_6     0 1 output_padding
-pnnx.Input              input_7     0 1 groups
-prim::Constant          op_0        0 1 transposed value=True
-prim::Constant          op_1        0 1 benchmark value=*
-prim::Constant          op_2        0 1 deterministic value=*
-prim::Constant          op_3        0 1 cudnn_enabled value=*
-prim::Constant          op_4        0 1 allow_tf32 value=*
-aten::_convolution      op_5        13 1 input weight bias stride padding dilation transposed output_padding groups benchmark deterministic cudnn_enabled allow_tf32 out
-pnnx.Output             output      1 0 out
-)PNNXIR";
-    }
-
-    const char* type_str() const
-    {
-        return "F.conv_transpose3d";
-    }
-};
-
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_conv_transpose3d, 10)
-
-} // namespace pnnx

tools/pnnx/src/pass_level3.cpp

@@ -21,6 +21,7 @@
 #include "pass_level3/fuse_chunk_split_unpack.h"
 #include "pass_level3/fuse_expression.h"
 #include "pass_level3/fuse_rnn_unpack.h"
+#include "pass_level3/rename_F_conv_transposend.h"
 #include "pass_level3/rename_F_convmode.h"
 
 // #include "pass_level4/canonicalize.h"
@@ -43,6 +44,8 @@ void pass_level3(Graph& g)
 
     eliminate_tuple_pair(g);
 
+    rename_F_conv_transposend(g);
+
     rename_F_convmode(g);
 
     fuse_expression(g);

tools/pnnx/src/pass_level3/rename_F_conv_transposend.cpp

@@ -0,0 +1,49 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "rename_F_conv_transposend.h"
+#include <algorithm>
+
+namespace pnnx {
+
+void rename_F_conv_transposend(Graph& graph)
+{
+    for (size_t i = 0; i < graph.ops.size(); i++)
+    {
+        Operator* op = graph.ops[i];
+
+        if (op->type != "F.conv_transposend")
+            continue;
+
+        Operator* stride = op->inputs[3]->producer;
+        if (stride->type != "prim::ListConstruct")
+            continue;
+
+        int n = stride->inputs.size();
+        if (n == 1)
+        {
+            op->type = "F.conv_transpose1d";
+        }
+        if (n == 2)
+        {
+            op->type = "F.conv_transpose2d";
+        }
+        if (n == 3)
+        {
+            op->type = "F.conv_transpose3d";
+        }
+    }
+}
+
+} // namespace pnnx

tools/pnnx/src/pass_level3/rename_F_conv_transposend.h

@@ -0,0 +1,21 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "ir.h"
+
+namespace pnnx {
+
+void rename_F_conv_transposend(Graph& graph);
+
+} // namespace pnnx

tools/pnnx/src/pass_ncnn/F_conv2d.cpp

@@ -52,8 +52,8 @@ pnnx.Output             output      1 0 out
         }
 
         op->params["0"] = weight.shape[0];
-        op->params["1"] = weight.shape[2];
-        op->params["11"] = weight.shape[3];
+        op->params["1"] = weight.shape[3];
+        op->params["11"] = weight.shape[2];
         op->params["2"] = captured_params.at("dilation").ai[1];
         op->params["12"] = captured_params.at("dilation").ai[0];
         op->params["3"] = captured_params.at("stride").ai[1];
@@ -119,8 +119,8 @@ pnnx.Output             output      1 0 out
         }
 
         op->params["0"] = weight.shape[0];
-        op->params["1"] = weight.shape[2];
-        op->params["11"] = weight.shape[3];
+        op->params["1"] = weight.shape[3];
+        op->params["11"] = weight.shape[2];
         op->params["2"] = captured_params.at("dilation").ai[1];
         op->params["12"] = captured_params.at("dilation").ai[0];
         op->params["3"] = captured_params.at("stride").ai[1];
@@ -183,8 +183,8 @@ pnnx.Output             output      1 0 out
         }
 
         op->params["0"] = weight.shape[0];
-        op->params["1"] = weight.shape[2];
-        op->params["11"] = weight.shape[3];
+        op->params["1"] = weight.shape[3];
+        op->params["11"] = weight.shape[2];
         op->params["2"] = captured_params.at("dilation").ai[1];
         op->params["12"] = captured_params.at("dilation").ai[0];
         op->params["3"] = captured_params.at("stride").ai[1];
@@ -251,8 +251,8 @@ pnnx.Output             output      1 0 out
         }
 
         op->params["0"] = weight.shape[0];
-        op->params["1"] = weight.shape[2];
-        op->params["11"] = weight.shape[3];
+        op->params["1"] = weight.shape[3];
+        op->params["11"] = weight.shape[2];
         op->params["2"] = captured_params.at("dilation").ai[1];
         op->params["12"] = captured_params.at("dilation").ai[0];
         op->params["3"] = captured_params.at("stride").ai[1];
@@ -315,8 +315,8 @@ pnnx.Output             output      1 0 out
         }
 
         op->params["0"] = weight_shape[0];
-        op->params["1"] = weight_shape[2];
-        op->params["11"] = weight_shape[3];
+        op->params["1"] = weight_shape[3];
+        op->params["11"] = weight_shape[2];
         op->params["2"] = captured_params.at("dilation").ai[1];
         op->params["12"] = captured_params.at("dilation").ai[0];
         op->params["3"] = captured_params.at("stride").ai[1];
@@ -375,8 +375,8 @@ pnnx.Output             output      1 0 out
         }
 
         op->params["0"] = weight_shape[0];
-        op->params["1"] = weight_shape[2];
-        op->params["11"] = weight_shape[3];
+        op->params["1"] = weight_shape[3];
+        op->params["11"] = weight_shape[2];
         op->params["2"] = captured_params.at("dilation").ai[1];
         op->params["12"] = captured_params.at("dilation").ai[0];
         op->params["3"] = captured_params.at("stride").ai[1];
@@ -434,8 +434,8 @@ pnnx.Output             output      1 0 out
         }
 
         op->params["0"] = weight_shape[0];
-        op->params["1"] = weight_shape[2];
-        op->params["11"] = weight_shape[3];
+        op->params["1"] = weight_shape[3];
+        op->params["11"] = weight_shape[2];
         op->params["2"] = captured_params.at("dilation").ai[1];
         op->params["12"] = captured_params.at("dilation").ai[0];
         op->params["3"] = captured_params.at("stride").ai[1];
@@ -495,8 +495,8 @@ pnnx.Output             output      1 0 out
         }
 
         op->params["0"] = weight_shape[0];
-        op->params["1"] = weight_shape[2];
-        op->params["11"] = weight_shape[3];
+        op->params["1"] = weight_shape[3];
+        op->params["11"] = weight_shape[2];
         op->params["2"] = captured_params.at("dilation").ai[1];
         op->params["12"] = captured_params.at("dilation").ai[0];
         op->params["3"] = captured_params.at("stride").ai[1];

tools/pnnx/src/pass_ncnn/F_conv3d.cpp

@@ -0,0 +1,303 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "pass_ncnn.h"
+
+namespace pnnx {
+
+namespace ncnn {
+
+class F_conv3d : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_weight   0 1 weight @qwq
+F.conv3d                op_0        2 1 input weight out bias=None stride=%stride padding=%padding dilation=%dilation groups=1
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "Convolution3D";
+    }
+
+    const char* name_str() const
+    {
+        return "conv3d";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        Attribute weight;
+        for (const auto& x : captured_attrs)
+        {
+            if (x.first.substr(0, 10) == "op_weight.")
+                weight = x.second;
+        }
+
+        op->params["0"] = weight.shape[0];
+        op->params["1"] = weight.shape[4];
+        op->params["11"] = weight.shape[3];
+        op->params["21"] = weight.shape[2];
+        op->params["2"] = captured_params.at("dilation").ai[2];
+        op->params["12"] = captured_params.at("dilation").ai[1];
+        op->params["22"] = captured_params.at("dilation").ai[0];
+        op->params["3"] = captured_params.at("stride").ai[2];
+        op->params["13"] = captured_params.at("stride").ai[1];
+        op->params["23"] = captured_params.at("stride").ai[0];
+        if (captured_params.at("padding").type == 4)
+        {
+            if (captured_params.at("padding").s == "same")
+                op->params["4"] = -233;
+            else if (captured_params.at("padding").s == "valid")
+                op->params["4"] = 0;
+        }
+        else
+        {
+            op->params["4"] = captured_params.at("padding").ai[2];
+            op->params["14"] = captured_params.at("padding").ai[1];
+            op->params["24"] = captured_params.at("padding").ai[0];
+        }
+        op->params["5"] = 0;
+        op->params["6"] = (int)(weight.data.size() / sizeof(float));
+
+        op->attrs["0"] = Attribute();
+        op->attrs["0"].data = {0, 0, 0, 0};
+        op->attrs["1"] = weight;
+    }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(F_conv3d, 20)
+
+class F_conv3d_1 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_weight   0 1 weight @qwq
+pnnx.Attribute          op_bias     0 1 bias @qwq
+F.conv3d                op_0        3 1 input weight bias out stride=%stride padding=%padding dilation=%dilation groups=1
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "Convolution3D";
+    }
+
+    const char* name_str() const
+    {
+        return "conv3d";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        Attribute weight;
+        Attribute bias;
+        for (const auto& x : captured_attrs)
+        {
+            if (x.first.substr(0, 10) == "op_weight.")
+                weight = x.second;
+            if (x.first.substr(0, 8) == "op_bias.")
+                bias = x.second;
+        }
+
+        op->params["0"] = weight.shape[0];
+        op->params["1"] = weight.shape[4];
+        op->params["11"] = weight.shape[3];
+        op->params["21"] = weight.shape[2];
+        op->params["2"] = captured_params.at("dilation").ai[2];
+        op->params["12"] = captured_params.at("dilation").ai[1];
+        op->params["22"] = captured_params.at("dilation").ai[0];
+        op->params["3"] = captured_params.at("stride").ai[2];
+        op->params["13"] = captured_params.at("stride").ai[1];
+        op->params["23"] = captured_params.at("stride").ai[0];
+        if (captured_params.at("padding").type == 4)
+        {
+            if (captured_params.at("padding").s == "same")
+                op->params["4"] = -233;
+            else if (captured_params.at("padding").s == "valid")
+                op->params["4"] = 0;
+        }
+        else
+        {
+            op->params["4"] = captured_params.at("padding").ai[2];
+            op->params["14"] = captured_params.at("padding").ai[1];
+            op->params["24"] = captured_params.at("padding").ai[0];
+        }
+        op->params["5"] = 1;
+        op->params["6"] = (int)(weight.data.size() / sizeof(float));
+
+        op->attrs["0"] = Attribute();
+        op->attrs["0"].data = {0, 0, 0, 0};
+        op->attrs["1"] = weight;
+        op->attrs["2"] = bias;
+    }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(F_conv3d_1, 20)
+
+class F_conv3d_2 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_weight   0 1 weight @qwq
+F.conv3d                op_0        2 1 input weight out bias=None stride=%stride padding=%padding dilation=%dilation groups=%groups
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "ConvolutionDepthWise3D";
+    }
+
+    const char* name_str() const
+    {
+        return "convdw3d";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        Attribute weight;
+        for (const auto& x : captured_attrs)
+        {
+            if (x.first.substr(0, 10) == "op_weight.")
+                weight = x.second;
+        }
+
+        op->params["0"] = weight.shape[0];
+        op->params["1"] = weight.shape[4];
+        op->params["11"] = weight.shape[3];
+        op->params["21"] = weight.shape[2];
+        op->params["2"] = captured_params.at("dilation").ai[2];
+        op->params["12"] = captured_params.at("dilation").ai[1];
+        op->params["22"] = captured_params.at("dilation").ai[0];
+        op->params["3"] = captured_params.at("stride").ai[2];
+        op->params["13"] = captured_params.at("stride").ai[1];
+        op->params["23"] = captured_params.at("stride").ai[0];
+        if (captured_params.at("padding").type == 4)
+        {
+            if (captured_params.at("padding").s == "same")
+                op->params["4"] = -233;
+            else if (captured_params.at("padding").s == "valid")
+                op->params["4"] = 0;
+        }
+        else
+        {
+            op->params["4"] = captured_params.at("padding").ai[2];
+            op->params["14"] = captured_params.at("padding").ai[1];
+            op->params["24"] = captured_params.at("padding").ai[0];
+        }
+        op->params["5"] = 0;
+        op->params["6"] = (int)(weight.data.size() / sizeof(float));
+        op->params["7"] = captured_params.at("groups");
+
+        op->attrs["0"] = Attribute();
+        op->attrs["0"].data = {0, 0, 0, 0};
+        op->attrs["1"] = weight;
+    }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(F_conv3d_2, 21)
+
+class F_conv3d_3 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_weight   0 1 weight @qwq
+pnnx.Attribute          op_bias     0 1 bias @qwq
+F.conv3d                op_0        3 1 input weight bias out stride=%stride padding=%padding dilation=%dilation groups=%groups
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "ConvolutionDepthWise3D";
+    }
+
+    const char* name_str() const
+    {
+        return "convdw3d";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        Attribute weight;
+        Attribute bias;
+        for (const auto& x : captured_attrs)
+        {
+            if (x.first.substr(0, 10) == "op_weight.")
+                weight = x.second;
+            if (x.first.substr(0, 8) == "op_bias.")
+                bias = x.second;
+        }
+
+        op->params["0"] = weight.shape[0];
+        op->params["1"] = weight.shape[4];
+        op->params["11"] = weight.shape[3];
+        op->params["21"] = weight.shape[2];
+        op->params["2"] = captured_params.at("dilation").ai[2];
+        op->params["12"] = captured_params.at("dilation").ai[1];
+        op->params["22"] = captured_params.at("dilation").ai[0];
+        op->params["3"] = captured_params.at("stride").ai[2];
+        op->params["13"] = captured_params.at("stride").ai[1];
+        op->params["23"] = captured_params.at("stride").ai[0];
+        if (captured_params.at("padding").type == 4)
+        {
+            if (captured_params.at("padding").s == "same")
+                op->params["4"] = -233;
+            else if (captured_params.at("padding").s == "valid")
+                op->params["4"] = 0;
+        }
+        else
+        {
+            op->params["4"] = captured_params.at("padding").ai[2];
+            op->params["14"] = captured_params.at("padding").ai[1];
+            op->params["24"] = captured_params.at("padding").ai[0];
+        }
+        op->params["5"] = 1;
+        op->params["6"] = (int)(weight.data.size() / sizeof(float));
+        op->params["7"] = captured_params.at("groups");
+
+        op->attrs["0"] = Attribute();
+        op->attrs["0"].data = {0, 0, 0, 0};
+        op->attrs["1"] = weight;
+        op->attrs["2"] = bias;
+    }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(F_conv3d_3, 21)
+
+} // namespace ncnn
+
+} // namespace pnnx

tools/pnnx/src/pass_ncnn/F_conv_transpose2d.cpp

@@ -0,0 +1,377 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "pass_ncnn.h"
+
+namespace pnnx {
+
+namespace ncnn {
+
+class F_conv_transpose2d : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_weight   0 1 weight @qwq
+F.conv_transpose2d      op_0        2 1 input weight out bias=None stride=%stride padding=%padding dilation=%dilation output_padding=%output_padding groups=1
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "Deconvolution";
+    }
+
+    const char* name_str() const
+    {
+        return "conv_transpose2d";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        Attribute weight;
+        for (const auto& x : captured_attrs)
+        {
+            if (x.first.substr(0, 10) == "op_weight.")
+                weight = x.second;
+        }
+
+        op->params["0"] = weight.shape[1];
+        op->params["1"] = weight.shape[3];
+        op->params["11"] = weight.shape[2];
+        op->params["2"] = captured_params.at("dilation").ai[1];
+        op->params["12"] = captured_params.at("dilation").ai[0];
+        op->params["3"] = captured_params.at("stride").ai[1];
+        op->params["13"] = captured_params.at("stride").ai[0];
+        op->params["4"] = captured_params.at("padding").ai[1];
+        op->params["14"] = captured_params.at("padding").ai[0];
+        op->params["18"] = captured_params.at("output_padding").ai[1];
+        op->params["19"] = captured_params.at("output_padding").ai[0];
+        op->params["5"] = 0;
+        op->params["6"] = (int)(weight.data.size() / sizeof(float));
+
+        // transpose inch-outch-kh-kw to outch-inch-kh-kw
+        const int inch = weight.shape[0];
+        const int outch = weight.shape[1];
+        const int kh = weight.shape[2];
+        const int kw = weight.shape[3];
+        std::vector<float> new_weight;
+        {
+            const float* w = (const float*)weight.data.data();
+
+            new_weight.resize(outch * inch * kh * kw);
+            float* w2 = (float*)new_weight.data();
+            const int maxk = kh * kw;
+
+            // reorder weight from inch-outch to outch-inch
+            for (int i = 0; i < outch; i++)
+            {
+                for (int j = 0; j < inch; j++)
+                {
+                    for (int k = 0; k < maxk; k++)
+                    {
+                        w2[(i * inch + j) * maxk + k] = w[(j * outch + i) * maxk + k];
+                    }
+                }
+            }
+        }
+
+        op->attrs["0"] = Attribute();
+        op->attrs["0"].data = {0, 0, 0, 0};
+        op->attrs["1"] = Attribute({outch, inch, kh, kw}, new_weight);
+    }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(F_conv_transpose2d, 20)
+
+class F_conv_transpose2d_1 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_weight   0 1 weight @qwq
+pnnx.Attribute          op_bias     0 1 bias @qwq
+F.conv_transpose2d      op_0        3 1 input weight bias out stride=%stride padding=%padding dilation=%dilation output_padding=%output_padding groups=1
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "Deconvolution";
+    }
+
+    const char* name_str() const
+    {
+        return "conv_transpose2d";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        Attribute weight;
+        Attribute bias;
+        for (const auto& x : captured_attrs)
+        {
+            if (x.first.substr(0, 10) == "op_weight.")
+                weight = x.second;
+            if (x.first.substr(0, 8) == "op_bias.")
+                bias = x.second;
+        }
+
+        op->params["0"] = weight.shape[1];
+        op->params["1"] = weight.shape[3];
+        op->params["11"] = weight.shape[2];
+        op->params["2"] = captured_params.at("dilation").ai[1];
+        op->params["12"] = captured_params.at("dilation").ai[0];
+        op->params["3"] = captured_params.at("stride").ai[1];
+        op->params["13"] = captured_params.at("stride").ai[0];
+        op->params["4"] = captured_params.at("padding").ai[1];
+        op->params["14"] = captured_params.at("padding").ai[0];
+        op->params["18"] = captured_params.at("output_padding").ai[1];
+        op->params["19"] = captured_params.at("output_padding").ai[0];
+        op->params["5"] = 1;
+        op->params["6"] = (int)(weight.data.size() / sizeof(float));
+
+        // transpose inch-outch-kh-kw to outch-inch-kh-kw
+        const int inch = weight.shape[0];
+        const int outch = weight.shape[1];
+        const int kh = weight.shape[2];
+        const int kw = weight.shape[3];
+        std::vector<float> new_weight;
+        {
+            const float* w = (const float*)weight.data.data();
+
+            new_weight.resize(outch * inch * kh * kw);
+            float* w2 = (float*)new_weight.data();
+            const int maxk = kh * kw;
+
+            // reorder weight from inch-outch to outch-inch
+            for (int i = 0; i < outch; i++)
+            {
+                for (int j = 0; j < inch; j++)
+                {
+                    for (int k = 0; k < maxk; k++)
+                    {
+                        w2[(i * inch + j) * maxk + k] = w[(j * outch + i) * maxk + k];
+                    }
+                }
+            }
+        }
+
+        op->attrs["0"] = Attribute();
+        op->attrs["0"].data = {0, 0, 0, 0};
+        op->attrs["1"] = Attribute({outch, inch, kh, kw}, new_weight);
+        op->attrs["2"] = bias;
+    }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(F_conv_transpose2d_1, 20)
+
+class F_conv_transpose2d_2 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_weight   0 1 weight @qwq
+F.conv_transpose2d      op_0        2 1 input weight out bias=None stride=%stride padding=%padding dilation=%dilation output_padding=%output_padding groups=%groups
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "DeconvolutionDepthWise";
+    }
+
+    const char* name_str() const
+    {
+        return "deconvdw2d";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        Attribute weight;
+        for (const auto& x : captured_attrs)
+        {
+            if (x.first.substr(0, 10) == "op_weight.")
+                weight = x.second;
+        }
+
+        const int groups = captured_params.at("groups").i;
+
+        op->params["0"] = weight.shape[1] * groups;
+        op->params["1"] = weight.shape[3];
+        op->params["11"] = weight.shape[2];
+        op->params["2"] = captured_params.at("dilation").ai[1];
+        op->params["12"] = captured_params.at("dilation").ai[0];
+        op->params["3"] = captured_params.at("stride").ai[1];
+        op->params["13"] = captured_params.at("stride").ai[0];
+        op->params["4"] = captured_params.at("padding").ai[1];
+        op->params["14"] = captured_params.at("padding").ai[0];
+        op->params["18"] = captured_params.at("output_padding").ai[1];
+        op->params["19"] = captured_params.at("output_padding").ai[0];
+        op->params["5"] = 0;
+        op->params["6"] = (int)(weight.data.size() / sizeof(float));
+        op->params["7"] = groups;
+
+        // transpose group-inch/group-outch/group-kh-kw to group-outch/group-inch/group-kh-kw
+        const int inch = weight.shape[0];
+        const int outch = weight.shape[1] * groups;
+        const int kh = weight.shape[2];
+        const int kw = weight.shape[3];
+        std::vector<float> new_weight;
+        {
+            const float* w = (const float*)weight.data.data();
+
+            new_weight.resize(outch / groups * inch * kh * kw);
+            float* w2 = (float*)new_weight.data();
+            const int outch_g = outch / groups;
+            const int inch_g = inch / groups;
+            const int maxk = kh * kw;
+
+            for (int g = 0; g < groups; g++)
+            {
+                // reorder weight from inch-outch to outch-inch
+                float* wg2 = w2 + g * outch_g * inch_g * maxk;
+                const float* wg = w + g * inch_g * outch_g * maxk;
+                for (int i = 0; i < outch_g; i++)
+                {
+                    for (int j = 0; j < inch_g; j++)
+                    {
+                        for (int k = 0; k < maxk; k++)
+                        {
+                            wg2[(i * inch_g + j) * maxk + k] = wg[(j * outch_g + i) * maxk + k];
+                        }
+                    }
+                }
+            }
+        }
+
+        op->attrs["0"] = Attribute();
+        op->attrs["0"].data = {0, 0, 0, 0};
+        op->attrs["1"] = Attribute({outch / groups, inch, kh, kw}, new_weight);
+    }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(F_conv_transpose2d_2, 21)
+
+class F_conv_transpose2d_3 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_weight   0 1 weight @qwq
+pnnx.Attribute          op_bias     0 1 bias @qwq
+F.conv_transpose2d      op_0        3 1 input weight bias out stride=%stride padding=%padding dilation=%dilation output_padding=%output_padding groups=%groups
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "DeconvolutionDepthWise";
+    }
+
+    const char* name_str() const
+    {
+        return "deconvdw2d";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        Attribute weight;
+        Attribute bias;
+        for (const auto& x : captured_attrs)
+        {
+            if (x.first.substr(0, 10) == "op_weight.")
+                weight = x.second;
+            if (x.first.substr(0, 8) == "op_bias.")
+                bias = x.second;
+        }
+
+        const int groups = captured_params.at("groups").i;
+
+        op->params["0"] = weight.shape[1] * groups;
+        op->params["1"] = weight.shape[3];
+        op->params["11"] = weight.shape[2];
+        op->params["2"] = captured_params.at("dilation").ai[1];
+        op->params["12"] = captured_params.at("dilation").ai[0];
+        op->params["3"] = captured_params.at("stride").ai[1];
+        op->params["13"] = captured_params.at("stride").ai[0];
+        op->params["4"] = captured_params.at("padding").ai[1];
+        op->params["14"] = captured_params.at("padding").ai[0];
+        op->params["18"] = captured_params.at("output_padding").ai[1];
+        op->params["19"] = captured_params.at("output_padding").ai[0];
+        op->params["5"] = 1;
+        op->params["6"] = (int)(weight.data.size() / sizeof(float));
+        op->params["7"] = groups;
+
+        // transpose group-inch/group-outch/group-kh-kw to group-outch/group-inch/group-kh-kw
+        const int inch = weight.shape[0];
+        const int outch = weight.shape[1] * groups;
+        const int kh = weight.shape[2];
+        const int kw = weight.shape[3];
+        std::vector<float> new_weight;
+        {
+            const float* w = (const float*)weight.data.data();
+
+            new_weight.resize(outch / groups * inch * kh * kw);
+            float* w2 = (float*)new_weight.data();
+            const int outch_g = outch / groups;
+            const int inch_g = inch / groups;
+            const int maxk = kh * kw;
+
+            for (int g = 0; g < groups; g++)
+            {
+                // reorder weight from inch-outch to outch-inch
+                float* wg2 = w2 + g * outch_g * inch_g * maxk;
+                const float* wg = w + g * inch_g * outch_g * maxk;
+                for (int i = 0; i < outch_g; i++)
+                {
+                    for (int j = 0; j < inch_g; j++)
+                    {
+                        for (int k = 0; k < maxk; k++)
+                        {
+                            wg2[(i * inch_g + j) * maxk + k] = wg[(j * outch_g + i) * maxk + k];
+                        }
+                    }
+                }
+            }
+        }
+
+        op->attrs["0"] = Attribute();
+        op->attrs["0"].data = {0, 0, 0, 0};
+        op->attrs["1"] = Attribute({outch / groups, inch, kh, kw}, new_weight);
+        op->attrs["2"] = bias;
+    }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(F_conv_transpose2d_3, 21)
+
+} // namespace ncnn
+
+} // namespace pnnx

tools/pnnx/tests/ncnn/CMakeLists.txt

@@ -14,14 +14,21 @@ pnnx_ncnn_add_test(F_adaptive_max_pool3d)
 pnnx_ncnn_add_test(F_avg_pool1d)
 pnnx_ncnn_add_test(F_avg_pool2d)
 pnnx_ncnn_add_test(F_avg_pool3d)
+pnnx_ncnn_add_test(F_batch_norm)
+#pnnx_ncnn_add_test(F_conv_transpose1d) # TODO
+pnnx_ncnn_add_test(F_conv_transpose2d)
+#pnnx_ncnn_add_test(F_conv_transpose3d) # TODO
 pnnx_ncnn_add_test(F_conv1d)
 pnnx_ncnn_add_test(F_conv2d)
+pnnx_ncnn_add_test(F_conv3d)
 pnnx_ncnn_add_test(F_elu)
 pnnx_ncnn_add_test(F_gelu)
+pnnx_ncnn_add_test(F_group_norm)
 pnnx_ncnn_add_test(F_hardsigmoid)
 pnnx_ncnn_add_test(F_hardswish)
 pnnx_ncnn_add_test(F_hardtanh)
 pnnx_ncnn_add_test(F_interpolate)
+pnnx_ncnn_add_test(F_layer_norm)
 pnnx_ncnn_add_test(F_leaky_relu)
 pnnx_ncnn_add_test(F_local_response_norm)
 pnnx_ncnn_add_test(F_max_pool1d)
@@ -31,6 +38,7 @@ pnnx_ncnn_add_test(F_normalize)
 pnnx_ncnn_add_test(F_pad)
 pnnx_ncnn_add_test(F_pixel_shuffle)
 pnnx_ncnn_add_test(F_pixel_unshuffle)
+pnnx_ncnn_add_test(F_prelu)
 pnnx_ncnn_add_test(F_relu)
 pnnx_ncnn_add_test(F_relu6)
 pnnx_ncnn_add_test(F_sigmoid)

tools/pnnx/tests/ncnn/test_F_batch_norm.py

@@ -0,0 +1,76 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.m3 = torch.rand(16)
+        self.v3 = torch.rand(16)
+        self.w3 = nn.Parameter(torch.rand(16))
+        self.b3 = nn.Parameter(torch.rand(16))
+        self.m4 = torch.rand(2)
+        self.v4 = torch.rand(2)
+        self.w4 = nn.Parameter(torch.rand(2))
+        self.b4 = nn.Parameter(torch.rand(2))
+        self.m5 = torch.rand(3)
+        self.v5 = torch.rand(3)
+        self.w5 = nn.Parameter(torch.rand(3))
+        self.b5 = nn.Parameter(torch.rand(3))
+
+    def forward(self, x, y, z):
+        x = F.batch_norm(x, self.m3, self.v3, self.w3, self.b3)
+
+        y = F.batch_norm(y, self.m4, self.v4, self.w4, self.b4, eps=1e-3)
+
+        z = F.batch_norm(z, self.m5, self.v5, self.w5, self.b5, eps=1e-2)
+        return x, y, z
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+    y = torch.rand(1, 2, 16)
+    z = torch.rand(1, 3, 12, 16)
+
+    a = net(x, y, z)
+
+    # export torchscript
+    mod = torch.jit.trace(net, (x, y, z))
+    mod.save("test_F_batch_norm.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../../src/pnnx test_F_batch_norm.pt inputshape=[1,16],[1,2,16],[1,3,12,16]")
+
+    # ncnn inference
+    import test_F_batch_norm_ncnn
+    b = test_F_batch_norm_ncnn.test_inference()
+
+    for a0, b0 in zip(a, b):
+        if not torch.allclose(a0, b0, 1e-4, 1e-4):
+            return False
+    return True
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/ncnn/test_F_conv1d.py

@@ -20,13 +20,20 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
-    def forward(self, x, w0, w1, b1):
+        self.w2 = nn.Parameter(torch.rand(12, 6, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(6, 4, 3))
+
+    def forward(self, x, w0, w1, b1, y):
         x = F.conv1d(x, w0, None, stride=2, padding=1)
         if torch.__version__ < '1.9':
             x = F.conv1d(x, w1, b1, stride=1, padding=1, dilation=2, groups=2)
         else:
             x = F.conv1d(x, w1, b1, stride=1, padding='same', dilation=2, groups=2)
-        return x
+
+        y = F.conv1d(y, self.w2, self.b2, stride=2, padding=2)
+        y = F.conv1d(y, self.w3, None, stride=2, padding=1, groups=3)
+        return x, y
 
 def test():
     net = Model()
@@ -37,22 +44,23 @@ def test():
     w0 = torch.rand(16, 12, 3)
     w1 = torch.rand(16, 8, 5)
     b1 = torch.rand(16)
+    y = torch.rand(1, 6, 25)
 
-    a = net(x, w0, w1, b1)
+    a0, a1 = net(x, w0, w1, b1, y)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, w0, w1, b1))
+    mod = torch.jit.trace(net, (x, w0, w1, b1, y))
     mod.save("test_F_conv1d.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../../src/pnnx test_F_conv1d.pt inputshape=[1,12,52],[16,12,3],[16,8,5],[16]")
+    os.system("../../src/pnnx test_F_conv1d.pt inputshape=[1,12,52],[16,12,3],[16,8,5],[16],[1,6,25]")
 
     # ncnn inference
     import test_F_conv1d_ncnn
-    b = test_F_conv1d_ncnn.test_inference()
+    b0, b1 = test_F_conv1d_ncnn.test_inference()
 
-    return torch.allclose(a, b, 1e-4, 1e-4)
+    return torch.allclose(a0, b0, 1e-4, 1e-4) and torch.allclose(a1, b1, 1e-4, 1e-4)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/ncnn/test_F_conv2d.py

@@ -20,13 +20,20 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
-    def forward(self, x, w0, w1, b1):
+        self.w2 = nn.Parameter(torch.rand(12, 6, 4, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(6, 4, 3, 3))
+
+    def forward(self, x, w0, w1, b1, y):
         x = F.conv2d(x, w0, None, stride=(2,2), padding=(1,1))
         if torch.__version__ < '1.9':
             x = F.conv2d(x, w1, b1, stride=(1,1), padding=(1,1), dilation=(2,1), groups=2)
         else:
             x = F.conv2d(x, w1, b1, stride=(1,1), padding='same', dilation=(2,1), groups=2)
-        return x
+
+        y = F.conv2d(y, self.w2, self.b2, stride=(2,2), padding=(2,2))
+        y = F.conv2d(y, self.w3, None, stride=(2,2), padding=(1,1), groups=3)
+        return x, y
 
 def test():
     net = Model()
@@ -37,22 +44,23 @@ def test():
     w0 = torch.rand(16, 12, 3, 3)
     w1 = torch.rand(16, 8, 5, 5)
     b1 = torch.rand(16)
+    y = torch.rand(1, 6, 32, 25)
 
-    a = net(x, w0, w1, b1)
+    a0, a1 = net(x, w0, w1, b1, y)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, w0, w1, b1))
+    mod = torch.jit.trace(net, (x, w0, w1, b1, y))
     mod.save("test_F_conv2d.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../../src/pnnx test_F_conv2d.pt inputshape=[1,12,52,64],[16,12,3,3],[16,8,5,5],[16]")
+    os.system("../../src/pnnx test_F_conv2d.pt inputshape=[1,12,52,64],[16,12,3,3],[16,8,5,5],[16],[1,6,32,25]")
 
     # ncnn inference
     import test_F_conv2d_ncnn
-    b = test_F_conv2d_ncnn.test_inference()
+    b0, b1 = test_F_conv2d_ncnn.test_inference()
 
-    return torch.allclose(a, b, 1e-4, 1e-4)
+    return torch.allclose(a0, b0, 1e-4, 1e-4) and torch.allclose(a1, b1, 1e-4, 1e-4)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/ncnn/test_F_conv3d.py

@@ -0,0 +1,59 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.w2 = nn.Parameter(torch.rand(12, 6, 4, 4, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(6, 4, 3, 3, 3))
+
+    def forward(self, y):
+        y = F.conv3d(y, self.w2, self.b2, stride=(2,2,2), padding=(2,2,2))
+        y = F.conv3d(y, self.w3, None, stride=(2,2,2), padding=(1,1,1), groups=3)
+        return y
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    y = torch.rand(1, 6, 12, 11, 10)
+
+    a = net(y)
+
+    # export torchscript
+    mod = torch.jit.trace(net, y)
+    mod.save("test_F_conv3d.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../../src/pnnx test_F_conv3d.pt inputshape=[1,6,12,11,10]")
+
+    # ncnn inference
+    import test_F_conv3d_ncnn
+    b = test_F_conv3d_ncnn.test_inference()
+
+    return torch.allclose(a, b, 1e-4, 1e-4)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/ncnn/test_F_conv_transpose1d.py

@@ -0,0 +1,59 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.w2 = nn.Parameter(torch.rand(6, 12, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(12, 2, 3))
+
+    def forward(self, y):
+        y = F.conv_transpose1d(y, self.w2, self.b2, stride=2, padding=1, output_padding=1)
+        y = F.conv_transpose1d(y, self.w3, None, stride=1, padding=2, dilation=2, groups=3)
+        return y
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    y = torch.rand(1, 6, 5)
+
+    a = net(y)
+
+    # export torchscript
+    mod = torch.jit.trace(net, y)
+    mod.save("test_F_conv_transpose1d.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../../src/pnnx test_F_conv_transpose1d.pt inputshape=[1,6,5]")
+
+    # ncnn inference
+    import test_F_conv_transpose1d_ncnn
+    b = test_F_conv_transpose1d_ncnn.test_inference()
+
+    return torch.allclose(a, b, 1e-4, 1e-4)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/ncnn/test_F_conv_transpose2d.py

@@ -0,0 +1,59 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.w2 = nn.Parameter(torch.rand(6, 12, 4, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(12, 2, 3, 3))
+
+    def forward(self, y):
+        y = F.conv_transpose2d(y, self.w2, self.b2, stride=(2,2), padding=(1,1), output_padding=(1,1))
+        y = F.conv_transpose2d(y, self.w3, None, stride=(1,2), padding=(2,1), dilation=(2,1), groups=3)
+        return y
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    y = torch.rand(1, 6, 5, 6)
+
+    a = net(y)
+
+    # export torchscript
+    mod = torch.jit.trace(net, y)
+    mod.save("test_F_conv_transpose2d.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../../src/pnnx test_F_conv_transpose2d.pt inputshape=[1,6,5,6]")
+
+    # ncnn inference
+    import test_F_conv_transpose2d_ncnn
+    b = test_F_conv_transpose2d_ncnn.test_inference()
+
+    return torch.allclose(a, b, 1e-4, 1e-4)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/ncnn/test_F_conv_transpose3d.py

@@ -0,0 +1,59 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.w2 = nn.Parameter(torch.rand(6, 12, 4, 4, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(12, 2, 3, 3, 3))
+
+    def forward(self, y):
+        y = F.conv_transpose3d(y, self.w2, self.b2, stride=(2,2, 2), padding=(1,0, 1), output_padding=(1,1, 0))
+        y = F.conv_transpose3d(y, self.w3, None, stride=(1,1,2), padding=(2,2,1), dilation=(2,2,1), groups=3)
+        return y
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    y = torch.rand(1, 6, 4, 5, 6)
+
+    a = net(y)
+
+    # export torchscript
+    mod = torch.jit.trace(net, y)
+    mod.save("test_F_conv_transpose3d.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../../src/pnnx test_F_conv_transpose3d.pt inputshape=[1,6,4,5,6]")
+
+    # ncnn inference
+    import test_F_conv_transpose3d_ncnn
+    b = test_F_conv_transpose3d_ncnn.test_inference()
+
+    return torch.allclose(a, b, 1e-4, 1e-4)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/ncnn/test_F_group_norm.py

@@ -0,0 +1,60 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.w5 = nn.Parameter(torch.rand(32))
+        self.b5 = nn.Parameter(torch.rand(32))
+
+    def forward(self, z):
+        z = F.group_norm(z, 8, self.w5, self.b5, eps=1e-2)
+        return z
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    z = torch.rand(1, 32, 12, 16)
+
+    a = net(z)
+
+    # export torchscript
+    mod = torch.jit.trace(net, z)
+    mod.save("test_F_group_norm.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../../src/pnnx test_F_group_norm.pt inputshape=[1,32,12,16]")
+
+    # ncnn inference
+    import test_F_group_norm_ncnn
+    b = test_F_group_norm_ncnn.test_inference()
+
+    for a0, b0 in zip(a, b):
+        if not torch.allclose(a0, b0, 1e-4, 1e-4):
+            return False
+    return True
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/ncnn/test_F_layer_norm.py

@@ -0,0 +1,65 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.w3 = nn.Parameter(torch.rand(24))
+        self.b3 = nn.Parameter(torch.rand(24))
+        self.w4 = nn.Parameter(torch.rand(12, 16))
+        self.b4 = nn.Parameter(torch.rand(12, 16))
+
+    def forward(self, x, y):
+        x = F.layer_norm(x, (24,), self.w3, self.b3)
+
+        y = F.layer_norm(y, (12,16), self.w4, self.b4, eps=1e-3)
+        return x, y
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(12, 24)
+    y = torch.rand(3, 12, 16)
+
+    a = net(x, y)
+
+    # export torchscript
+    mod = torch.jit.trace(net, (x, y))
+    mod.save("test_F_layer_norm.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../../src/pnnx test_F_layer_norm.pt inputshape=[12,24],[3,12,16]")
+
+    # ncnn inference
+    import test_F_layer_norm_ncnn
+    b = test_F_layer_norm_ncnn.test_inference()
+
+    for a0, b0 in zip(a, b):
+        if not torch.allclose(a0, b0, 1e-4, 1e-4):
+            return False
+    return True
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/ncnn/test_F_prelu.py

@@ -0,0 +1,65 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.w4 = nn.Parameter(torch.rand(16))
+        self.w5 = nn.Parameter(torch.rand(2))
+        self.w6 = nn.Parameter(torch.rand(3))
+
+    def forward(self, x, y, z):
+        x = F.prelu(x, self.w4)
+        y = F.prelu(y, self.w5)
+        z = F.prelu(z, self.w6)
+        return x, y, z
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 16)
+    y = torch.rand(12, 2, 16)
+    z = torch.rand(1, 3, 12, 16)
+
+    a = net(x, y, z)
+
+    # export torchscript
+    mod = torch.jit.trace(net, (x, y, z))
+    mod.save("test_F_prelu.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../../src/pnnx test_F_prelu.pt inputshape=[1,16],[12,2,16],[1,3,12,16]")
+
+    # ncnn inference
+    import test_F_prelu_ncnn
+    b = test_F_prelu_ncnn.test_inference()
+
+    for a0, b0 in zip(a, b):
+        if not torch.allclose(a0, b0, 1e-4, 1e-4):
+            return False
+    return True
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/test_F_batch_norm.py

@@ -20,15 +20,31 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
+        self.m3 = torch.rand(16)
+        self.v3 = torch.rand(16)
+        self.w3 = nn.Parameter(torch.rand(16))
+        self.b3 = nn.Parameter(torch.rand(16))
+        self.m4 = torch.rand(2)
+        self.v4 = torch.rand(2)
+        self.w4 = nn.Parameter(torch.rand(2))
+        self.b4 = nn.Parameter(torch.rand(2))
+        self.m5 = torch.rand(3)
+        self.v5 = torch.rand(3)
+        self.w5 = nn.Parameter(torch.rand(3))
+        self.b5 = nn.Parameter(torch.rand(3))
+
     def forward(self, x, y, z, m0, v0, w0, b0, m1, v1, w1, b1, m2, v2, w2, b2):
         x = F.batch_norm(x, m0, v0, w0, b0)
         x = F.batch_norm(x, m0, v0, None, None)
+        x = F.batch_norm(x, self.m3, self.v3, self.w3, self.b3)
 
         y = F.batch_norm(y, m1, v1, w1, b1, eps=1e-3)
         y = F.batch_norm(y, m1, v1, None, None)
+        y = F.batch_norm(y, self.m4, self.v4, self.w4, self.b4)
 
         z = F.batch_norm(z, m2, v2, w2, b2)
         z = F.batch_norm(z, m2, v2, None, None, eps=1e-2)
+        z = F.batch_norm(z, self.m5, self.v5, self.w5, self.b5)
         return x, y, z
 
 def test():

tools/pnnx/tests/test_F_conv1d.py

@@ -20,13 +20,20 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
-    def forward(self, x, w0, w1, b1):
+        self.w2 = nn.Parameter(torch.rand(12, 6, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(6, 4, 3))
+
+    def forward(self, x, w0, w1, b1, y):
         x = F.conv1d(x, w0, None, stride=2, padding=1)
         if torch.__version__ < '1.9':
             x = F.conv1d(x, w1, b1, stride=1, padding=1, dilation=2, groups=2)
         else:
             x = F.conv1d(x, w1, b1, stride=1, padding='same', dilation=2, groups=2)
-        return x
+
+        y = F.conv1d(y, self.w2, self.b2, stride=2, padding=2)
+        y = F.conv1d(y, self.w3, None, stride=2, padding=1, groups=3)
+        return x, y
 
 def test():
     net = Model()
@@ -37,22 +44,23 @@ def test():
     w0 = torch.rand(16, 12, 3)
     w1 = torch.rand(16, 8, 5)
     b1 = torch.rand(16)
+    y = torch.rand(1, 6, 25)
 
-    a = net(x, w0, w1, b1)
+    a0, a1 = net(x, w0, w1, b1, y)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, w0, w1, b1))
+    mod = torch.jit.trace(net, (x, w0, w1, b1, y))
     mod.save("test_F_conv1d.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../src/pnnx test_F_conv1d.pt inputshape=[1,12,52],[16,12,3],[16,8,5],[16]")
+    os.system("../src/pnnx test_F_conv1d.pt inputshape=[1,12,52],[16,12,3],[16,8,5],[16],[1,6,25]")
 
     # pnnx inference
     import test_F_conv1d_pnnx
-    b = test_F_conv1d_pnnx.test_inference()
+    b0, b1 = test_F_conv1d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.equal(a0, b0) and torch.equal(a1, b1)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_F_conv2d.py

@@ -20,13 +20,20 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
-    def forward(self, x, w0, w1, b1):
+        self.w2 = nn.Parameter(torch.rand(12, 6, 4, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(6, 4, 3, 3))
+
+    def forward(self, x, w0, w1, b1, y):
         x = F.conv2d(x, w0, None, stride=(2,2), padding=(1,1))
         if torch.__version__ < '1.9':
             x = F.conv2d(x, w1, b1, stride=(1,1), padding=(1,1), dilation=(2,1), groups=2)
         else:
             x = F.conv2d(x, w1, b1, stride=(1,1), padding='same', dilation=(2,1), groups=2)
-        return x
+
+        y = F.conv2d(y, self.w2, self.b2, stride=(2,2), padding=(2,2))
+        y = F.conv2d(y, self.w3, None, stride=(2,2), padding=(1,1), groups=3)
+        return x, y
 
 def test():
     net = Model()
@@ -37,22 +44,23 @@ def test():
     w0 = torch.rand(16, 12, 3, 3)
     w1 = torch.rand(16, 8, 5, 5)
     b1 = torch.rand(16)
+    y = torch.rand(1, 6, 32, 25)
 
-    a = net(x, w0, w1, b1)
+    a0, a1 = net(x, w0, w1, b1, y)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, w0, w1, b1))
+    mod = torch.jit.trace(net, (x, w0, w1, b1, y))
     mod.save("test_F_conv2d.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../src/pnnx test_F_conv2d.pt inputshape=[1,12,52,64],[16,12,3,3],[16,8,5,5],[16]")
+    os.system("../src/pnnx test_F_conv2d.pt inputshape=[1,12,52,64],[16,12,3,3],[16,8,5,5],[16],[1,6,32,25]")
 
     # pnnx inference
     import test_F_conv2d_pnnx
-    b = test_F_conv2d_pnnx.test_inference()
+    b0, b1 = test_F_conv2d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.equal(a0, b0) and torch.equal(a1, b1)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_F_conv3d.py

@@ -20,13 +20,20 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
-    def forward(self, x, w0, w1, b1):
+        self.w2 = nn.Parameter(torch.rand(12, 6, 4, 4, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(6, 4, 3, 3, 3))
+
+    def forward(self, x, w0, w1, b1, y):
         x = F.conv3d(x, w0, None, stride=(2,2,2), padding=(1,0,1))
         if torch.__version__ < '1.9':
             x = F.conv3d(x, w1, b1, stride=(1,1,1), padding=(1,1,1), dilation=(2,2,1), groups=2)
         else:
             x = F.conv3d(x, w1, b1, stride=(1,1,1), padding='same', dilation=(2,2,1), groups=2)
-        return x
+
+        y = F.conv3d(y, self.w2, self.b2, stride=(2,2,2), padding=(2,2,2))
+        y = F.conv3d(y, self.w3, None, stride=(2,2,2), padding=(1,1,1), groups=3)
+        return x, y
 
 def test():
     net = Model()
@@ -37,22 +44,23 @@ def test():
     w0 = torch.rand(16, 12, 3, 2, 3)
     w1 = torch.rand(16, 8, 5, 4, 5)
     b1 = torch.rand(16)
+    y = torch.rand(1, 6, 12, 11, 10)
 
-    a = net(x, w0, w1, b1)
+    a0, a1 = net(x, w0, w1, b1, y)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, w0, w1, b1))
+    mod = torch.jit.trace(net, (x, w0, w1, b1, y))
     mod.save("test_F_conv3d.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../src/pnnx test_F_conv3d.pt inputshape=[1,12,20,32,40],[16,12,3,2,3],[16,8,5,4,5],[16]")
+    os.system("../src/pnnx test_F_conv3d.pt inputshape=[1,12,20,32,40],[16,12,3,2,3],[16,8,5,4,5],[16],[1,6,12,11,10]")
 
     # pnnx inference
     import test_F_conv3d_pnnx
-    b = test_F_conv3d_pnnx.test_inference()
+    b0, b1 = test_F_conv3d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.equal(a0, b0) and torch.equal(a1, b1)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_F_conv_transpose1d.py

@@ -20,10 +20,17 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
-    def forward(self, x, w0, w1, b1):
+        self.w2 = nn.Parameter(torch.rand(6, 12, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(12, 2, 3))
+
+    def forward(self, x, w0, w1, b1, y):
         x = F.conv_transpose1d(x, w0, None, stride=2, padding=1, output_padding=1)
         x = F.conv_transpose1d(x, w1, b1, stride=1, padding=2, dilation=2, groups=2)
-        return x
+
+        y = F.conv_transpose1d(y, self.w2, self.b2, stride=2, padding=1, output_padding=1)
+        y = F.conv_transpose1d(y, self.w3, None, stride=1, padding=2, dilation=2, groups=3)
+        return x, y
 
 def test():
     net = Model()
@@ -34,22 +41,23 @@ def test():
     w0 = torch.rand(12, 16, 3)
     w1 = torch.rand(16, 8, 5)
     b1 = torch.rand(16)
+    y = torch.rand(1, 6, 5)
 
-    a = net(x, w0, w1, b1)
+    a0, a1 = net(x, w0, w1, b1, y)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, w0, w1, b1))
+    mod = torch.jit.trace(net, (x, w0, w1, b1, y))
     mod.save("test_F_conv_transpose1d.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../src/pnnx test_F_conv_transpose1d.pt inputshape=[1,12,22],[12,16,3],[16,8,5],[16]")
+    os.system("../src/pnnx test_F_conv_transpose1d.pt inputshape=[1,12,22],[12,16,3],[16,8,5],[16],[1,6,5]")
 
     # pnnx inference
     import test_F_conv_transpose1d_pnnx
-    b = test_F_conv_transpose1d_pnnx.test_inference()
+    b0, b1 = test_F_conv_transpose1d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.equal(a0, b0) and torch.equal(a1, b1)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_F_conv_transpose2d.py

@@ -20,10 +20,17 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
-    def forward(self, x, w0, w1, b1):
+        self.w2 = nn.Parameter(torch.rand(6, 12, 4, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(12, 2, 3, 3))
+
+    def forward(self, x, w0, w1, b1, y):
         x = F.conv_transpose2d(x, w0, None, stride=(2,2), padding=(1,1), output_padding=(1,1))
         x = F.conv_transpose2d(x, w1, b1, stride=(1,2), padding=(2,1), dilation=(2,1), groups=2)
-        return x
+
+        y = F.conv_transpose2d(y, self.w2, self.b2, stride=(2,2), padding=(1,1), output_padding=(1,1))
+        y = F.conv_transpose2d(y, self.w3, None, stride=(1,2), padding=(2,1), dilation=(2,1), groups=3)
+        return x, y
 
 def test():
     net = Model()
@@ -34,22 +41,23 @@ def test():
     w0 = torch.rand(12, 16, 3, 3)
     w1 = torch.rand(16, 8, 5, 5)
     b1 = torch.rand(16)
+    y = torch.rand(1, 6, 5, 6)
 
-    a = net(x, w0, w1, b1)
+    a0, a1 = net(x, w0, w1, b1, y)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, w0, w1, b1))
+    mod = torch.jit.trace(net, (x, w0, w1, b1, y))
     mod.save("test_F_conv_transpose2d.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../src/pnnx test_F_conv_transpose2d.pt inputshape=[1,12,22,32],[12,16,3,3],[16,8,5,5],[16]")
+    os.system("../src/pnnx test_F_conv_transpose2d.pt inputshape=[1,12,22,32],[12,16,3,3],[16,8,5,5],[16],[1,6,5,6]")
 
     # pnnx inference
     import test_F_conv_transpose2d_pnnx
-    b = test_F_conv_transpose2d_pnnx.test_inference()
+    b0, b1 = test_F_conv_transpose2d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.equal(a0, b0) and torch.equal(a1, b1)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_F_conv_transpose3d.py

@@ -20,10 +20,17 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
-    def forward(self, x, w0, w1, b1):
+        self.w2 = nn.Parameter(torch.rand(6, 12, 4, 4, 4))
+        self.b2 = nn.Parameter(torch.rand(12))
+        self.w3 = nn.Parameter(torch.rand(12, 2, 3, 3, 3))
+
+    def forward(self, x, w0, w1, b1, y):
         x = F.conv_transpose3d(x, w0, None, stride=(2,2,2), padding=(1,0,1), output_padding=(1,1,0))
         x = F.conv_transpose3d(x, w1, b1, stride=(1,1,2), padding=(2,2,1), dilation=(2,2,1), groups=2)
-        return x
+
+        y = F.conv_transpose3d(y, self.w2, self.b2, stride=(2,2, 2), padding=(1,0, 1), output_padding=(1,1, 0))
+        y = F.conv_transpose3d(y, self.w3, None, stride=(1,1,2), padding=(2,2,1), dilation=(2,2,1), groups=3)
+        return x, y
 
 def test():
     net = Model()
@@ -34,22 +41,23 @@ def test():
     w0 = torch.rand(12, 16, 3, 2, 3)
     w1 = torch.rand(16, 8, 5, 4, 5)
     b1 = torch.rand(16)
+    y = torch.rand(1, 6, 4, 5, 6)
 
-    a = net(x, w0, w1, b1)
+    a0, a1 = net(x, w0, w1, b1, y)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, w0, w1, b1))
+    mod = torch.jit.trace(net, (x, w0, w1, b1, y))
     mod.save("test_F_conv_transpose3d.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../src/pnnx test_F_conv_transpose3d.pt inputshape=[1,12,10,12,14],[12,16,3,2,3],[16,8,5,4,5],[16]")
+    os.system("../src/pnnx test_F_conv_transpose3d.pt inputshape=[1,12,10,12,14],[12,16,3,2,3],[16,8,5,4,5],[16],[1,6,4,5,6]")
 
     # pnnx inference
     import test_F_conv_transpose3d_pnnx
-    b = test_F_conv_transpose3d_pnnx.test_inference()
+    b0, b1 = test_F_conv_transpose3d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.equal(a0, b0) and torch.equal(a1, b1)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_F_group_norm.py

@@ -20,15 +20,25 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
+        self.w3 = nn.Parameter(torch.rand(16))
+        self.b3 = nn.Parameter(torch.rand(16))
+        self.w4 = nn.Parameter(torch.rand(12))
+        self.b4 = nn.Parameter(torch.rand(12))
+        self.w5 = nn.Parameter(torch.rand(32))
+        self.b5 = nn.Parameter(torch.rand(32))
+
     def forward(self, x, y, z, w0, b0, w1, b1, w2, b2):
         x = F.group_norm(x, 2, w0, b0)
         x = F.group_norm(x, 1, None, None)
+        x = F.group_norm(x, 4, self.w3, self.b3)
 
         y = F.group_norm(y, 3, w1, b1, eps=1e-4)
         y = F.group_norm(y, 4, None, None)
+        y = F.group_norm(y, 6, self.w4, self.b4)
 
         z = F.group_norm(z, 32, w2, b2)
         z = F.group_norm(z, 4, None, None, eps=1e-2)
+        z = F.group_norm(z, 8, self.w5, self.b5)
         return x, y, z
 
 def test():

tools/pnnx/tests/test_F_instance_norm.py

@@ -20,15 +20,31 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
+        self.m3 = torch.rand(12)
+        self.v3 = torch.rand(12)
+        self.w3 = nn.Parameter(torch.rand(12))
+        self.b3 = nn.Parameter(torch.rand(12))
+        self.m4 = torch.rand(3)
+        self.v4 = torch.rand(3)
+        self.w4 = nn.Parameter(torch.rand(3))
+        self.b4 = nn.Parameter(torch.rand(3))
+        self.m5 = torch.rand(10)
+        self.v5 = torch.rand(10)
+        self.w5 = nn.Parameter(torch.rand(10))
+        self.b5 = nn.Parameter(torch.rand(10))
+
     def forward(self, x, y, z, m0, v0, w0, b0, m1, v1, w1, b1, m2, v2, w2, b2):
         x = F.instance_norm(x, m0, v0, w0, b0)
         x = F.instance_norm(x, m0, v0, None, None)
+        x = F.instance_norm(x, self.m3, self.v3, self.w3, self.b3)
 
         y = F.instance_norm(y, m1, v1, w1, b1, eps=1e-3)
         y = F.instance_norm(y, m1, v1, None, None)
+        y = F.instance_norm(y, self.m4, self.v4, self.w4, self.b4)
 
         z = F.instance_norm(z, m2, v2, w2, b2)
         z = F.instance_norm(z, m2, v2, None, None, eps=1e-2)
+        z = F.instance_norm(z, self.m5, self.v5, self.w5, self.b5)
         return x, y, z
 
 def test():

tools/pnnx/tests/test_F_layer_norm.py

@@ -20,15 +20,25 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
+        self.w3 = nn.Parameter(torch.rand(24))
+        self.b3 = nn.Parameter(torch.rand(24))
+        self.w4 = nn.Parameter(torch.rand(12, 16))
+        self.b4 = nn.Parameter(torch.rand(12, 16))
+        self.w5 = nn.Parameter(torch.rand(24))
+        self.b5 = nn.Parameter(torch.rand(24))
+
     def forward(self, x, y, z, w0, b0, w1, b1, w2, b2):
         x = F.layer_norm(x, (24,), w0, b0)
         x = F.layer_norm(x, (12,24), None, None)
+        x = F.layer_norm(x, (24,), self.w3, self.b3)
 
         y = F.layer_norm(y, (16,), None, None, eps=1e-3)
         y = F.layer_norm(y, (12,16), w1, b1)
+        y = F.layer_norm(y, (12,16), self.w4, self.b4)
 
         z = F.layer_norm(z, (24,), w2, b2)
         z = F.layer_norm(z, (12,16,24), None, None, eps=1e-2)
+        z = F.layer_norm(z, (24,), self.w5, self.b5)
         return x, y, z
 
 def test():

tools/pnnx/tests/test_F_prelu.py

@@ -20,11 +20,20 @@ class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
 
+        self.w4 = nn.Parameter(torch.rand(16))
+        self.w5 = nn.Parameter(torch.rand(2))
+        self.w6 = nn.Parameter(torch.rand(3))
+        self.w7 = nn.Parameter(torch.rand(1))
+
     def forward(self, x, y, z, w, w0, w1, w2, w3):
         x = F.prelu(x, w0)
+        x = F.prelu(x, self.w4)
         y = F.prelu(y, w1)
+        y = F.prelu(y, self.w5)
         z = F.prelu(z, w2)
+        z = F.prelu(z, self.w6)
         w = F.prelu(w, w3)
+        w = F.prelu(w, self.w7)
         return x, y, z, w
 
 def test():