diff --git a/tools/pnnx/src/CMakeLists.txt b/tools/pnnx/src/CMakeLists.txt index f8cbfdc72..5645b6384 100644 --- a/tools/pnnx/src/CMakeLists.txt +++ b/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 diff --git a/tools/pnnx/src/pass_level2/torch_argmax.cpp b/tools/pnnx/src/pass_level2/torch_argmax.cpp new file mode 100644 index 000000000..c7ea31044 --- /dev/null +++ b/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 diff --git a/tools/pnnx/src/pass_level2/torch_argmin.cpp b/tools/pnnx/src/pass_level2/torch_argmin.cpp new file mode 100644 index 000000000..d20779f04 --- /dev/null +++ b/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 diff --git a/tools/pnnx/src/pass_level5/fuse_conv2d_batchnorm2d.cpp b/tools/pnnx/src/pass_level5/fuse_conv2d_batchnorm2d.cpp index 049b03f36..346e1b214 100644 --- a/tools/pnnx/src/pass_level5/fuse_conv2d_batchnorm2d.cpp +++ b/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"); diff --git a/tools/pnnx/tests/CMakeLists.txt b/tools/pnnx/tests/CMakeLists.txt index 78a9c5b96..d98c4c4c3 100644 --- a/tools/pnnx/tests/CMakeLists.txt +++ b/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) diff --git a/tools/pnnx/tests/ncnn/test_Tensor_reshape.py b/tools/pnnx/tests/ncnn/test_Tensor_reshape.py index 19fafb87d..af5eca8f9 100644 --- a/tools/pnnx/tests/ncnn/test_Tensor_reshape.py +++ b/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(): diff --git a/tools/pnnx/tests/ncnn/test_Tensor_view.py b/tools/pnnx/tests/ncnn/test_Tensor_view.py index 048ea8ee9..7bf967621 100644 --- a/tools/pnnx/tests/ncnn/test_Tensor_view.py +++ b/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(): diff --git a/tools/pnnx/tests/ncnn/test_torch_permute.py b/tools/pnnx/tests/ncnn/test_torch_permute.py index ce6040797..bf6d49778 100644 --- a/tools/pnnx/tests/ncnn/test_torch_permute.py +++ b/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(): diff --git a/tools/pnnx/tests/ncnn/test_torch_transpose.py b/tools/pnnx/tests/ncnn/test_torch_transpose.py index 365fa6726..cfd7204d8 100644 --- a/tools/pnnx/tests/ncnn/test_torch_transpose.py +++ b/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(): diff --git a/tools/pnnx/tests/test_torch_argmax.py b/tools/pnnx/tests/test_torch_argmax.py new file mode 100644 index 000000000..68f13e46b --- /dev/null +++ b/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) diff --git a/tools/pnnx/tests/test_torch_argmin.py b/tools/pnnx/tests/test_torch_argmin.py new file mode 100644 index 000000000..d080b3ccf --- /dev/null +++ b/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)