add diag layer and its converter

src/CMakeLists.txt

@@ -71,6 +71,7 @@ ncnn_add_layer(Concat)
 ncnn_add_layer(Convolution)
 ncnn_add_layer(Crop)
 ncnn_add_layer(Deconvolution)
+ncnn_add_layer(Diag)
 ncnn_add_layer(Dropout)
 ncnn_add_layer(Eltwise)
 ncnn_add_layer(ELU)

src/layer/diag.cpp

@@ -0,0 +1,82 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2023 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 "diag.h"
+
+namespace ncnn {
+
+Diag::Diag()
+{
+ one_blob_only = true;
+ support_inplace = false;
+}
+
+int Diag::load_param(const ParamDict& pd)
+{
+ diagonal = pd.get(0, 0);
+
+ return 0;
+}
+
+int Diag::forward(const Mat& bottom_blob, Mat& top_blob, const Option& opt) const
+{
+ int dims = bottom_blob.dims;
+ size_t elemsize = bottom_blob.elemsize;
+
+ if (dims == 1)
+ {
+ int w = bottom_blob.w;
+ int top_w = w + std::abs(diagonal);
+ int stride = top_w + 1;
+
+ top_blob.create(top_w, top_w, elemsize, opt.blob_allocator);
+ if (top_blob.empty())
+ return -100;
+
+ top_blob.fill(0.0f);
+
+ int bias_r = -std::min(diagonal, 0);
+ int bias_c = std::max(diagonal, 0);
+
+ for (int i = 0; i < w; i++)
+ {
+ top_blob.row(i + bias_r)[i + bias_c] = bottom_blob[i];
+ }
+ }
+ if (dims == 2)
+ {
+ int w = bottom_blob.w;
+ int h = bottom_blob.h;
+ float tmp = (w - h) / 2.0;
+
+ int len = std::min(w, h) - (int)std::max(std::abs(diagonal - tmp) - std::abs(tmp), 0.0f);
+ len = std::max(len, 0);
+
+ top_blob.create(len, elemsize, opt.blob_allocator);
+ if (top_blob.empty())
+ return -100;
+
+ int bias_r = -std::min(diagonal, 0);
+ int bias_c = std::max(diagonal, 0);
+
+ for (int i = 0; i < len; i++)
+ {
+ top_blob[i] = bottom_blob.row(i + bias_r)[i + bias_c];
+ }
+ }
+
+ return 0;
+}
+
+} // namespace ncnn

src/layer/diag.h

@@ -0,0 +1,37 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2023 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.
+
+#ifndef LAYER_DIAG_H
+#define LAYER_DIAG_H
+
+#include "layer.h"
+
+namespace ncnn {
+
+class Diag : public Layer
+{
+public:
+ Diag();
+
+ virtual int load_param(const ParamDict& pd);
+
+ virtual int forward(const Mat& bottom_blob, Mat& top_blob, const Option& opt) const;
+
+public:
+ int diagonal;
+};
+
+} // namespace ncnn
+
+#endif // LAYER_DIAG_H

tests/CMakeLists.txt

@@ -92,6 +92,7 @@ ncnn_add_layer_test(DeconvolutionDepthWise3D)
 ncnn_add_layer_test(DeepCopy)
 ncnn_add_layer_test(DeformableConv2D)
 ncnn_add_layer_test(Dequantize)
+ncnn_add_layer_test(Diag)
 ncnn_add_layer_test(Dropout)
 ncnn_add_layer_test(Einsum)
 ncnn_add_layer_test(Eltwise)

tests/test_diag.cpp

@@ -0,0 +1,57 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2023 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 "layer/diag.h"
+#include "testutil.h"
+
+static int test_diag(const ncnn::Mat& a, int diagonal)
+{
+ ncnn::ParamDict pd;
+ pd.set(0, diagonal);
+
+ std::vector<ncnn::Mat> weights(0);
+
+ int ret = test_layer<ncnn::Diag>("Diag", pd, weights, a);
+ if (ret != 0)
+ {
+ fprintf(stderr, "test_diag failed a.dims=%d a=(%d %d %d %d)\n", a.dims, a.w, a.h, a.d, a.c);
+ }
+
+ return ret;
+}
+
+static int test_diag_0()
+{
+ return 0
+ || test_diag(RandomMat(5, 24), 3)
+ || test_diag(RandomMat(7, 12), 0)
+ || test_diag(RandomMat(3, 4), -6);
+}
+
+static int test_diag_1()
+{
+ return 0
+ || test_diag(RandomMat(5), -1)
+ || test_diag(RandomMat(7), 0)
+ || test_diag(RandomMat(3), 2);
+}
+
+int main()
+{
+ SRAND(7767517);
+
+ return 0
+ || test_diag_0()
+ || test_diag_1();
+}

tools/pnnx/src/CMakeLists.txt

@@ -216,6 +216,7 @@ set(pnnx_pass_level2_SRCS
  pass_level2/torch_cross.cpp
  pass_level2/torch_cumsum.cpp
  pass_level2/torch_dequantize.cpp
+ pass_level2/torch_diag.cpp
  pass_level2/torch_einsum.cpp
  pass_level2/torch_empty.cpp
  pass_level2/torch_empty_like.cpp

tools/pnnx/src/pass_level2/torch_diag.cpp

@@ -0,0 +1,41 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2023 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_diag : 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 diagonal
+aten::diag op_0 2 1 input diagonal out
+pnnx.Output output 1 0 out
+)PNNXIR";
+ }
+
+ const char* type_str() const
+ {
+ return "torch.diag";
+ }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_diag, 20)
+
+} // namespace pnnx

tools/pnnx/src/pass_ncnn/torch_diag.cpp

@@ -0,0 +1,63 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2023 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_ncnn.h"
+
+namespace pnnx {
+
+namespace ncnn {
+
+class torch_diag : public GraphRewriterPass
+{
+public:
+ const char* match_pattern_graph() const
+ {
+ return R"PNNXIR(7767517
+3 2
+pnnx.Input input 0 1 input
+torch.diag op_0 1 1 input out diagonal=%diagonal
+pnnx.Output output 1 0 out
+)PNNXIR";
+ }
+
+ const char* type_str() const
+ {
+ return "Diag";
+ }
+
+ const char* name_str() const
+ {
+ return "diag";
+ }
+
+ void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
+ {
+ int diagonal = captured_params.at("diagonal").i;
+ int input_rank = op->inputs[0]->shape.size();
+
+ if (input_rank > 2)
+ {
+ fprintf(stderr, "diag %d-rank tensor is not supported yet!\n", input_rank);
+ return;
+ }
+
+ op->params["0"] = diagonal;
+ }
+};
+
+REGISTER_GLOBAL_PNNX_NCNN_GRAPH_REWRITER_PASS(torch_diag, 20)
+
+} // namespace ncnn
+
+} // namespace pnnx

tools/pnnx/tests/CMakeLists.txt

@@ -197,6 +197,7 @@ pnnx_add_test(torch_cross)
 pnnx_add_test(torch_cumsum)
 pnnx_add_test(torch_einsum)
 pnnx_add_test(torch_eq)
+pnnx_add_test(torch_diag)
 pnnx_add_test(torch_flatten)
 pnnx_add_test(torch_full)
 pnnx_add_test(torch_full_like)

tools/pnnx/tests/test_torch_diag.py

@@ -0,0 +1,60 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2023 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.diag(x, -1)
+ y = torch.diag(y)
+ z = torch.diag(z, 3)
+ return x, y, z
+
+def test():
+ net = Model()
+ net.eval()
+
+ torch.manual_seed(0)
+ x = torch.rand(7)
+ y = torch.rand(5, 5)
+ z = torch.rand(4, 8)
+ a = net(x, y, z)
+
+ # export torchscript
+ mod = torch.jit.trace(net, (x, y, z))
+ mod.save("test_torch_diag.pt")
+
+ # torchscript to pnnx
+ import os
+ os.system("../src/pnnx test_torch_diag.pt inputshape=[7],[5,5],[4,8]")
+
+ # pnnx inference
+ import test_torch_diag_pnnx
+ b = test_torch_diag_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)