u

tools/pnnx/src/pass_level1/nn_Conv1d.cpp

@@ -122,9 +122,8 @@ public:
             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;
-            int maxk = op->params.at("kernel_size").ai[0];
-            apply_weight_norm(weight_data, weight_g_data, outch, inch, maxk);
+            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

tools/pnnx/src/pass_level1/nn_Conv2d.cpp

@@ -122,9 +122,8 @@ public:
             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;
-            int maxk = op->params.at("kernel_size").ai[0] * op->params.at("kernel_size").ai[1];
-            apply_weight_norm(weight_data, weight_g_data, outch, inch, maxk);
+            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

tools/pnnx/src/pass_level1/nn_Conv3d.cpp

@@ -122,9 +122,8 @@ public:
             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;
-            int maxk = 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, maxk);
+            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

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/utils.cpp

@@ -112,23 +112,23 @@ 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 outch, int inch, int maxk)
+void apply_weight_norm(std::vector<float>& weight, const std::vector<float>& weight_g, int outch, int inch)
 {
     const float eps = 1e-12f;
 
     for (int i = 0; i < outch; i++)
     {
-        float* pw = weight.data() + i * inch * maxk;
+        float* pw = weight.data() + i * inch;
 
         float norm = 0.f;
-        for (int j = 0; j < inch * maxk; j++)
+        for (int j = 0; j < inch; j++)
         {
             float w = pw[j];
             norm += w * w;
         }
         norm = sqrt(norm) + eps;
 
-        for (int j = 0; j < inch * maxk; j++)
+        for (int j = 0; j < inch; j++)
         {
             pw[j] = weight_g[i] * pw[j] / norm;
         }

tools/pnnx/src/utils.h

@@ -12,7 +12,7 @@ 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 outch, int inch, int maxk);
+void apply_weight_norm(std::vector<float>& weight, const std::vector<float>& weight_g, int outch, int inch);
 
 } // namespace pnnx
 

tools/pnnx/tests/test_nn_ConvTranspose1d.py

@@ -18,6 +18,8 @@ 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)
 
+        self.deconv_7 = torch.nn.utils.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 +59,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-4, 1e-4)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_nn_ConvTranspose2d.py

@@ -18,6 +18,8 @@ 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)
 
+        self.deconv_7 = torch.nn.utils.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 +59,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-4, 1e-4)
 
 if __name__ == "__main__":
     if test():

tools/pnnx/tests/test_nn_ConvTranspose3d.py

@@ -18,6 +18,8 @@ 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)
 
+        self.deconv_7 = torch.nn.utils.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 +59,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-4, 1e-4)
 
 if __name__ == "__main__":
     if test():