Merge branch 'master' into shader_cache2

tools/pnnx/src/pass_level1/nn_Conv1d.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "fuse_module_pass.h"
+#include "utils.h"
 
 namespace pnnx {
 
@@ -41,7 +42,7 @@ public:
             convolution = convolution_mode;
         }
 
-        const TorchTensorProxy& weight = mod.attr("weight");
+        const TorchTensorProxy& weight = mod.hasattr("weight") ? mod.attr("weight") : mod.attr("weight_v");
 
         op->params["groups"] = convolution->namedInput("groups");
         op->params["in_channels"] = weight.size(1) * op->params["groups"].i;
@@ -114,6 +115,21 @@ public:
         op->params["bias"] = mod.hasattr("bias");
 
         op->attrs["weight"] = weight;
+        if (!mod.hasattr("weight"))
+        {
+            // weight norm
+            Attribute weight_g = mod.attr("weight_g");
+            std::vector<float> weight_data = op->attrs["weight"].get_float32_data();
+            std::vector<float> weight_g_data = weight_g.get_float32_data();
+            int outch = op->params.at("out_channels").i;
+            int inch = op->params.at("in_channels").i * op->params.at("kernel_size").ai[0];
+            apply_weight_norm(weight_data, weight_g_data, outch, inch);
+            op->attrs["weight"].set_float32_data(weight_data);
+
+            // drop the additional weight input
+            op->inputs[1]->remove_consumer(op);
+            op->inputs.resize(1);
+        }
         if (mod.hasattr("bias"))
         {
             op->attrs["bias"] = mod.attr("bias");

tools/pnnx/src/pass_level1/nn_Conv2d.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "fuse_module_pass.h"
+#include "utils.h"
 
 namespace pnnx {
 
@@ -41,7 +42,7 @@ public:
             convolution = convolution_mode;
         }
 
-        const TorchTensorProxy& weight = mod.attr("weight");
+        const TorchTensorProxy& weight = mod.hasattr("weight") ? mod.attr("weight") : mod.attr("weight_v");
 
         op->params["groups"] = convolution->namedInput("groups");
         op->params["in_channels"] = weight.size(1) * op->params["groups"].i;
@@ -114,6 +115,21 @@ public:
         op->params["bias"] = mod.hasattr("bias");
 
         op->attrs["weight"] = weight;
+        if (!mod.hasattr("weight"))
+        {
+            // weight norm
+            Attribute weight_g = mod.attr("weight_g");
+            std::vector<float> weight_data = op->attrs["weight"].get_float32_data();
+            std::vector<float> weight_g_data = weight_g.get_float32_data();
+            int outch = op->params.at("out_channels").i;
+            int inch = op->params.at("in_channels").i * op->params.at("kernel_size").ai[0] * op->params.at("kernel_size").ai[1];
+            apply_weight_norm(weight_data, weight_g_data, outch, inch);
+            op->attrs["weight"].set_float32_data(weight_data);
+
+            // drop the additional weight input
+            op->inputs[1]->remove_consumer(op);
+            op->inputs.resize(1);
+        }
         if (mod.hasattr("bias"))
         {
             op->attrs["bias"] = mod.attr("bias");

tools/pnnx/src/pass_level1/nn_Conv3d.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "fuse_module_pass.h"
+#include "utils.h"
 
 namespace pnnx {
 
@@ -41,7 +42,7 @@ public:
             convolution = convolution_mode;
         }
 
-        const TorchTensorProxy& weight = mod.attr("weight");
+        const TorchTensorProxy& weight = mod.hasattr("weight") ? mod.attr("weight") : mod.attr("weight_v");
 
         op->params["groups"] = convolution->namedInput("groups");
         op->params["in_channels"] = weight.size(1) * op->params["groups"].i;
@@ -114,6 +115,21 @@ public:
         op->params["bias"] = mod.hasattr("bias");
 
         op->attrs["weight"] = weight;
+        if (!mod.hasattr("weight"))
+        {
+            // weight norm
+            Attribute weight_g = mod.attr("weight_g");
+            std::vector<float> weight_data = op->attrs["weight"].get_float32_data();
+            std::vector<float> weight_g_data = weight_g.get_float32_data();
+            int outch = op->params.at("out_channels").i;
+            int inch = op->params.at("in_channels").i * op->params.at("kernel_size").ai[0] * op->params.at("kernel_size").ai[1] * op->params.at("kernel_size").ai[2];
+            apply_weight_norm(weight_data, weight_g_data, outch, inch);
+            op->attrs["weight"].set_float32_data(weight_data);
+
+            // drop the additional weight input
+            op->inputs[1]->remove_consumer(op);
+            op->inputs.resize(1);
+        }
         if (mod.hasattr("bias"))
         {
             op->attrs["bias"] = mod.attr("bias");

tools/pnnx/src/pass_level1/nn_ConvTranspose1d.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "fuse_module_pass.h"
+#include "utils.h"
 
 namespace pnnx {
 
@@ -22,7 +23,7 @@ public:
     {
         const TorchNodeProxy* convolution = graph.find_node_by_kind("aten::_convolution");
 
-        const TorchTensorProxy& weight = mod.attr("weight");
+        const TorchTensorProxy& weight = mod.hasattr("weight") ? mod.attr("weight") : mod.attr("weight_v");
 
         op->params["groups"] = convolution->namedInput("groups");
         op->params["in_channels"] = weight.size(0);
@@ -35,6 +36,21 @@ public:
         op->params["bias"] = mod.hasattr("bias");
 
         op->attrs["weight"] = weight;
+        if (!mod.hasattr("weight"))
+        {
+            // weight norm
+            Attribute weight_g = mod.attr("weight_g");
+            std::vector<float> weight_data = op->attrs["weight"].get_float32_data();
+            std::vector<float> weight_g_data = weight_g.get_float32_data();
+            int inch = op->params.at("in_channels").i;
+            int outch = op->params.at("out_channels").i * op->params.at("kernel_size").ai[0];
+            apply_weight_norm(weight_data, weight_g_data, inch, outch);
+            op->attrs["weight"].set_float32_data(weight_data);
+
+            // drop the additional weight input
+            op->inputs[1]->remove_consumer(op);
+            op->inputs.resize(1);
+        }
         if (mod.hasattr("bias"))
         {
             op->attrs["bias"] = mod.attr("bias");

tools/pnnx/src/pass_level1/nn_ConvTranspose2d.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "fuse_module_pass.h"
+#include "utils.h"
 
 namespace pnnx {
 
@@ -22,7 +23,7 @@ public:
     {
         const TorchNodeProxy* convolution = graph.find_node_by_kind("aten::_convolution");
 
-        const TorchTensorProxy& weight = mod.attr("weight");
+        const TorchTensorProxy& weight = mod.hasattr("weight") ? mod.attr("weight") : mod.attr("weight_v");
 
         op->params["groups"] = convolution->namedInput("groups");
         op->params["in_channels"] = weight.size(0);
@@ -35,6 +36,21 @@ public:
         op->params["bias"] = mod.hasattr("bias");
 
         op->attrs["weight"] = weight;
+        if (!mod.hasattr("weight"))
+        {
+            // weight norm
+            Attribute weight_g = mod.attr("weight_g");
+            std::vector<float> weight_data = op->attrs["weight"].get_float32_data();
+            std::vector<float> weight_g_data = weight_g.get_float32_data();
+            int inch = op->params.at("in_channels").i;
+            int outch = op->params.at("out_channels").i * op->params.at("kernel_size").ai[0] * op->params.at("kernel_size").ai[1];
+            apply_weight_norm(weight_data, weight_g_data, inch, outch);
+            op->attrs["weight"].set_float32_data(weight_data);
+
+            // drop the additional weight input
+            op->inputs[1]->remove_consumer(op);
+            op->inputs.resize(1);
+        }
         if (mod.hasattr("bias"))
         {
             op->attrs["bias"] = mod.attr("bias");

tools/pnnx/src/pass_level1/nn_ConvTranspose3d.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "fuse_module_pass.h"
+#include "utils.h"
 
 namespace pnnx {
 
@@ -22,7 +23,7 @@ public:
     {
         const TorchNodeProxy* convolution = graph.find_node_by_kind("aten::_convolution");
 
-        const TorchTensorProxy& weight = mod.attr("weight");
+        const TorchTensorProxy& weight = mod.hasattr("weight") ? mod.attr("weight") : mod.attr("weight_v");
 
         op->params["groups"] = convolution->namedInput("groups");
         op->params["in_channels"] = weight.size(0);
@@ -35,6 +36,21 @@ public:
         op->params["bias"] = mod.hasattr("bias");
 
         op->attrs["weight"] = weight;
+        if (!mod.hasattr("weight"))
+        {
+            // weight norm
+            Attribute weight_g = mod.attr("weight_g");
+            std::vector<float> weight_data = op->attrs["weight"].get_float32_data();
+            std::vector<float> weight_g_data = weight_g.get_float32_data();
+            int inch = op->params.at("in_channels").i;
+            int outch = op->params.at("out_channels").i * op->params.at("kernel_size").ai[0] * op->params.at("kernel_size").ai[1] * op->params.at("kernel_size").ai[2];
+            apply_weight_norm(weight_data, weight_g_data, inch, outch);
+            op->attrs["weight"].set_float32_data(weight_data);
+
+            // drop the additional weight input
+            op->inputs[1]->remove_consumer(op);
+            op->inputs.resize(1);
+        }
         if (mod.hasattr("bias"))
         {
             op->attrs["bias"] = mod.attr("bias");

tools/pnnx/src/pass_level1/nn_Linear.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "fuse_module_pass.h"
+#include "utils.h"
 
 namespace pnnx {
 
@@ -22,13 +23,28 @@ public:
     {
         // const TorchNodeProxy* addmm = graph.find_node_by_kind("aten::addmm");
 
-        const TorchTensorProxy& weight = mod.attr("weight");
+        const TorchTensorProxy& weight = mod.hasattr("weight") ? mod.attr("weight") : mod.attr("weight_v");
 
         op->params["in_features"] = weight.size(1);
         op->params["out_features"] = weight.size(0);
         op->params["bias"] = mod.hasattr("bias");
 
         op->attrs["weight"] = weight;
+        if (!mod.hasattr("weight"))
+        {
+            // weight norm
+            Attribute weight_g = mod.attr("weight_g");
+            std::vector<float> weight_data = op->attrs["weight"].get_float32_data();
+            std::vector<float> weight_g_data = weight_g.get_float32_data();
+            int outch = op->params.at("out_features").i;
+            int inch = op->params.at("in_features").i;
+            apply_weight_norm(weight_data, weight_g_data, outch, inch);
+            op->attrs["weight"].set_float32_data(weight_data);
+
+            // drop the additional weight input
+            op->inputs[1]->remove_consumer(op);
+            op->inputs.resize(1);
+        }
         if (mod.hasattr("bias"))
         {
             op->attrs["bias"] = mod.attr("bias");

tools/pnnx/src/pass_level5/fuse_static_linear.cpp

@@ -29,7 +29,7 @@ pnnx.Output             output      1 0 out
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-nn.Linear               linear      1 1 input out in_features=%in_features out_features=%out_features bias=False @weight=%op_weight.data
+nn.Linear               F_linear    1 1 input out in_features=%in_features out_features=%out_features bias=False @weight=%op_weight.data
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }
@@ -55,7 +55,7 @@ pnnx.Output             output      1 0 out
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-nn.Linear               linear      1 1 input out in_features=%in_features out_features=%out_features bias=True @weight=%op_weight.data @bias=%op_bias.data
+nn.Linear               F_linear    1 1 input out in_features=%in_features out_features=%out_features bias=True @weight=%op_weight.data @bias=%op_bias.data
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }
@@ -82,7 +82,7 @@ pnnx.Output             output      1 0 out
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-nn.Linear               linear      1 1 input out in_features=%in_features out_features=%out_features bias=True @weight=%op_weight.data @bias=%op_bias.data
+nn.Linear               F_linear    1 1 input out in_features=%in_features out_features=%out_features bias=True @weight=%op_weight.data @bias=%op_bias.data
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }

tools/pnnx/src/utils.cpp

@@ -3,6 +3,8 @@
 
 #include "utils.h"
 
+#include <math.h>
+
 namespace pnnx {
 
 unsigned short float32_to_float16(float value)
@@ -110,4 +112,25 @@ float float16_to_float32(unsigned short value)
     return tmp.f;
 }
 
+void apply_weight_norm(std::vector<float>& weight, const std::vector<float>& weight_g, int dim0, int size)
+{
+    for (int i = 0; i < dim0; i++)
+    {
+        float* pw = weight.data() + i * size;
+
+        double norm = 0.f;
+        for (int j = 0; j < size; j++)
+        {
+            float w = pw[j];
+            norm += w * w;
+        }
+        norm = sqrt(norm);
+
+        for (int j = 0; j < size; j++)
+        {
+            pw[j] = pw[j] * (weight_g[i] / norm);
+        }
+    }
+}
+
 } // namespace pnnx

tools/pnnx/src/utils.h

@@ -4,12 +4,16 @@
 #ifndef PNNX_UTILS_H
 #define PNNX_UTILS_H
 
+#include <vector>
+
 namespace pnnx {
 
 unsigned short float32_to_float16(float value);
 
 float float16_to_float32(unsigned short value);
 
+void apply_weight_norm(std::vector<float>& weight, const std::vector<float>& weight_g, int dim0, int size);
+
 } // namespace pnnx
 
 #endif // PNNX_UTILS_H

tools/pnnx/tests/ncnn/test_nn_Conv1d.py

@@ -22,6 +22,12 @@ class Model(nn.Module):
         self.conv_5 = nn.Conv1d(in_channels=32, out_channels=32, kernel_size=2, stride=2, padding=3, dilation=1, groups=32, bias=True, padding_mode='reflect')
         self.conv_6 = nn.Conv1d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, dilation=1, groups=1, bias=False, padding_mode='replicate')
 
+        self.conv_7 = nn.Conv1d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5), dilation=2, groups=1, bias=True)
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.conv_7 = torch.nn.utils.weight_norm(self.conv_7)
+        else:
+            self.conv_7 = torch.nn.utils.parametrizations.weight_norm(self.conv_7)
+
     def forward(self, x):
         x = self.conv_0(x)
         x = self.conv_1(x)
@@ -30,6 +36,7 @@ class Model(nn.Module):
         x = self.conv_4(x)
         x = self.conv_5(x)
         x = self.conv_6(x)
+        x = self.conv_7(x)
 
         return x
 
@@ -54,7 +61,7 @@ def test():
     import test_nn_Conv1d_ncnn
     b = test_nn_Conv1d_ncnn.test_inference()
 
-    return torch.allclose(a, b, 1e-4, 1e-4)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/ncnn/test_nn_Conv2d.py

@@ -22,6 +22,12 @@ class Model(nn.Module):
         self.conv_5 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=2, stride=2, padding=3, dilation=1, groups=32, bias=True, padding_mode='reflect')
         self.conv_6 = nn.Conv2d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, dilation=1, groups=1, bias=False, padding_mode='replicate')
 
+        self.conv_7 = nn.Conv2d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6), dilation=2, groups=1, bias=True)
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.conv_7 = torch.nn.utils.weight_norm(self.conv_7)
+        else:
+            self.conv_7 = torch.nn.utils.parametrizations.weight_norm(self.conv_7)
+
     def forward(self, x):
         x = self.conv_0(x)
         x = self.conv_1(x)
@@ -30,6 +36,7 @@ class Model(nn.Module):
         x = self.conv_4(x)
         x = self.conv_5(x)
         x = self.conv_6(x)
+        x = self.conv_7(x)
 
         return x
 
@@ -54,7 +61,7 @@ def test():
     import test_nn_Conv2d_ncnn
     b = test_nn_Conv2d_ncnn.test_inference()
 
-    return torch.allclose(a, b, 1e-4, 1e-4)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/ncnn/test_nn_Conv3d.py

@@ -22,7 +22,13 @@ class Model(nn.Module):
         if version.parse(torch.__version__) >= version.parse('1.10'):
             self.conv_5 = nn.Conv3d(in_channels=32, out_channels=32, kernel_size=2, stride=2, padding=3, dilation=1, groups=32, bias=True, padding_mode='reflect')
             self.conv_6 = nn.Conv3d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, dilation=1, groups=1, bias=False, padding_mode='replicate')
-            # self.conv_7 = nn.Conv3d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6), dilation=2, groups=1, bias=True, padding_mode='circular')
+            # self.conv_7 = nn.Conv3d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6,7), dilation=2, groups=1, bias=True, padding_mode='circular')
+
+        self.conv_8 = nn.Conv3d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6,7), dilation=2, groups=1, bias=True)
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.conv_8 = torch.nn.utils.weight_norm(self.conv_8)
+        else:
+            self.conv_8 = torch.nn.utils.parametrizations.weight_norm(self.conv_8)
 
     def forward(self, x):
         x = self.conv_0(x)
@@ -36,6 +42,7 @@ class Model(nn.Module):
         x = self.conv_5(x)
         x = self.conv_6(x)
         #x = self.conv_7(x)
+        x = self.conv_8(x)
 
         return x
 
@@ -60,7 +67,7 @@ def test():
     import test_nn_Conv3d_ncnn
     b = test_nn_Conv3d_ncnn.test_inference()
 
-    return torch.allclose(a, b, 1e-4, 1e-4)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/ncnn/test_nn_ConvTranspose1d.py

@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from packaging import version
 
 class Model(nn.Module):
     def __init__(self):
@@ -18,6 +19,11 @@ class Model(nn.Module):
         self.deconv_6 = nn.ConvTranspose1d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, output_padding=0, dilation=1, groups=1, bias=False)
         self.deconv_7 = nn.ConvTranspose1d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(6), output_padding=(1), dilation=2, groups=1, bias=True)
 
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.deconv_7 = torch.nn.utils.weight_norm(self.deconv_7)
+        else:
+            self.deconv_7 = torch.nn.utils.parametrizations.weight_norm(self.deconv_7)
+
         self.downsample = nn.Conv1d(24, 16, 3, stride=2, padding=1)
         self.upsample = nn.ConvTranspose1d(16, 24, 3, stride=2, padding=1)
 
@@ -57,7 +63,7 @@ def test():
     import test_nn_ConvTranspose1d_ncnn
     b = test_nn_ConvTranspose1d_ncnn.test_inference()
 
-    return torch.allclose(a, b, 1e-4, 1e-4)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/ncnn/test_nn_ConvTranspose2d.py

@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from packaging import version
 
 class Model(nn.Module):
     def __init__(self):
@@ -18,6 +19,11 @@ class Model(nn.Module):
         self.deconv_6 = nn.ConvTranspose2d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, output_padding=0, dilation=1, groups=1, bias=False)
         self.deconv_7 = nn.ConvTranspose2d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6), output_padding=(1,0), dilation=2, groups=1, bias=True)
 
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.deconv_7 = torch.nn.utils.weight_norm(self.deconv_7)
+        else:
+            self.deconv_7 = torch.nn.utils.parametrizations.weight_norm(self.deconv_7)
+
         self.downsample = nn.Conv2d(24, 16, 3, stride=2, padding=1)
         self.upsample = nn.ConvTranspose2d(16, 24, 3, stride=2, padding=1)
 
@@ -57,7 +63,7 @@ def test():
     import test_nn_ConvTranspose2d_ncnn
     b = test_nn_ConvTranspose2d_ncnn.test_inference()
 
-    return torch.allclose(a, b, 1e-4, 1e-4)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/ncnn/test_nn_ConvTranspose3d.py

@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from packaging import version
 
 class Model(nn.Module):
     def __init__(self):
@@ -18,6 +19,11 @@ class Model(nn.Module):
         self.deconv_6 = nn.ConvTranspose3d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, output_padding=0, dilation=1, groups=1, bias=False)
         self.deconv_7 = nn.ConvTranspose3d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6,7), output_padding=(1,0,1), dilation=2, groups=1, bias=True)
 
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.deconv_7 = torch.nn.utils.weight_norm(self.deconv_7)
+        else:
+            self.deconv_7 = torch.nn.utils.parametrizations.weight_norm(self.deconv_7)
+
         self.downsample = nn.Conv3d(24, 16, 3, stride=2, padding=1)
         self.upsample = nn.ConvTranspose3d(16, 24, 3, stride=2, padding=1)
 
@@ -57,7 +63,7 @@ def test():
     import test_nn_ConvTranspose3d_ncnn
     b = test_nn_ConvTranspose3d_ncnn.test_inference()
 
-    return torch.allclose(a, b, 1e-4, 1e-4)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/ncnn/test_nn_Linear.py

@@ -4,27 +4,41 @@
 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.linear_0 = nn.Linear(in_features=64, out_features=16, bias=False)
-        self.linear_1 = nn.Linear(in_features=16, out_features=3, bias=True)
+        self.linear_1 = nn.Linear(in_features=16, out_features=13, bias=True)
+
+        self.linear_2 = nn.Linear(in_features=13, out_features=17, bias=True)
+        if version.parse(torch.__version__) < version.parse('1.9'):
+            # weight_norm on torch 1.8 produces wrong output shape, skip it
+            pass
+        elif version.parse(torch.__version__) < version.parse('2.1'):
+            self.linear_2 = torch.nn.utils.weight_norm(self.linear_2)
+        else:
+            self.linear_2 = torch.nn.utils.parametrizations.weight_norm(self.linear_2)
 
     def forward(self, x, y, z, w):
         x = self.linear_0(x)
         x = self.linear_1(x)
+        x = self.linear_2(x)
 
         y = self.linear_0(y)
         y = self.linear_1(y)
+        y = self.linear_2(y)
 
         z = self.linear_0(z)
         z = self.linear_1(z)
+        z = self.linear_2(z)
         z = F.relu(z)
 
         w = self.linear_0(w)
         w = self.linear_1(w)
+        w = self.linear_2(w)
         return x, y, z, w
 
 def test():
@@ -52,7 +66,7 @@ def test():
     b = test_nn_Linear_ncnn.test_inference()
 
     for a0, b0 in zip(a, b):
-        if not torch.allclose(a0, b0, 1e-4, 1e-4):
+        if not torch.allclose(a0, b0, 1e-3, 1e-3):
             return False
     return True
 

tools/pnnx/tests/test_nn_Conv1d.py

@@ -21,7 +21,13 @@ class Model(nn.Module):
             self.conv_4 = nn.Conv1d(in_channels=28, out_channels=32, kernel_size=3, stride=1, padding='same', dilation=2, groups=2, bias=False, padding_mode='zeros')
         self.conv_5 = nn.Conv1d(in_channels=32, out_channels=32, kernel_size=2, stride=2, padding=3, dilation=1, groups=32, bias=True, padding_mode='reflect')
         self.conv_6 = nn.Conv1d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, dilation=1, groups=1, bias=False, padding_mode='replicate')
-        #self.conv_7 = nn.Conv1d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6), dilation=2, groups=1, bias=True, padding_mode='circular')
+        #self.conv_7 = nn.Conv1d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5), dilation=2, groups=1, bias=True, padding_mode='circular')
+
+        self.conv_8 = nn.Conv1d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5), dilation=2, groups=1, bias=True)
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.conv_8 = torch.nn.utils.weight_norm(self.conv_8)
+        else:
+            self.conv_8 = torch.nn.utils.parametrizations.weight_norm(self.conv_8)
 
     def forward(self, x):
         x = self.conv_0(x)
@@ -32,6 +38,7 @@ class Model(nn.Module):
         x = self.conv_5(x)
         x = self.conv_6(x)
         #x = self.conv_7(x)
+        x = self.conv_8(x)
 
         return x
 
@@ -56,7 +63,7 @@ def test():
     import test_nn_Conv1d_pnnx
     b = test_nn_Conv1d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_nn_Conv2d.py

@@ -23,6 +23,12 @@ class Model(nn.Module):
         self.conv_6 = nn.Conv2d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, dilation=1, groups=1, bias=False, padding_mode='replicate')
         #self.conv_7 = nn.Conv2d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6), dilation=2, groups=1, bias=True, padding_mode='circular')
 
+        self.conv_8 = nn.Conv2d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6), dilation=2, groups=1, bias=True)
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.conv_8 = torch.nn.utils.weight_norm(self.conv_8)
+        else:
+            self.conv_8 = torch.nn.utils.parametrizations.weight_norm(self.conv_8)
+
     def forward(self, x):
         x = self.conv_0(x)
         x = self.conv_1(x)
@@ -32,6 +38,7 @@ class Model(nn.Module):
         x = self.conv_5(x)
         x = self.conv_6(x)
         #x = self.conv_7(x)
+        x = self.conv_8(x)
 
         return x
 
@@ -56,7 +63,7 @@ def test():
     import test_nn_Conv2d_pnnx
     b = test_nn_Conv2d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_nn_Conv3d.py

@@ -24,6 +24,12 @@ class Model(nn.Module):
             self.conv_6 = nn.Conv3d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, dilation=1, groups=1, bias=False, padding_mode='replicate')
             # self.conv_7 = nn.Conv3d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6), dilation=2, groups=1, bias=True, padding_mode='circular')
 
+        self.conv_8 = nn.Conv3d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6,7), dilation=2, groups=1, bias=True)
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.conv_8 = torch.nn.utils.weight_norm(self.conv_8)
+        else:
+            self.conv_8 = torch.nn.utils.parametrizations.weight_norm(self.conv_8)
+
     def forward(self, x):
         x = self.conv_0(x)
         x = self.conv_1(x)
@@ -36,6 +42,7 @@ class Model(nn.Module):
         x = self.conv_5(x)
         x = self.conv_6(x)
         #x = self.conv_7(x)
+        x = self.conv_8(x)
 
         return x
 
@@ -60,7 +67,7 @@ def test():
     import test_nn_Conv3d_pnnx
     b = test_nn_Conv3d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_nn_ConvTranspose1d.py

@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from packaging import version
 
 class Model(nn.Module):
     def __init__(self):
@@ -18,6 +19,11 @@ class Model(nn.Module):
         self.deconv_6 = nn.ConvTranspose1d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, output_padding=0, dilation=1, groups=1, bias=False)
         self.deconv_7 = nn.ConvTranspose1d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(6), output_padding=(1), dilation=2, groups=1, bias=True)
 
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.deconv_7 = torch.nn.utils.weight_norm(self.deconv_7)
+        else:
+            self.deconv_7 = torch.nn.utils.parametrizations.weight_norm(self.deconv_7)
+
         self.downsample = nn.Conv1d(24, 16, 3, stride=2, padding=1)
         self.upsample = nn.ConvTranspose1d(16, 24, 3, stride=2, padding=1)
 
@@ -57,7 +63,7 @@ def test():
     import test_nn_ConvTranspose1d_pnnx
     b = test_nn_ConvTranspose1d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_nn_ConvTranspose2d.py

@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from packaging import version
 
 class Model(nn.Module):
     def __init__(self):
@@ -18,6 +19,11 @@ class Model(nn.Module):
         self.deconv_6 = nn.ConvTranspose2d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, output_padding=0, dilation=1, groups=1, bias=False)
         self.deconv_7 = nn.ConvTranspose2d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6), output_padding=(1,0), dilation=2, groups=1, bias=True)
 
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.deconv_7 = torch.nn.utils.weight_norm(self.deconv_7)
+        else:
+            self.deconv_7 = torch.nn.utils.parametrizations.weight_norm(self.deconv_7)
+
         self.downsample = nn.Conv2d(24, 16, 3, stride=2, padding=1)
         self.upsample = nn.ConvTranspose2d(16, 24, 3, stride=2, padding=1)
 
@@ -57,7 +63,7 @@ def test():
     import test_nn_ConvTranspose2d_pnnx
     b = test_nn_ConvTranspose2d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_nn_ConvTranspose3d.py

@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from packaging import version
 
 class Model(nn.Module):
     def __init__(self):
@@ -18,6 +19,11 @@ class Model(nn.Module):
         self.deconv_6 = nn.ConvTranspose3d(in_channels=32, out_channels=28, kernel_size=2, stride=1, padding=2, output_padding=0, dilation=1, groups=1, bias=False)
         self.deconv_7 = nn.ConvTranspose3d(in_channels=28, out_channels=24, kernel_size=3, stride=2, padding=(5,6,7), output_padding=(1,0,1), dilation=2, groups=1, bias=True)
 
+        if version.parse(torch.__version__) < version.parse('2.1'):
+            self.deconv_7 = torch.nn.utils.weight_norm(self.deconv_7)
+        else:
+            self.deconv_7 = torch.nn.utils.parametrizations.weight_norm(self.deconv_7)
+
         self.downsample = nn.Conv3d(24, 16, 3, stride=2, padding=1)
         self.upsample = nn.ConvTranspose3d(16, 24, 3, stride=2, padding=1)
 
@@ -57,7 +63,7 @@ def test():
     import test_nn_ConvTranspose3d_pnnx
     b = test_nn_ConvTranspose3d_pnnx.test_inference()
 
-    return torch.equal(a, b)
+    return torch.allclose(a, b, 1e-3, 1e-3)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_nn_Linear.py

@@ -4,23 +4,36 @@
 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.linear_0 = nn.Linear(in_features=64, out_features=16, bias=False)
-        self.linear_1 = nn.Linear(in_features=16, out_features=3, bias=True)
+        self.linear_1 = nn.Linear(in_features=16, out_features=13, bias=True)
+
+        self.linear_2 = nn.Linear(in_features=13, out_features=17, bias=True)
+        if version.parse(torch.__version__) < version.parse('1.9'):
+            # weight_norm on torch 1.8 produces wrong output shape, skip it
+            pass
+        elif version.parse(torch.__version__) < version.parse('2.1'):
+            self.linear_2 = torch.nn.utils.weight_norm(self.linear_2)
+        else:
+            self.linear_2 = torch.nn.utils.parametrizations.weight_norm(self.linear_2)
 
     def forward(self, x, y, z):
         x = self.linear_0(x)
         x = self.linear_1(x)
+        x = self.linear_2(x)
 
         y = self.linear_0(y)
         y = self.linear_1(y)
+        y = self.linear_2(y)
 
         z = self.linear_0(z)
         z = self.linear_1(z)
+        z = self.linear_2(z)
         return x, y, z
 
 def test():
@@ -32,7 +45,7 @@ def test():
     y = torch.rand(12, 64)
     z = torch.rand(1, 3, 12, 64)
 
-    a0, a1, a2 = net(x, y, z)
+    a = net(x, y, z)
 
     # export torchscript
     mod = torch.jit.trace(net, (x, y, z))
@@ -44,9 +57,13 @@ def test():
 
     # pnnx inference
     import test_nn_Linear_pnnx
-    b0, b1, b2 = test_nn_Linear_pnnx.test_inference()
+    b = test_nn_Linear_pnnx.test_inference()
 
-    return torch.equal(a0, b0) and torch.equal(a1, b1) and torch.equal(a2, b2)
+    for a0, b0 in zip(a, b):
+        b0 = b0.reshape_as(a0)
+        if not torch.allclose(a0, b0, 1e-3, 1e-3):
+            return False
+    return True
 
 if __name__ == "__main__":
     if test():