pnnx convert onnx layernorm instancenorm groupnorm (#5533)

* pnnx convert onnx layernorm

* fuse early

* skip layernorm affine false test for torch 2.1

* pnnx convert onnx layernorm instancenorm groupnorm

* take num_features from input shape for instancenorm module

* torch < 1.10 can not handle track_running_stats=True

tools/pnnx/src/pass_level1/nn_InstanceNorm1d.cpp

@@ -65,6 +65,12 @@ public:
             op->attrs["running_mean"] = running_mean;
             op->attrs["running_var"] = mod.attr("running_var").toTensor();
         }
+
+        // take num_features from input shape
+        if (!op->has_param("num_features") && !op->inputs[0]->shape.empty())
+        {
+            op->params["num_features"] = op->inputs[0]->shape[op->inputs[0]->shape.size() - 2];
+        }
     }
 };
 

tools/pnnx/src/pass_level1/nn_InstanceNorm2d.cpp

@@ -65,6 +65,12 @@ public:
             op->attrs["running_mean"] = running_mean;
             op->attrs["running_var"] = mod.attr("running_var").toTensor();
         }
+
+        // take num_features from input shape
+        if (!op->has_param("num_features") && !op->inputs[0]->shape.empty())
+        {
+            op->params["num_features"] = op->inputs[0]->shape[op->inputs[0]->shape.size() - 2];
+        }
     }
 };
 

tools/pnnx/src/pass_level1/nn_InstanceNorm3d.cpp

@@ -65,6 +65,12 @@ public:
             op->attrs["running_mean"] = running_mean;
             op->attrs["running_var"] = mod.attr("running_var").toTensor();
         }
+
+        // take num_features from input shape
+        if (!op->has_param("num_features") && !op->inputs[0]->shape.empty())
+        {
+            op->params["num_features"] = op->inputs[0]->shape[op->inputs[0]->shape.size() - 2];
+        }
     }
 };
 

tools/pnnx/src/pass_level2/F_group_norm.cpp

@@ -42,4 +42,87 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_group_norm, 10)
 
+class F_group_norm_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+7 6
+pnnx.Input              input_0     0 1 input
+Reshape                 op_0        1 1 input r1 allowzero=0 shape=(0,%num_groups,-1)
+pnnx.Attribute          op_1        0 1 ones @data
+pnnx.Attribute          op_2        0 1 zeros @data
+InstanceNormalization   op_3        3 1 r1 ones zeros in epsilon=%epsilon
+Reshape                 op_4        1 1 in out allowzero=0 shape=%shape
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.group_norm";
+    }
+
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        const Operator* op_reshape = matched_operators.at("op_0");
+        const std::vector<int>& inputshape = op_reshape->inputs[0]->shape;
+        if (inputshape != captured_params.at("shape").ai)
+            return false;
+
+        const int num_groups = captured_params.at("num_groups").i;
+
+        const Attribute& ones = captured_attrs.at("op_1.data");
+        const Attribute& zeros = captured_attrs.at("op_2.data");
+
+        if (ones.shape.size() != 1 || ones.shape[0] != num_groups)
+            return false;
+        if (zeros.shape.size() != 1 || zeros.shape[0] != num_groups)
+            return false;
+
+        for (auto x : ones.get_float32_data())
+        {
+            if (x != 1.f)
+                return false;
+        }
+
+        for (auto x : zeros.get_float32_data())
+        {
+            if (x != 0.f)
+                return false;
+        }
+
+        return true;
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
+    {
+        op->params["num_groups"] = captured_params.at("num_groups");
+        op->params["eps"] = captured_params.at("epsilon");
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_group_norm_onnx, 9)
+
+class F_group_norm_onnx_1 : public F_group_norm_onnx
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+7 6
+pnnx.Input              input_0     0 1 input
+Reshape                 op_0        1 1 input r1 shape=(0,%num_groups,-1)
+pnnx.Attribute          op_1        0 1 ones @data
+pnnx.Attribute          op_2        0 1 zeros @data
+InstanceNormalization   op_3        3 1 r1 ones zeros in epsilon=%epsilon
+Reshape                 op_4        1 1 in out shape=%shape
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_group_norm_onnx_1, 9)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/F_instance_norm.cpp

@@ -45,4 +45,34 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_instance_norm, 10)
 
+class F_instance_norm_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 weight
+pnnx.Input              input_2     0 1 bias
+InstanceNormalization   op_0        3 1 input weight bias out epsilon=%epsilon
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.instance_norm";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
+    {
+        op->params["eps"] = captured_params.at("epsilon");
+        op->params["running_mean"] = Parameter();
+        op->params["running_var"] = Parameter();
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_instance_norm_onnx, 10)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/F_layer_norm.cpp

@@ -134,4 +134,176 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_layer_norm_onnx_2, 10)
 
+class F_layer_norm_onnx_3 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+11 10
+pnnx.Input              input       0 1 input
+torch.mean              mean        1 1 input mean dim=%dim keepdim=True
+aten::sub               op_1        2 1 input mean pnnx_1
+prim::Constant          op_2        0 1 two value=2.000000e+00
+aten::pow               op_3        2 1 pnnx_1 two pnnx_2
+torch.mean              op_4        1 1 pnnx_2 var dim=%dim keepdim=True
+prim::Constant          op_5        0 1 eps value=%eps
+aten::add               op_6        2 1 var eps pnnx_4
+aten::sqrt              op_7        1 1 pnnx_4 pnnx_5
+aten::div               op_8        2 1 pnnx_1 pnnx_5 out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.layer_norm";
+    }
+
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
+    {
+        const Operator* op_mean = matched_operators.at("mean");
+        const std::vector<int>& inputshape = op_mean->inputs[0]->shape;
+        if (inputshape.empty())
+            return false;
+
+        // dim must be the last N dimensions
+        std::vector<int> dim = captured_params.at("dim").ai;
+
+        const int input_rank = (int)inputshape.size();
+        const int dim_count = (int)dim.size();
+
+        for (int i = 0; i < dim_count; i++)
+        {
+            if (dim[i] < 0)
+                dim[i] += input_rank;
+
+            if (dim[i] < input_rank - dim_count)
+                return false;
+        }
+
+        std::vector<int> normalized_shape(dim_count);
+        for (int i = 0; i < dim_count; i++)
+        {
+            normalized_shape[i] = inputshape[input_rank - dim_count + i];
+        }
+
+        return true;
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
+    {
+        const std::vector<int>& inputshape = op->inputs[0]->shape;
+        const std::vector<int>& dim = captured_params.at("dim").ai;
+        const int input_rank = (int)inputshape.size();
+        const int dim_count = (int)dim.size();
+
+        std::vector<int> normalized_shape(dim_count);
+        for (int i = 0; i < dim_count; i++)
+        {
+            normalized_shape[i] = inputshape[input_rank - dim_count + i];
+        }
+
+        op->params["normalized_shape"] = normalized_shape;
+        op->params["eps"] = captured_params.at("eps");
+        op->params["weight"] = Parameter();
+        op->params["bias"] = Parameter();
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_layer_norm_onnx_3, 30)
+
+class F_layer_norm_onnx_4 : 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
+torch.mean              mean        1 1 input mean dim=%dim keepdim=True
+aten::sub               op_1        2 1 input mean pnnx_1
+prim::Constant          op_2        0 1 two value=2.000000e+00
+aten::pow               op_3        2 1 pnnx_1 two pnnx_2
+torch.mean              op_4        1 1 pnnx_2 var dim=%dim keepdim=True
+prim::Constant          op_5        0 1 eps value=%eps
+aten::add               op_6        2 1 var eps pnnx_4
+aten::sqrt              op_7        1 1 pnnx_4 pnnx_5
+aten::div               op_8        2 1 pnnx_1 pnnx_5 pnnx_6
+aten::mul               mul         2 1 pnnx_6 weight pnnx_7
+aten::add               add         2 1 pnnx_7 bias out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.layer_norm";
+    }
+
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
+    {
+        const Operator* op_mean = matched_operators.at("mean");
+        const std::vector<int>& inputshape = op_mean->inputs[0]->shape;
+        if (inputshape.empty())
+            return false;
+
+        // dim must be the last N dimensions
+        std::vector<int> dim = captured_params.at("dim").ai;
+
+        const int input_rank = (int)inputshape.size();
+        const int dim_count = (int)dim.size();
+
+        for (int i = 0; i < dim_count; i++)
+        {
+            if (dim[i] < 0)
+                dim[i] += input_rank;
+
+            if (dim[i] < input_rank - dim_count)
+                return false;
+        }
+
+        std::vector<int> normalized_shape(dim_count);
+        for (int i = 0; i < dim_count; i++)
+        {
+            normalized_shape[i] = inputshape[input_rank - dim_count + i];
+        }
+
+        // check weight and bias shape
+        const Operator* op_mul = matched_operators.at("mul");
+        const Operator* op_add = matched_operators.at("add");
+        const std::vector<int>& weight_shape = op_mul->inputs[1]->shape;
+        const std::vector<int>& bias_shape = op_add->inputs[1]->shape;
+
+        if (weight_shape != normalized_shape)
+            return false;
+
+        if (bias_shape != normalized_shape)
+            return false;
+
+        return true;
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
+    {
+        const std::vector<int>& inputshape = op->inputs[0]->shape;
+        const std::vector<int>& dim = captured_params.at("dim").ai;
+        const int input_rank = (int)inputshape.size();
+        const int dim_count = (int)dim.size();
+
+        std::vector<int> normalized_shape(dim_count);
+        for (int i = 0; i < dim_count; i++)
+        {
+            normalized_shape[i] = inputshape[input_rank - dim_count + i];
+        }
+
+        op->params["normalized_shape"] = normalized_shape;
+        op->params["eps"] = captured_params.at("eps");
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_layer_norm_onnx_4, 29)
+
 } // namespace pnnx

tools/pnnx/src/pass_level5/fuse_layernorm.cpp

@@ -19,8 +19,6 @@
 #include <math.h>
 #include <string.h>
 
-#include <torch/csrc/api/include/torch/version.h>
-
 namespace pnnx {
 
 class fuse_layernorm_pass : public GraphRewriterPass
@@ -65,139 +63,12 @@ pnnx.Output             output      1 0 out
     }
 };
 
-class fuse_layernorm_pass_1 : public GraphRewriterPass
-{
-public:
-    const char* match_pattern_graph() const
-    {
-        // clang-format off
-        // *INDENT-OFF*
-
-        return R"PNNXIR(7767517
-9 8
-pnnx.Input              input       0 1 input #input=(1,?,%c)f32
-pnnx.Attribute          op_0        0 1 weight @data #weight=(%c)f32
-pnnx.Attribute          op_1        0 1 bias @data #bias=(%c)f32
-torch.mean              op_2        1 1 input mean dim=(-1) keepdim=True
-pnnx.Expression         op_3        2 1 input mean 173 expr=sub(@0,@1)
-pnnx.Expression         op_4        1 1 173 174 expr=pow(@0,2.000000e+00)
-torch.mean              op_5        1 1 174 var dim=(-1) keepdim=True
-pnnx.Expression         op_6        4 1 173 var weight bias out expr=add(mul(div(@0,sqrt(add(@1,%eps))),@2),@3)
-pnnx.Output             output      1 0 out
-)PNNXIR";
-
-        // *INDENT-ON*
-        // clang-format on
-    }
-
-    const char* replace_pattern_graph() const
-    {
-        return R"PNNXIR(7767517
-3 2
-pnnx.Input              input       0 1 input
-nn.LayerNorm            ln          1 1 input out elementwise_affine=True eps=%eps normalized_shape=(%c) @weight=%op_0.data @bias=%op_1.data
-pnnx.Output             output      1 0 out
-)PNNXIR";
-    }
-};
-
-class fuse_layernorm_pass_1_1 : public fuse_layernorm_pass_1
-{
-public:
-    const char* match_pattern_graph() const
-    {
-        // clang-format off
-        // *INDENT-OFF*
-        return R"PNNXIR(7767517
-9 8
-pnnx.Input              input       0 1 input #input=(1,?,%c)f32
-pnnx.Attribute          op_0        0 1 weight @data #weight=(%c)f32
-pnnx.Attribute          op_1        0 1 bias @data #bias=(%c)f32
-torch.mean              op_2        1 1 input mean dim=(2) keepdim=True #input=(1,?,%c)f32
-pnnx.Expression         op_3        2 1 input mean 75 expr=sub(@0,@1)
-pnnx.Expression         op_4        1 1 75 76 expr=pow(@0,2.000000e+00)
-torch.mean              op_5        1 1 76 var dim=(2) keepdim=True
-pnnx.Expression         op_6        4 1 75 var weight bias out expr=add(mul(div(@0,sqrt(add(@1,%eps))),@2),@3)
-pnnx.Output             output      1 0 out
-)PNNXIR";
-
-        // *INDENT-ON*
-        // clang-format on
-    }
-};
-
-class fuse_layernorm_pass_2 : public GraphRewriterPass
-{
-public:
-    const char* match_pattern_graph() const
-    {
-        // clang-format off
-        // *INDENT-OFF*
-
-        return R"PNNXIR(7767517
-7 6
-pnnx.Input              input       0 1 input #input=(1,?,%c)f32
-torch.mean              op_0        1 1 input mean dim=(-1) keepdim=True
-pnnx.Expression         op_1        2 1 input mean 25 expr=sub(@0,@1)
-pnnx.Expression         op_2        1 1 25 26 expr=pow(@0,2.000000e+00)
-torch.mean              op_3        1 1 26 var dim=(-1) keepdim=True
-pnnx.Expression         op_4        2 1 25 var out expr=div(@0,sqrt(add(@1,%eps)))
-pnnx.Output             output      1 0 out
-)PNNXIR";
-
-        // *INDENT-ON*
-        // clang-format on
-    }
-
-    const char* replace_pattern_graph() const
-    {
-        return R"PNNXIR(7767517
-3 2
-pnnx.Input              input       0 1 input
-nn.LayerNorm            ln          1 1 input out elementwise_affine=False eps=%eps normalized_shape=(%c)
-pnnx.Output             output      1 0 out
-)PNNXIR";
-    }
-};
-
-class fuse_layernorm_pass_2_1 : public fuse_layernorm_pass_2
-{
-public:
-    const char* match_pattern_graph() const
-    {
-        // clang-format off
-        // *INDENT-OFF*
-
-        return R"PNNXIR(7767517
-7 6
-pnnx.Input              input       0 1 input #input=(1,?,?,%c)f32
-torch.mean              op_0        1 1 input mean dim=(-1) keepdim=True
-pnnx.Expression         op_1        2 1 input mean 32 expr=sub(@0,@1)
-pnnx.Expression         op_2        1 1 32 33 expr=pow(@0,2.000000e+00)
-torch.mean              op_3        1 1 33 var dim=(-1) keepdim=True
-pnnx.Expression         op_4        2 1 32 var out expr=div(@0,sqrt(add(@1,%eps)))
-pnnx.Output             output      1 0 out
-)PNNXIR";
-
-        // *INDENT-ON*
-        // clang-format on
-    }
-};
-
 void fuse_layernorm(Graph& graph)
 {
     fuse_layernorm_pass a;
-    fuse_layernorm_pass_1 b;
-    fuse_layernorm_pass_1_1 b1;
-    fuse_layernorm_pass_2 c;
-    fuse_layernorm_pass_2_1 c1;
     int opindex = 0;
 
     pnnx_graph_rewrite(graph, &a, opindex);
-    pnnx_graph_rewrite(graph, &b, opindex);
-    pnnx_graph_rewrite(graph, &b1, opindex);
-    pnnx_graph_rewrite(graph, &c, opindex);
-    pnnx_graph_rewrite(graph, &c1, opindex);
 }
 
 } // namespace pnnx

tools/pnnx/src/pass_onnx.cpp

@@ -816,6 +816,8 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
                         is_attr_weight = true;
                     if (sim_op_type == "ConvTranspose" && (j == 1 || j == 2))
                         is_attr_weight = true;
+                    if (sim_op_type == "InstanceNormalization" && (j == 1 || j == 2))
+                        is_attr_weight = true;
                 }
 
                 int64_t numel = 1;

tools/pnnx/tests/ncnn/test_nn_InstanceNorm2d.py

@@ -21,7 +21,9 @@ class Model(nn.Module):
         super(Model, self).__init__()
 
         self.in_0 = nn.InstanceNorm2d(num_features=12, affine=True)
-        self.in_1 = nn.InstanceNorm2d(num_features=12, eps=1e-2, affine=True)
+        self.in_0.weight = nn.Parameter(torch.rand(12))
+        self.in_0.bias = nn.Parameter(torch.rand(12))
+        self.in_1 = nn.InstanceNorm2d(num_features=12, eps=1e-2, affine=False)
 
     def forward(self, x):
         x = self.in_0(x)

tools/pnnx/tests/ncnn/test_nn_LayerNorm.py

@@ -21,6 +21,8 @@ class Model(nn.Module):
         super(Model, self).__init__()
 
         self.ln_0 = nn.LayerNorm(64)
+        self.ln_0.weight = nn.Parameter(torch.rand(64))
+        self.ln_0.bias = nn.Parameter(torch.rand(64))
         self.ln_1 = nn.LayerNorm(normalized_shape=(24,64), eps=1e-2, elementwise_affine=False)
 
     def forward(self, x, y):

tools/pnnx/tests/onnx/CMakeLists.txt

@@ -18,7 +18,7 @@ pnnx_onnx_add_test(F_conv3d)
 # pnnx_onnx_add_test(F_group_norm)
 # pnnx_onnx_add_test(F_instance_norm)
 pnnx_onnx_add_test(F_interpolate)
-# pnnx_onnx_add_test(F_layer_norm)
+pnnx_onnx_add_test(F_layer_norm)
 pnnx_onnx_add_test(F_linear)
 pnnx_onnx_add_test(F_local_response_norm)
 pnnx_onnx_add_test(F_max_pool1d)
@@ -47,11 +47,11 @@ pnnx_onnx_add_test(nn_Conv3d)
 pnnx_onnx_add_test(nn_ConvTranspose1d)
 pnnx_onnx_add_test(nn_ConvTranspose2d)
 pnnx_onnx_add_test(nn_ConvTranspose3d)
-# pnnx_onnx_add_test(nn_GroupNorm)
-# pnnx_onnx_add_test(nn_InstanceNorm1d)
-# pnnx_onnx_add_test(nn_InstanceNorm2d)
-# pnnx_onnx_add_test(nn_InstanceNorm3d)
-# pnnx_onnx_add_test(nn_LayerNorm)
+pnnx_onnx_add_test(nn_GroupNorm)
+pnnx_onnx_add_test(nn_InstanceNorm1d)
+pnnx_onnx_add_test(nn_InstanceNorm2d)
+pnnx_onnx_add_test(nn_InstanceNorm3d)
+pnnx_onnx_add_test(nn_LayerNorm)
 pnnx_onnx_add_test(nn_Linear)
 pnnx_onnx_add_test(nn_LocalResponseNorm)
 pnnx_onnx_add_test(nn_MaxPool1d)

tools/pnnx/tests/onnx/test_F_layer_norm.py

@@ -0,0 +1,82 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2024 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
+from packaging import version
+
+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)
+        if version.parse(torch.__version__) < version.parse('2.1') or version.parse(torch.__version__) >= version.parse('2.2'):
+            x = F.layer_norm(x, (12,24), None, None)
+        x = F.layer_norm(x, (24,), self.w3, self.b3)
+
+        if version.parse(torch.__version__) < version.parse('2.1') or version.parse(torch.__version__) >= version.parse('2.2'):
+            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)
+        if version.parse(torch.__version__) < version.parse('2.1') or version.parse(torch.__version__) >= version.parse('2.2'):
+            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():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 12, 24)
+    y = torch.rand(2, 3, 12, 16)
+    z = torch.rand(1, 10, 12, 16, 24)
+    w0 = torch.rand(24)
+    b0 = torch.rand(24)
+    w1 = torch.rand(12, 16)
+    b1 = torch.rand(12, 16)
+    w2 = torch.rand(24)
+    b2 = torch.rand(24)
+
+    a0, a1, a2 = net(x, y, z, w0, b0, w1, b1, w2, b2)
+
+    # export onnx
+    torch.onnx.export(net, (x, y, z, w0, b0, w1, b1, w2, b2), "test_F_layer_norm.onnx")
+
+    # onnx to pnnx
+    import os
+    os.system("../../src/pnnx test_F_layer_norm.onnx inputshape=[1,12,24],[2,3,12,16],[1,10,12,16,24],[24],[24],[12,16],[12,16],[24],[24]")
+
+    # pnnx inference
+    import test_F_layer_norm_pnnx
+    b0, b1, b2 = test_F_layer_norm_pnnx.test_inference()
+
+    return torch.equal(a0, b0) and torch.equal(a1, b1) and torch.equal(a2, b2)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/onnx/test_nn_GroupNorm.py

@@ -0,0 +1,73 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2024 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.gn_0 = nn.GroupNorm(num_groups=4, num_channels=12)
+        self.gn_0.weight = nn.Parameter(torch.rand(12))
+        self.gn_0.bias = nn.Parameter(torch.rand(12))
+        self.gn_1 = nn.GroupNorm(num_groups=12, num_channels=12, eps=1e-2, affine=False)
+        self.gn_2 = nn.GroupNorm(num_groups=1, num_channels=12, eps=1e-4, affine=True)
+        self.gn_2.weight = nn.Parameter(torch.rand(12))
+        self.gn_2.bias = nn.Parameter(torch.rand(12))
+
+    def forward(self, x, y, z):
+        x = self.gn_0(x)
+        x = self.gn_1(x)
+        x = self.gn_2(x)
+
+        y = self.gn_0(y)
+        y = self.gn_1(y)
+        y = self.gn_2(y)
+
+        z = self.gn_0(z)
+        z = self.gn_1(z)
+        z = self.gn_2(z)
+        return x, y, z
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 12, 64)
+    y = torch.rand(1, 12, 24, 64)
+    z = torch.rand(1, 12, 24, 32, 64)
+
+    a0, a1, a2 = net(x, y, z)
+
+    # export onnx
+    torch.onnx.export(net, (x, y, z), "test_nn_GroupNorm.onnx")
+
+    # onnx to pnnx
+    import os
+    os.system("../../src/pnnx test_nn_GroupNorm.onnx inputshape=[1,12,64],[1,12,24,64],[1,12,24,32,64]")
+
+    # pnnx inference
+    import test_nn_GroupNorm_pnnx
+    b0, b1, b2 = test_nn_GroupNorm_pnnx.test_inference()
+
+    return torch.allclose(a0, b0, 1e-4, 1e-4) and torch.allclose(a1, b1, 1e-4, 1e-4) and torch.allclose(a2, b2, 1e-4, 1e-4)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/onnx/test_nn_InstanceNorm1d.py

@@ -0,0 +1,67 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2024 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
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.in_0 = nn.InstanceNorm1d(num_features=12, affine=True)
+        self.in_0.weight = nn.Parameter(torch.rand(12))
+        self.in_0.bias = nn.Parameter(torch.rand(12))
+        self.in_1 = nn.InstanceNorm1d(num_features=12, eps=1e-2, affine=False)
+        if version.parse(torch.__version__) < version.parse('1.10'):
+            self.in_2 = nn.InstanceNorm1d(num_features=12, eps=1e-4, affine=True, track_running_stats=False)
+        else:
+            self.in_2 = nn.InstanceNorm1d(num_features=12, eps=1e-4, affine=True, track_running_stats=True)
+        self.in_2.weight = nn.Parameter(torch.rand(12))
+        self.in_2.bias = nn.Parameter(torch.rand(12))
+
+    def forward(self, x):
+        x = self.in_0(x)
+        x = self.in_1(x)
+        x = self.in_2(x)
+        return x
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 12, 24)
+
+    a = net(x)
+
+    # export onnx
+    torch.onnx.export(net, (x,), "test_nn_InstanceNorm1d.onnx")
+
+    # onnx to pnnx
+    import os
+    os.system("../../src/pnnx test_nn_InstanceNorm1d.onnx inputshape=[1,12,24]")
+
+    # pnnx inference
+    import test_nn_InstanceNorm1d_pnnx
+    b = test_nn_InstanceNorm1d_pnnx.test_inference()
+
+    return torch.equal(a, b)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/onnx/test_nn_InstanceNorm2d.py

@@ -0,0 +1,67 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2024 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
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.in_0 = nn.InstanceNorm2d(num_features=12, affine=True)
+        self.in_0.weight = nn.Parameter(torch.rand(12))
+        self.in_0.bias = nn.Parameter(torch.rand(12))
+        self.in_1 = nn.InstanceNorm2d(num_features=12, eps=1e-2, affine=False)
+        if version.parse(torch.__version__) < version.parse('1.10'):
+            self.in_2 = nn.InstanceNorm2d(num_features=12, eps=1e-4, affine=True, track_running_stats=False)
+        else:
+            self.in_2 = nn.InstanceNorm2d(num_features=12, eps=1e-4, affine=True, track_running_stats=True)
+        self.in_2.weight = nn.Parameter(torch.rand(12))
+        self.in_2.bias = nn.Parameter(torch.rand(12))
+
+    def forward(self, x):
+        x = self.in_0(x)
+        x = self.in_1(x)
+        x = self.in_2(x)
+        return x
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 12, 24, 64)
+
+    a = net(x)
+
+    # export onnx
+    torch.onnx.export(net, (x,), "test_nn_InstanceNorm2d.onnx")
+
+    # onnx to pnnx
+    import os
+    os.system("../../src/pnnx test_nn_InstanceNorm2d.onnx inputshape=[1,12,24,64]")
+
+    # pnnx inference
+    import test_nn_InstanceNorm2d_pnnx
+    b = test_nn_InstanceNorm2d_pnnx.test_inference()
+
+    return torch.equal(a, b)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/onnx/test_nn_InstanceNorm3d.py

@@ -0,0 +1,67 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2024 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
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.in_0 = nn.InstanceNorm3d(num_features=12, affine=True)
+        self.in_0.weight = nn.Parameter(torch.rand(12))
+        self.in_0.bias = nn.Parameter(torch.rand(12))
+        self.in_1 = nn.InstanceNorm3d(num_features=12, eps=1e-2, affine=False)
+        if version.parse(torch.__version__) < version.parse('1.10'):
+            self.in_2 = nn.InstanceNorm3d(num_features=12, eps=1e-4, affine=True, track_running_stats=False)
+        else:
+            self.in_2 = nn.InstanceNorm3d(num_features=12, eps=1e-4, affine=True, track_running_stats=True)
+        self.in_2.weight = nn.Parameter(torch.rand(12))
+        self.in_2.bias = nn.Parameter(torch.rand(12))
+
+    def forward(self, x):
+        x = self.in_0(x)
+        x = self.in_1(x)
+        x = self.in_2(x)
+        return x
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 12, 24, 32, 64)
+
+    a = net(x)
+
+    # export onnx
+    torch.onnx.export(net, (x,), "test_nn_InstanceNorm3d.onnx")
+
+    # onnx to pnnx
+    import os
+    os.system("../../src/pnnx test_nn_InstanceNorm3d.onnx inputshape=[1,12,24,32,64]")
+
+    # pnnx inference
+    import test_nn_InstanceNorm3d_pnnx
+    b = test_nn_InstanceNorm3d_pnnx.test_inference()
+
+    return torch.equal(a, b)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/onnx/test_nn_LayerNorm.py

@@ -0,0 +1,73 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2024 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
+from packaging import version
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.ln_0 = nn.LayerNorm(64)
+        self.ln_0.weight = nn.Parameter(torch.rand(64))
+        self.ln_0.bias = nn.Parameter(torch.rand(64))
+        if version.parse(torch.__version__) >= version.parse('2.1') and version.parse(torch.__version__) < version.parse('2.2'):
+            self.ln_1 = nn.LayerNorm(normalized_shape=(24,64), eps=1e-2, elementwise_affine=True)
+            self.ln_1.weight = nn.Parameter(torch.rand(24,64))
+            self.ln_1.bias = nn.Parameter(torch.rand(24,64))
+        else:
+            self.ln_1 = nn.LayerNorm(normalized_shape=(24,64), eps=1e-2, elementwise_affine=False)
+
+    def forward(self, x, y, z):
+        x = self.ln_0(x)
+        x = self.ln_1(x)
+
+        y = self.ln_0(y)
+        y = self.ln_1(y)
+
+        z = self.ln_0(z)
+        z = self.ln_1(z)
+        return x, y, z
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 24, 64)
+    y = torch.rand(1, 12, 24, 64)
+    z = torch.rand(1, 12, 16, 24, 64)
+
+    a0, a1, a2 = net(x, y, z)
+
+    # export onnx
+    torch.onnx.export(net, (x, y, z), "test_nn_LayerNorm.onnx")
+
+    # onnx to pnnx
+    import os
+    os.system("../../src/pnnx test_nn_LayerNorm.onnx inputshape=[1,24,64],[1,12,24,64],[1,12,16,24,64]")
+
+    # pnnx inference
+    import test_nn_LayerNorm_pnnx
+    b0, b1, b2 = test_nn_LayerNorm_pnnx.test_inference()
+
+    return torch.equal(a0, b0) and torch.equal(a1, b1) and torch.equal(a2, b2)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/test_nn_InstanceNorm1d.py

@@ -21,8 +21,12 @@ class Model(nn.Module):
         super(Model, self).__init__()
 
         self.in_0 = nn.InstanceNorm1d(num_features=12, affine=True)
-        self.in_1 = nn.InstanceNorm1d(num_features=12, eps=1e-2, affine=True)
+        self.in_0.weight = nn.Parameter(torch.rand(12))
+        self.in_0.bias = nn.Parameter(torch.rand(12))
+        self.in_1 = nn.InstanceNorm1d(num_features=12, eps=1e-2, affine=False)
         self.in_2 = nn.InstanceNorm1d(num_features=12, eps=1e-4, affine=True, track_running_stats=True)
+        self.in_2.weight = nn.Parameter(torch.rand(12))
+        self.in_2.bias = nn.Parameter(torch.rand(12))
 
     def forward(self, x):
         x = self.in_0(x)

tools/pnnx/tests/test_nn_InstanceNorm2d.py

@@ -21,8 +21,12 @@ class Model(nn.Module):
         super(Model, self).__init__()
 
         self.in_0 = nn.InstanceNorm2d(num_features=12, affine=True)
-        self.in_1 = nn.InstanceNorm2d(num_features=12, eps=1e-2, affine=True)
+        self.in_0.weight = nn.Parameter(torch.rand(12))
+        self.in_0.bias = nn.Parameter(torch.rand(12))
+        self.in_1 = nn.InstanceNorm2d(num_features=12, eps=1e-2, affine=False)
         self.in_2 = nn.InstanceNorm2d(num_features=12, eps=1e-4, affine=True, track_running_stats=True)
+        self.in_2.weight = nn.Parameter(torch.rand(12))
+        self.in_2.bias = nn.Parameter(torch.rand(12))
 
     def forward(self, x):
         x = self.in_0(x)

tools/pnnx/tests/test_nn_InstanceNorm3d.py

@@ -21,8 +21,12 @@ class Model(nn.Module):
         super(Model, self).__init__()
 
         self.in_0 = nn.InstanceNorm3d(num_features=12, affine=True)
-        self.in_1 = nn.InstanceNorm3d(num_features=12, eps=1e-2, affine=True)
+        self.in_0.weight = nn.Parameter(torch.rand(12))
+        self.in_0.bias = nn.Parameter(torch.rand(12))
+        self.in_1 = nn.InstanceNorm3d(num_features=12, eps=1e-2, affine=False)
         self.in_2 = nn.InstanceNorm3d(num_features=12, eps=1e-4, affine=True, track_running_stats=True)
+        self.in_2.weight = nn.Parameter(torch.rand(12))
+        self.in_2.bias = nn.Parameter(torch.rand(12))
 
     def forward(self, x):
         x = self.in_0(x)

tools/pnnx/tests/test_nn_LayerNorm.py

@@ -21,6 +21,8 @@ class Model(nn.Module):
         super(Model, self).__init__()
 
         self.ln_0 = nn.LayerNorm(64)
+        self.ln_0.weight = nn.Parameter(torch.rand(64))
+        self.ln_0.bias = nn.Parameter(torch.rand(64))
         self.ln_1 = nn.LayerNorm(normalized_shape=(24,64), eps=1e-2, elementwise_affine=False)
 
     def forward(self, x, y, z):