convert pnnx argmax argmin, fix #3435

tools/pnnx/src/CMakeLists.txt

@@ -169,6 +169,8 @@ set(pnnx_pass_level2_SRCS
     pass_level2/Tensor_select.cpp
     pass_level2/Tensor_slice.cpp
     pass_level2/Tensor_view.cpp
+    pass_level2/torch_argmax.cpp
+    pass_level2/torch_argmin.cpp
     pass_level2/torch_cat.cpp
     pass_level2/torch_chunk.cpp
     pass_level2/torch_clamp.cpp

tools/pnnx/src/pass_level2/torch_argmax.cpp

@@ -0,0 +1,42 @@
+// 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 torch_argmax : 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 dim
+pnnx.Input              input_2     0 1 keepdim
+aten::argmax            op_0        3 1 input dim keepdim out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.argmax";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_argmax, 20)
+
+} // namespace pnnx

tools/pnnx/src/pass_level2/torch_argmin.cpp

@@ -0,0 +1,42 @@
+// 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 torch_argmin : 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 dim
+pnnx.Input              input_2     0 1 keepdim
+aten::argmin            op_0        3 1 input dim keepdim out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.argmin";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_argmin, 20)
+
+} // namespace pnnx

tools/pnnx/src/pass_level5/fuse_conv2d_batchnorm2d.cpp

@@ -29,7 +29,7 @@ public:
         return R"PNNXIR(7767517
 4 3
 pnnx.Input              input       0 1 input
-nn.Conv2d               op_0        1 1 input a in_channels=%in_channels out_channels=%out_channels kernel_size=%kernel_size stride=%stride padding=%padding dilation=%dilation groups=%groups bias=%bias @weight @bias
+nn.Conv2d               op_0        1 1 input a in_channels=%in_channels out_channels=%out_channels kernel_size=%kernel_size stride=%stride padding_mode=%padding_mode padding=%padding dilation=%dilation groups=%groups bias=%bias @weight @bias
 nn.BatchNorm2d          op_1        1 1 a out num_features=%num_features eps=%eps affine=%affine @running_mean @running_var @weight @bias
 pnnx.Output             output      1 0 out
 )PNNXIR";
@@ -50,6 +50,7 @@ pnnx.Output             output      1 0 out
         op->params["in_channels"] = captured_params.at("in_channels");
         op->params["out_channels"] = captured_params.at("out_channels");
         op->params["kernel_size"] = captured_params.at("kernel_size");
+        op->params["padding_mode"] = captured_params.at("padding_mode");
         op->params["stride"] = captured_params.at("stride");
         op->params["padding"] = captured_params.at("padding");
         op->params["dilation"] = captured_params.at("dilation");

tools/pnnx/tests/CMakeLists.txt

@@ -161,6 +161,8 @@ pnnx_add_test(Tensor_select)
 pnnx_add_test(Tensor_slice)
 pnnx_add_test(Tensor_view)
 
+pnnx_add_test(torch_argmax)
+pnnx_add_test(torch_argmin)
 pnnx_add_test(torch_cat)
 pnnx_add_test(torch_chunk)
 pnnx_add_test(torch_clamp)

tools/pnnx/tests/ncnn/test_Tensor_reshape.py

@@ -27,6 +27,9 @@ class Model(nn.Module):
         y = y.reshape(99, 5)
         z = z.reshape(4, 3, 6, 10)
         z = z.reshape(15, 6, 8)
+        x = F.relu(x)
+        y = F.relu(y)
+        z = F.relu(z)
         return x, y, z
 
 def test():

tools/pnnx/tests/ncnn/test_Tensor_view.py

@@ -27,6 +27,9 @@ class Model(nn.Module):
         y = y.reshape(99, 5)
         z = z.reshape(4, 3, 6, 10)
         z = z.reshape(15, 6, 8)
+        x = F.relu(x)
+        y = F.relu(y)
+        z = F.relu(z)
         return x, y, z
 
 def test():

tools/pnnx/tests/ncnn/test_torch_permute.py

@@ -35,6 +35,9 @@ class Model(nn.Module):
             y = torch.permute(y, (1, 0, 2))
             z = torch.permute(z, (1, 3, 0, 2))
             z = torch.permute(z, (2, 0, 3, 1))
+        x = F.relu(x)
+        y = F.relu(y)
+        z = F.relu(z)
         return x, y, z
 
 def test():

tools/pnnx/tests/ncnn/test_torch_transpose.py

@@ -24,6 +24,9 @@ class Model(nn.Module):
         x = torch.transpose(x, 0, 1)
         y = torch.transpose(y, 1, 2)
         z = torch.transpose(z, 0, 2)
+        x = F.relu(x)
+        y = F.relu(y)
+        z = F.relu(z)
         return x, y, z
 
 def test():

tools/pnnx/tests/test_torch_argmax.py

@@ -0,0 +1,61 @@
+# 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__()
+
+    def forward(self, x, y, z):
+        x = torch.argmax(x)
+        y = torch.argmax(y, dim=1)
+        z = torch.argmax(z, dim=2, keepdim=True)
+        return x, y, z
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 3, 16)
+    y = torch.rand(1, 5, 9, 11)
+    z = torch.rand(14, 8, 5, 9, 10)
+
+    a = net(x, y, z)
+
+    # export torchscript
+    mod = torch.jit.trace(net, (x, y, z))
+    mod.save("test_torch_argmax.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../src/pnnx test_torch_argmax.pt inputshape=[1,3,16],[1,5,9,11],[14,8,5,9,10]")
+
+    # pnnx inference
+    import test_torch_argmax_pnnx
+    b = test_torch_argmax_pnnx.test_inference()
+
+    for a0, b0 in zip(a, b):
+        if not torch.equal(a0, b0):
+            return False
+    return True
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/test_torch_argmin.py

@@ -0,0 +1,61 @@
+# 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__()
+
+    def forward(self, x, y, z):
+        x = torch.argmin(x)
+        y = torch.argmin(y, dim=1)
+        z = torch.argmin(z, dim=2, keepdim=True)
+        return x, y, z
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 3, 16)
+    y = torch.rand(1, 5, 9, 11)
+    z = torch.rand(14, 8, 5, 9, 10)
+
+    a = net(x, y, z)
+
+    # export torchscript
+    mod = torch.jit.trace(net, (x, y, z))
+    mod.save("test_torch_argmin.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../src/pnnx test_torch_argmin.pt inputshape=[1,3,16],[1,5,9,11],[14,8,5,9,10]")
+
+    # pnnx inference
+    import test_torch_argmin_pnnx
+    b = test_torch_argmin_pnnx.test_inference()
+
+    for a0, b0 in zip(a, b):
+        if not torch.equal(a0, b0):
+            return False
+    return True
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)