From 5e9bf90455b02a25e7991dab4ab8b3035a4bebb3 Mon Sep 17 00:00:00 2001 From: nihui Date: Fri, 31 May 2024 16:06:59 +0800 Subject: [PATCH] pnnx onnx with inputshape (#5480) --- tools/pnnx/src/CMakeLists.txt | 2 + tools/pnnx/src/ir.cpp | 17 +- tools/pnnx/src/load_onnx.cpp | 167 ++++++++++- tools/pnnx/src/load_onnx.h | 6 +- tools/pnnx/src/main.cpp | 4 +- tools/pnnx/src/pass_level2/F_avg_pool2d.cpp | 20 +- tools/pnnx/src/pass_level2/F_batch_norm.cpp | 25 ++ tools/pnnx/src/pass_level2/F_conv2d.cpp | 28 ++ tools/pnnx/src/pass_level2/F_gelu.cpp | 2 +- tools/pnnx/src/pass_level2/F_hardswish.cpp | 28 ++ tools/pnnx/src/pass_level2/F_interpolate.cpp | 114 ++++++++ tools/pnnx/src/pass_level2/F_linear.cpp | 62 +++- tools/pnnx/src/pass_level2/F_pad.cpp | 22 ++ tools/pnnx/src/pass_level2/F_softmax.cpp | 21 ++ tools/pnnx/src/pass_level2/Tensor_reshape.cpp | 47 ++- tools/pnnx/src/pass_level2/Tensor_select.cpp | 24 +- tools/pnnx/src/pass_level2/Tensor_to.cpp | 43 +++ tools/pnnx/src/pass_level2/torch_clamp.cpp | 27 +- tools/pnnx/src/pass_level2/torch_split.cpp | 17 ++ tools/pnnx/src/pass_level2/torch_squeeze.cpp | 38 +++ .../pnnx/src/pass_level2/torch_unsqueeze.cpp | 38 +++ .../pnnx/src/pass_level3/fuse_expression.cpp | 9 +- tools/pnnx/src/pass_level5.cpp | 6 + .../pnnx/src/pass_level5/eval_expression.cpp | 55 ++-- tools/pnnx/src/pass_level5/fuse_layernorm.cpp | 100 +++++++ .../pass_level5/fuse_multiheadattention.cpp | 272 ++++++++++++++++-- .../fuse_slice_squeeze_to_select.cpp | 86 ++++++ .../fuse_slice_squeeze_to_select.h | 21 ++ .../fuse_slice_to_tensor_split.cpp | 36 +-- .../src/pass_level5/fuse_static_groupnorm.cpp | 2 +- .../pass_level5/fuse_static_instancenorm.cpp | 6 +- .../src/pass_level5/fuse_static_layernorm.cpp | 2 +- tools/pnnx/src/pass_level5/fuse_swish.cpp | 71 +++++ tools/pnnx/src/pass_level5/fuse_swish.h | 21 ++ tools/pnnx/src/pass_onnx.cpp | 79 ++++- tools/pnnx/src/pass_onnx/eliminate_noop.cpp | 2 +- tools/pnnx/src/pass_onnx/fold_constants.cpp | 52 +++- tools/pnnx/src/pass_onnx/fold_constants.h | 6 +- tools/pnnx/src/pass_onnx/shape_inference.cpp | 272 +++++++++++++++++- tools/pnnx/src/pass_onnx/shape_inference.h | 6 +- tools/pnnx/tests/onnx/CMakeLists.txt | 2 + tools/pnnx/tests/onnx/test_swin_t.py | 65 +++++ tools/pnnx/tests/onnx/test_vit_b_32.py | 68 +++++ 43 files changed, 1844 insertions(+), 147 deletions(-) create mode 100644 tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.cpp create mode 100644 tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.h create mode 100644 tools/pnnx/src/pass_level5/fuse_swish.cpp create mode 100644 tools/pnnx/src/pass_level5/fuse_swish.h create mode 100644 tools/pnnx/tests/onnx/test_swin_t.py create mode 100644 tools/pnnx/tests/onnx/test_vit_b_32.py diff --git a/tools/pnnx/src/CMakeLists.txt b/tools/pnnx/src/CMakeLists.txt index 7134179f0..c87794a0f 100644 --- a/tools/pnnx/src/CMakeLists.txt +++ b/tools/pnnx/src/CMakeLists.txt @@ -363,6 +363,7 @@ set(pnnx_pass_level5_SRCS pass_level5/fuse_slice_copy.cpp pass_level5/fuse_slice_indices.cpp pass_level5/fuse_slice_to_tensor_split.cpp + pass_level5/fuse_slice_squeeze_to_select.cpp pass_level5/fuse_static_batchnorm.cpp pass_level5/fuse_static_conv.cpp pass_level5/fuse_static_convtranspose.cpp @@ -371,6 +372,7 @@ set(pnnx_pass_level5_SRCS pass_level5/fuse_static_instancenorm.cpp pass_level5/fuse_static_layernorm.cpp pass_level5/fuse_static_linear.cpp + pass_level5/fuse_swish.cpp pass_level5/normalize_einsum_equation.cpp pass_level5/unroll_rnn_op.cpp ) diff --git a/tools/pnnx/src/ir.cpp b/tools/pnnx/src/ir.cpp index a1fbd4932..fe5424e14 100644 --- a/tools/pnnx/src/ir.cpp +++ b/tools/pnnx/src/ir.cpp @@ -1004,11 +1004,20 @@ static std::string expand_expression(const Operator* op) { std::string a = exprstack.top(); exprstack.pop(); - std::string b = exprstack.top(); - exprstack.pop(); - std::string r = a + ".size(" + b + ")"; - exprstack.push(r); + if (exprstack.empty()) + { + std::string r = a + ".shape"; + exprstack.push(r); + } + else + { + std::string b = exprstack.top(); + exprstack.pop(); + + std::string r = a + ".size(" + b + ")"; + exprstack.push(r); + } } else if (t == "int" || t == "abs" diff --git a/tools/pnnx/src/load_onnx.cpp b/tools/pnnx/src/load_onnx.cpp index 4a7199b88..6543ad343 100644 --- a/tools/pnnx/src/load_onnx.cpp +++ b/tools/pnnx/src/load_onnx.cpp @@ -23,6 +23,8 @@ #include +#include + #include "ir.h" #include "pass_onnx/canonicalize.h" @@ -304,7 +306,14 @@ Operand* Graph::new_operand(const onnx::ValueInfoProto& value) r->shape.resize(tensor_shape.dim_size()); for (int z = 0; z < tensor_shape.dim_size(); z++) { - r->shape[z] = tensor_shape.dim(z).dim_value(); + if (!tensor_shape.dim(z).has_dim_value()) + { + r->shape[z] = -1; + } + else + { + r->shape[z] = tensor_shape.dim(z).dim_value(); + } } operands.push_back(r); @@ -366,7 +375,147 @@ static double get_current_time() return usec.count() / 1000.0; } -int load_onnx(const std::string& onnxpath, Graph& pnnx_graph) +static const char* get_onnx_tensor_elem_data_type_str(ONNXTensorElementDataType type) +{ + switch (type) + { + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: + return "i8"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: + return "u8"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16: + return "i16"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16: + return "u16"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: + return "i32"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32: + return "u32"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64: + return "i64"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64: + return "u64"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: + return "f16"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: + return "f32"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE: + return "f64"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: + return "bf16"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64: + return "c64"; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128: + return "c128"; + default: + break; + } + + // unknown + fprintf(stderr, "unsupported tensor elem data type %d\n", (int)type); + return ""; +} + +static bool check_input_shape(onnx::ModelProto& model, const std::vector >& input_shapes, const std::vector& input_types) +{ + const onnx::GraphProto& graph = model.graph(); + + if (!input_shapes.empty() && (int)input_shapes.size() != graph.input_size()) + { + fprintf(stderr, "input_shape expect %d tensors but got %d\n", graph.input_size(), (int)input_shapes.size()); + return false; + } + + for (int i = 0; i < graph.input_size(); i++) + { + const onnx::ValueInfoProto& value = graph.input(i); + const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape(); + + ONNXTensorElementDataType datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type(); + + bool matched = true; + + if (input_shapes.empty()) + { + // dynamic dimension size + for (int j = 0; j < tsp.dim_size(); j++) + { + if (!tsp.dim(j).has_dim_value()) + { + matched = false; + break; + } + } + } + else + { + if ((int)input_shapes[i].size() != tsp.dim_size()) + { + matched = false; + } + else + { + for (int j = 0; j < tsp.dim_size(); j++) + { + if (!tsp.dim(j).has_dim_value()) + continue; + + int64_t ds = tsp.dim(j).dim_value(); + if (input_shapes[i][j] != ds) + matched = false; + } + } + + if (input_types[i] != get_onnx_tensor_elem_data_type_str(datatype)) + matched = false; + } + + if (!matched) + { + fprintf(stderr, "input_shapes[%d] expect [", i); + for (int j = 0; j < tsp.dim_size(); j++) + { + if (tsp.dim(j).has_dim_value()) + { + int64_t ds = tsp.dim(j).dim_value(); + fprintf(stderr, "%ld", ds); + } + else + { + fprintf(stderr, "?"); + } + if (j + 1 != tsp.dim_size()) + fprintf(stderr, ","); + } + fprintf(stderr, "]%s but got ", get_onnx_tensor_elem_data_type_str(datatype)); + if (input_shapes.empty()) + { + fprintf(stderr, "nothing\n"); + } + else + { + fprintf(stderr, "["); + for (size_t j = 0; j < input_shapes[i].size(); j++) + { + fprintf(stderr, "%ld", input_shapes[i][j]); + if (j + 1 != input_shapes[i].size()) + fprintf(stderr, ","); + } + fprintf(stderr, "]%s\n", input_types[i].c_str()); + } + + return false; + } + } + + return true; +} + +int load_onnx(const std::string& onnxpath, Graph& pnnx_graph, + const std::vector >& input_shapes, + const std::vector& input_types, + const std::vector >& input_shapes2, + const std::vector& input_types2) { onnx::ModelProto model; @@ -377,6 +526,16 @@ int load_onnx(const std::string& onnxpath, Graph& pnnx_graph) return -1; } + // input shape sanity check + if (!check_input_shape(model, input_shapes, input_types)) + { + return -1; + } + if (!input_shapes2.empty() && !check_input_shape(model, input_shapes2, input_types2)) + { + return -1; + } + fprintf(stderr, "############# pass_level0 onnx \n"); onnx2pnnx::ModelStat oldstat = onnx2pnnx::get_model_stat(model); @@ -415,7 +574,7 @@ int load_onnx(const std::string& onnxpath, Graph& pnnx_graph) t0 = get_current_time(); - onnx2pnnx::fold_constants(model); + onnx2pnnx::fold_constants(model, input_shapes, input_types, input_shapes2, input_types2); t1 = get_current_time(); @@ -445,7 +604,7 @@ int load_onnx(const std::string& onnxpath, Graph& pnnx_graph) t0 = get_current_time(); - onnx2pnnx::shape_inference(model); + onnx2pnnx::shape_inference(model, input_shapes, input_types, input_shapes2, input_types2); t1 = get_current_time(); diff --git a/tools/pnnx/src/load_onnx.h b/tools/pnnx/src/load_onnx.h index e1b99757e..2fdf6acd9 100644 --- a/tools/pnnx/src/load_onnx.h +++ b/tools/pnnx/src/load_onnx.h @@ -19,7 +19,11 @@ namespace pnnx { -int load_onnx(const std::string& onnxpath, Graph& g); +int load_onnx(const std::string& onnxpath, Graph& g, + const std::vector >& input_shapes, + const std::vector& input_types, + const std::vector >& input_shapes2, + const std::vector& input_types2); } // namespace pnnx diff --git a/tools/pnnx/src/main.cpp b/tools/pnnx/src/main.cpp index 93e7bafd2..c25128032 100644 --- a/tools/pnnx/src/main.cpp +++ b/tools/pnnx/src/main.cpp @@ -316,7 +316,9 @@ int main(int argc, char** argv) #if BUILD_ONNX2PNNX if (!model_file_maybe_torchscript(ptpath)) { - load_onnx(ptpath.c_str(), pnnx_graph); + load_onnx(ptpath.c_str(), pnnx_graph, + input_shapes, input_types, + input_shapes2, input_types2); } else #endif diff --git a/tools/pnnx/src/pass_level2/F_avg_pool2d.cpp b/tools/pnnx/src/pass_level2/F_avg_pool2d.cpp index 82b73c9c6..d83df9183 100644 --- a/tools/pnnx/src/pass_level2/F_avg_pool2d.cpp +++ b/tools/pnnx/src/pass_level2/F_avg_pool2d.cpp @@ -43,7 +43,7 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_avg_pool2d, 10) -class F_avg_pool2d_1 : public GraphRewriterPass +class F_avg_pool2d_onnx : public GraphRewriterPass { public: const char* match_pattern_graph() const @@ -83,6 +83,22 @@ pnnx.Output output 1 0 out } }; -REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_avg_pool2d_1, 10) +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_avg_pool2d_onnx, 10) + +class F_avg_pool2d_onnx_1 : public F_avg_pool2d_onnx +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +3 2 +pnnx.Input input_0 0 1 input +AveragePool op_0 1 1 input out kernel_shape=%kernel_shape strides=%strides pads=%pads ceil_mode=%ceil_mode count_include_pad=%count_include_pad +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_avg_pool2d_onnx_1, 10) } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/F_batch_norm.cpp b/tools/pnnx/src/pass_level2/F_batch_norm.cpp index 4f3cbae3b..83efcdc4e 100644 --- a/tools/pnnx/src/pass_level2/F_batch_norm.cpp +++ b/tools/pnnx/src/pass_level2/F_batch_norm.cpp @@ -104,4 +104,29 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_batch_norm_onnx, 10) +class F_batch_norm_onnx_1 : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +7 6 +pnnx.Input input_0 0 1 input +pnnx.Input input_1 0 1 weight +pnnx.Input input_2 0 1 bias +pnnx.Input input_3 0 1 running_mean +pnnx.Input input_4 0 1 running_var +BatchNormalization op_0 5 1 input weight bias running_mean running_var out epsilon=%eps momentum=* +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "F.batch_norm"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_batch_norm_onnx_1, 10) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/F_conv2d.cpp b/tools/pnnx/src/pass_level2/F_conv2d.cpp index 73d788d2b..5ba27f4bb 100644 --- a/tools/pnnx/src/pass_level2/F_conv2d.cpp +++ b/tools/pnnx/src/pass_level2/F_conv2d.cpp @@ -227,6 +227,34 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_conv2d_onnx, 10) +class F_conv2d_onnx_0 : public F_conv2d_onnx +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input_0 0 1 input +pnnx.Input input_1 0 1 weight +Conv op_0 2 1 input weight out kernel_shape=%kernel_shape strides=%strides pads=%pads dilations=%dilations group=%group +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* replace_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input_0 0 1 input +pnnx.Input input_1 0 1 weight +F.conv2d conv 2 1 input weight out bias=None +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_conv2d_onnx_0, 10) + class F_conv2d_onnx_1 : public GraphRewriterPass { public: diff --git a/tools/pnnx/src/pass_level2/F_gelu.cpp b/tools/pnnx/src/pass_level2/F_gelu.cpp index eb14032a8..7718dc658 100644 --- a/tools/pnnx/src/pass_level2/F_gelu.cpp +++ b/tools/pnnx/src/pass_level2/F_gelu.cpp @@ -214,7 +214,7 @@ public: pnnx.Input input 0 1 input prim::Constant op_0 0 1 sqrt2 value=%sqrt2 aten::div op_1 2 1 input sqrt2 16 -Erf op_2 1 1 16 17 +aten::erf op_2 1 1 16 17 prim::Constant op_3 0 1 one value=%1 aten::add op_4 2 1 17 one 22 aten::mul op_5 2 1 input 22 25 diff --git a/tools/pnnx/src/pass_level2/F_hardswish.cpp b/tools/pnnx/src/pass_level2/F_hardswish.cpp index 96b9d265d..e48a7e5b1 100644 --- a/tools/pnnx/src/pass_level2/F_hardswish.cpp +++ b/tools/pnnx/src/pass_level2/F_hardswish.cpp @@ -200,6 +200,34 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_hardswish_6, 8) +class F_hardswish_7 : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +10 9 +pnnx.Input input 0 1 input +prim::Constant op_0 0 1 25 value=3 +aten::add op_1 2 1 input 25 a +prim::Constant op_2 0 1 min value=0 +prim::Constant op_3 0 1 max value=6 +aten::clamp op_4 3 1 a min max b +aten::mul op_5 2 1 input b c +prim::Constant op_6 0 1 49 value=6 +aten::div op_7 2 1 c 49 out +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "F.hardswish"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_hardswish_7, 8) + class F_hardswish_onnx : public GraphRewriterPass { public: diff --git a/tools/pnnx/src/pass_level2/F_interpolate.cpp b/tools/pnnx/src/pass_level2/F_interpolate.cpp index 217d46913..41d33ee49 100644 --- a/tools/pnnx/src/pass_level2/F_interpolate.cpp +++ b/tools/pnnx/src/pass_level2/F_interpolate.cpp @@ -997,4 +997,118 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_interpolate_7, 10) +class F_interpolate_onnx : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input 0 1 input +pnnx.Attribute op_0 0 1 size @data +Resize op_1 2 1 input size out coordinate_transformation_mode=* mode=nearest nearest_mode=floor cubic_coeff_a=* +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "F.interpolate"; + } + + bool match(const std::map& matched_operators, const std::map& /*captured_params*/, const std::map& captured_attrs) const + { + auto size = captured_attrs.at("op_0.data"); + if (size.type != 5) + return false; + + int size_count = size.data.size() / sizeof(int64_t); + if (size_count < 3 || size_count > 5) + return false; + + const std::vector& input_shape = matched_operators.at("op_1")->inputs[0]->shape; + if (input_shape.size() < 3 || input_shape.size() > 5) + return false; + + const int64_t* ps = (const int64_t*)size.data.data(); + if (input_shape[0] != ps[0] || input_shape[1] != ps[1]) + return false; + + return true; + } + + void write(Operator* op, const std::map& /*captured_params*/, const std::map& captured_attrs) const + { + auto size = captured_attrs.at("op_0.data"); + int size_count = size.data.size() / sizeof(float); + const int64_t* ps = (const int64_t*)size.data.data(); + + op->params["mode"] = "nearest"; + op->params["align_corners"] = false; + if (size_count == 3) + op->params["size"] = {ps[2]}; + if (size_count == 4) + op->params["size"] = {ps[2], ps[3]}; + if (size_count == 5) + op->params["size"] = {ps[2], ps[3], ps[4]}; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_interpolate_onnx, 10) + +class F_interpolate_onnx_1 : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input 0 1 input +pnnx.Attribute op_0 0 1 scale_factor @data +Resize op_1 2 1 input scale_factor out coordinate_transformation_mode=* mode=nearest nearest_mode=floor cubic_coeff_a=* +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "F.interpolate"; + } + + bool match(const std::map& /*captured_params*/, const std::map& captured_attrs) const + { + auto scale_factor = captured_attrs.at("op_0.data"); + if (scale_factor.type != 1) + return false; + + int scale_factor_count = scale_factor.data.size() / sizeof(float); + if (scale_factor_count < 3 || scale_factor_count > 5) + return false; + + const float* ps = (const float*)scale_factor.data.data(); + if (ps[0] != 1.f || ps[1] != 1.f) + return false; + + return true; + } + + void write(Operator* op, const std::map& /*captured_params*/, const std::map& captured_attrs) const + { + auto scale_factor = captured_attrs.at("op_0.data"); + int scale_factor_count = scale_factor.data.size() / sizeof(float); + const float* ps = (const float*)scale_factor.data.data(); + + op->params["mode"] = "nearest"; + op->params["recompute_scale_factor"] = true; + if (scale_factor_count == 3) + op->params["scale_factor"] = {ps[2]}; + if (scale_factor_count == 4) + op->params["scale_factor"] = {ps[2], ps[3]}; + if (scale_factor_count == 5) + op->params["scale_factor"] = {ps[2], ps[3], ps[4]}; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_interpolate_onnx_1, 10) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/F_linear.cpp b/tools/pnnx/src/pass_level2/F_linear.cpp index 91ad0951e..4c454581e 100644 --- a/tools/pnnx/src/pass_level2/F_linear.cpp +++ b/tools/pnnx/src/pass_level2/F_linear.cpp @@ -198,7 +198,7 @@ public: const char* match_pattern_graph() const { return R"PNNXIR(7767517 -5 4 +4 3 pnnx.Input input_0 0 1 input pnnx.Attribute weight 0 1 weight @data=(%in_features,%out_features)f32 MatMul matmul 2 1 input weight out @@ -209,7 +209,7 @@ pnnx.Output output 1 0 out const char* replace_pattern_graph() const { return R"PNNXIR(7767517 -5 4 +4 3 pnnx.Input input_0 0 1 input pnnx.Attribute weight 0 1 weight F.linear linear 2 1 input weight out bias=None $weight=weight @@ -220,4 +220,62 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_linear_onnx_2, 9) +class F_linear_onnx_3 : public F_linear_onnx_1 +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +6 5 +pnnx.Input input_0 0 1 input +pnnx.Attribute weight 0 1 weight @data=(%in_features,%out_features)f32 +pnnx.Attribute bias 0 1 bias @data=(%out_features)f32 +MatMul matmul 2 1 input weight mm +aten::add add 2 1 mm bias out +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* replace_pattern_graph() const + { + return R"PNNXIR(7767517 +5 4 +pnnx.Input input_0 0 1 input +pnnx.Attribute weight 0 1 weight +pnnx.Attribute bias 0 1 bias +F.linear linear 3 1 input weight bias out $weight=weight $bias=bias +pnnx.Output output 1 0 out +)PNNXIR"; + } + + void write(const std::map& ops, const std::map& captured_params, const std::map& captured_attrs) const + { + F_linear_onnx_1::write(ops, captured_params, captured_attrs); + + Operator* op_bias = ops.at("bias"); + op_bias->attrs["data"] = captured_attrs.at("bias.data"); + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_linear_onnx_3, 8) + +class F_linear_onnx_4 : public F_linear_onnx_3 +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +6 5 +pnnx.Input input_0 0 1 input +pnnx.Attribute weight 0 1 weight @data=(%in_features,%out_features)f32 +pnnx.Attribute bias 0 1 bias @data=(%out_features)f32 +MatMul matmul 2 1 input weight mm +aten::add add 2 1 bias mm out +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_linear_onnx_4, 8) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/F_pad.cpp b/tools/pnnx/src/pass_level2/F_pad.cpp index c4007d25f..b30ea045a 100644 --- a/tools/pnnx/src/pass_level2/F_pad.cpp +++ b/tools/pnnx/src/pass_level2/F_pad.cpp @@ -259,4 +259,26 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_pad_8, 10) +class F_pad_onnx : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input_0 0 1 input +pnnx.Input input_1 0 1 pad +Pad op_0 2 1 input pad out mode=constant +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "F.pad"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_pad_onnx, 10) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/F_softmax.cpp b/tools/pnnx/src/pass_level2/F_softmax.cpp index 8d8068a43..956feeeb1 100644 --- a/tools/pnnx/src/pass_level2/F_softmax.cpp +++ b/tools/pnnx/src/pass_level2/F_softmax.cpp @@ -60,4 +60,25 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_softmax_1, 10) +class F_softmax_onnx : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +3 2 +pnnx.Input input_0 0 1 input +Softmax op_0 1 1 input out axis=%dim +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "F.softmax"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_softmax_onnx, 10) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/Tensor_reshape.cpp b/tools/pnnx/src/pass_level2/Tensor_reshape.cpp index 6cf95e06b..e71349e65 100644 --- a/tools/pnnx/src/pass_level2/Tensor_reshape.cpp +++ b/tools/pnnx/src/pass_level2/Tensor_reshape.cpp @@ -44,10 +44,11 @@ public: const char* match_pattern_graph() const { return R"PNNXIR(7767517 -4 3 +5 4 pnnx.Input input_0 0 1 input pnnx.Input input_1 0 1 shape -Reshape op_0 2 1 input shape out allowzero=* +aten::cat op_0 1 1 shape cat dim=0 +Reshape op_1 2 1 input cat out allowzero=* pnnx.Output output 1 0 out )PNNXIR"; } @@ -58,9 +59,27 @@ pnnx.Output output 1 0 out } }; -REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx, 20) +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx, 19) + +class Tensor_reshape_onnx_1 : public Tensor_reshape_onnx +{ +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 shape +aten::cat op_0 1 1 shape cat dim=0 +Reshape op_1 2 1 input cat out +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx_1, 19) -class Tensor_reshape_onnx_1 : public GraphRewriterPass +class Tensor_reshape_onnx_2 : public Tensor_reshape_onnx { public: const char* match_pattern_graph() const @@ -69,17 +88,29 @@ public: 4 3 pnnx.Input input_0 0 1 input pnnx.Input input_1 0 1 shape -Reshape op_0 2 1 input shape out +Reshape op_1 2 1 input shape out allowzero=* pnnx.Output output 1 0 out )PNNXIR"; } +}; - const char* type_str() const +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx_2, 20) + +class Tensor_reshape_onnx_3 : public Tensor_reshape_onnx +{ +public: + const char* match_pattern_graph() const { - return "Tensor.reshape"; + return R"PNNXIR(7767517 +4 3 +pnnx.Input input_0 0 1 input +pnnx.Input input_1 0 1 shape +Reshape op_1 2 1 input shape out +pnnx.Output output 1 0 out +)PNNXIR"; } }; -REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx_1, 20) +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx_3, 20) } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/Tensor_select.cpp b/tools/pnnx/src/pass_level2/Tensor_select.cpp index 07760fcfc..303616bce 100644 --- a/tools/pnnx/src/pass_level2/Tensor_select.cpp +++ b/tools/pnnx/src/pass_level2/Tensor_select.cpp @@ -46,7 +46,7 @@ public: { return R"PNNXIR(7767517 3 2 -pnnx.Input input_0 0 1 input +pnnx.Input input 0 1 input aten::select op_0 1 1 input out dim=%dim index=%index pnnx.Output output 1 0 out )PNNXIR"; @@ -60,4 +60,26 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_select_1, 20) +class Tensor_select_onnx : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input 0 1 input +prim::Constant op_0 0 1 indices value=%index +Gather op_1 2 1 input indices out axis=%dim +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "Tensor.select"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_select_onnx, 20) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/Tensor_to.cpp b/tools/pnnx/src/pass_level2/Tensor_to.cpp index bd1797d2a..4ea6f5336 100644 --- a/tools/pnnx/src/pass_level2/Tensor_to.cpp +++ b/tools/pnnx/src/pass_level2/Tensor_to.cpp @@ -119,4 +119,47 @@ REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_to, 20) REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_to_1, 20) REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_to_2, 20) +class Tensor_to_onnx : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +3 2 +pnnx.Input input 0 1 input +Cast op_0 1 1 input out to=%to +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "Tensor.to"; + } + + void write(Operator* op, const std::map& captured_params) const + { + const int to = captured_params.at("to").i; + + op->params["non_blocking"] = false; + op->params["copy"] = false; + op->params["memory_format"] = "torch.preserve_format"; + + if (to == 1) op->params["dtype"] = "torch.float"; + if (to == 2) op->params["dtype"] = "torch.uint8"; + if (to == 3) op->params["dtype"] = "torch.int8"; + if (to == 5) op->params["dtype"] = "torch.short"; + if (to == 6) op->params["dtype"] = "torch.int"; + if (to == 7) op->params["dtype"] = "torch.long"; + if (to == 9) op->params["dtype"] = "torch.bool"; + if (to == 10) op->params["dtype"] = "torch.half"; + if (to == 11) op->params["dtype"] = "torch.double"; + if (to == 14) op->params["dtype"] = "torch.complex64"; + if (to == 15) op->params["dtype"] = "torch.complex128"; + if (to == 16) op->params["dtype"] = "torch.bfloat16"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_to_onnx, 20) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/torch_clamp.cpp b/tools/pnnx/src/pass_level2/torch_clamp.cpp index dd4e0cbdd..a04d6341b 100644 --- a/tools/pnnx/src/pass_level2/torch_clamp.cpp +++ b/tools/pnnx/src/pass_level2/torch_clamp.cpp @@ -41,36 +41,13 @@ REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_clamp, 20) class torch_clamp_onnx : 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 min -pnnx.Input input_2 0 1 max -Clip op_0 3 1 input min max out -pnnx.Output output 1 0 out -)PNNXIR"; - } - - const char* type_str() const - { - return "torch.clamp"; - } -}; - -REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_clamp_onnx, 20) - -class torch_clamp_onnx_1 : public GraphRewriterPass -{ public: const char* match_pattern_graph() const { return R"PNNXIR(7767517 3 2 pnnx.Input input 0 1 input -Clip op_0 1 1 input out min=%min max=%max +aten::clamp op_0 1 1 input out min=%min max=%max pnnx.Output output 1 0 out )PNNXIR"; } @@ -81,6 +58,6 @@ pnnx.Output output 1 0 out } }; -REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_clamp_onnx_1, 20) +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_clamp_onnx, 20) } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/torch_split.cpp b/tools/pnnx/src/pass_level2/torch_split.cpp index f3369e17e..1b08a43e7 100644 --- a/tools/pnnx/src/pass_level2/torch_split.cpp +++ b/tools/pnnx/src/pass_level2/torch_split.cpp @@ -105,4 +105,21 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_split_onnx, 20) +class torch_split_onnx_1 : public torch_split_onnx +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input_0 0 1 tensor +pnnx.Input input_1 0 1 split_size_or_sections +aten::split op_0 2 1 tensor split_size_or_sections out dim=%dim indices=None +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_split_onnx_1, 20) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/torch_squeeze.cpp b/tools/pnnx/src/pass_level2/torch_squeeze.cpp index 4300b3ef6..2b6b685b0 100644 --- a/tools/pnnx/src/pass_level2/torch_squeeze.cpp +++ b/tools/pnnx/src/pass_level2/torch_squeeze.cpp @@ -80,4 +80,42 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_squeeze_0, 20) +class torch_squeeze_onnx : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input_0 0 1 input +pnnx.Input input_1 0 1 dim +Squeeze op_0 2 1 input dim out +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "torch.squeeze"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_squeeze_onnx, 20) + +class torch_squeeze_onnx_1 : public torch_squeeze_onnx +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +3 2 +pnnx.Input input 0 1 input +Squeeze op_0 1 1 input out axes=%dim +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_squeeze_onnx_1, 20) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level2/torch_unsqueeze.cpp b/tools/pnnx/src/pass_level2/torch_unsqueeze.cpp index c7f2d8ad4..f5780509b 100644 --- a/tools/pnnx/src/pass_level2/torch_unsqueeze.cpp +++ b/tools/pnnx/src/pass_level2/torch_unsqueeze.cpp @@ -59,4 +59,42 @@ pnnx.Output output 1 0 out REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_unsqueeze_01, 20) +class torch_unsqueeze_onnx : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input_0 0 1 input +pnnx.Input input_1 0 1 dim +Unsqueeze op_0 2 1 input dim out +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "torch.unsqueeze"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_unsqueeze_onnx, 20) + +class torch_unsqueeze_onnx_1 : public torch_unsqueeze_onnx +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +3 2 +pnnx.Input input 0 1 input +Unsqueeze op_0 1 1 input out axes=%dim +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_unsqueeze_onnx_1, 20) + } // namespace pnnx diff --git a/tools/pnnx/src/pass_level3/fuse_expression.cpp b/tools/pnnx/src/pass_level3/fuse_expression.cpp index 2ed20abe7..31db1f2f0 100644 --- a/tools/pnnx/src/pass_level3/fuse_expression.cpp +++ b/tools/pnnx/src/pass_level3/fuse_expression.cpp @@ -54,7 +54,7 @@ static bool operand_maybe_tensor(const Operand* operand) if (op->type == "aten::size") { - return false; + return op->inputs.size() == 1; } if (op->type == "aten::Int") @@ -508,8 +508,11 @@ static void fuse_expression(Graph& graph, Operand* operand, std::string& expr, s { expr += "size("; fuse_expression(graph, op->inputs[0], expr, inputs, foldable_constants, zip); - expr += ","; - fuse_expression(graph, op->inputs[1], expr, inputs, foldable_constants, zip); + if (op->inputs.size() == 2) + { + expr += ","; + fuse_expression(graph, op->inputs[1], expr, inputs, foldable_constants, zip); + } expr += ")"; } else if (op->type == "aten::Int") diff --git a/tools/pnnx/src/pass_level5.cpp b/tools/pnnx/src/pass_level5.cpp index 7b91eabfe..4a1a2760a 100644 --- a/tools/pnnx/src/pass_level5.cpp +++ b/tools/pnnx/src/pass_level5.cpp @@ -49,6 +49,7 @@ #include "pass_level5/fuse_slice_copy.h" #include "pass_level5/fuse_slice_indices.h" #include "pass_level5/fuse_slice_to_tensor_split.h" +#include "pass_level5/fuse_slice_squeeze_to_select.h" #include "pass_level5/fuse_static_batchnorm.h" #include "pass_level5/fuse_static_conv.h" #include "pass_level5/fuse_static_convtranspose.h" @@ -57,6 +58,7 @@ #include "pass_level5/fuse_static_instancenorm.h" #include "pass_level5/fuse_static_layernorm.h" #include "pass_level5/fuse_static_linear.h" +#include "pass_level5/fuse_swish.h" #include "pass_level5/normalize_einsum_equation.h" #include "pass_level4/dead_code_elimination.h" #include "pass_level4/canonicalize.h" @@ -77,6 +79,8 @@ void pass_level5(Graph& g, const std::set& foldable_constants, cons eliminate_noop_slice(g); + fuse_slice_squeeze_to_select(g); + fuse_slice_indices(g); normalize_einsum_equation(g); @@ -140,6 +144,8 @@ void pass_level5(Graph& g, const std::set& foldable_constants, cons fuse_multiheadattention(g); fuse_scaled_dot_product_attention(g); + fuse_swish(g); + fuse_index_expression(g); dead_code_elimination(g); diff --git a/tools/pnnx/src/pass_level5/eval_expression.cpp b/tools/pnnx/src/pass_level5/eval_expression.cpp index 10b38b9fc..44e1f7e36 100644 --- a/tools/pnnx/src/pass_level5/eval_expression.cpp +++ b/tools/pnnx/src/pass_level5/eval_expression.cpp @@ -119,40 +119,49 @@ static std::string eval_expression(const Operator* op) { std::string a = exprstack.top(); exprstack.pop(); - std::string b = exprstack.top(); - exprstack.pop(); - if (token_is_argument(a) && token_is_literal(b)) + if (exprstack.empty()) { - int input_index = std::stoi(a.substr(1)); - if (op->inputs[input_index]->shape.empty()) - { - std::string r = std::string("size(") + a + "," + b + ")"; - exprstack.push(r); - } - else + std::string r = std::string("size(") + a + ")"; + exprstack.push(r); + } + else + { + std::string b = exprstack.top(); + exprstack.pop(); + + if (token_is_argument(a) && token_is_literal(b)) { - int bi = std::stoi(b); - if (bi < 0) - bi = op->inputs[input_index]->shape.size() + bi; - int r = op->inputs[input_index]->shape[bi]; - if (r == -1) + int input_index = std::stoi(a.substr(1)); + if (op->inputs[input_index]->shape.empty()) { - // do not evaluate dynamic size info as -1 - // just keep the size expression std::string r = std::string("size(") + a + "," + b + ")"; exprstack.push(r); } else { - exprstack.push(std::to_string(r)); + int bi = std::stoi(b); + if (bi < 0) + bi = op->inputs[input_index]->shape.size() + bi; + int r = op->inputs[input_index]->shape[bi]; + if (r == -1) + { + // do not evaluate dynamic size info as -1 + // just keep the size expression + std::string r = std::string("size(") + a + "," + b + ")"; + exprstack.push(r); + } + else + { + exprstack.push(std::to_string(r)); + } } } - } - else - { - std::string r = std::string("size(") + a + "," + b + ")"; - exprstack.push(r); + else + { + std::string r = std::string("size(") + a + "," + b + ")"; + exprstack.push(r); + } } } else if (t == "int" diff --git a/tools/pnnx/src/pass_level5/fuse_layernorm.cpp b/tools/pnnx/src/pass_level5/fuse_layernorm.cpp index 2a410b43a..2c8636a87 100644 --- a/tools/pnnx/src/pass_level5/fuse_layernorm.cpp +++ b/tools/pnnx/src/pass_level5/fuse_layernorm.cpp @@ -65,12 +65,112 @@ pnnx.Output output 1 0 out } }; +class fuse_layernorm_pass_1 : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + // clang-format off + // *INDENT-OFF* + + return R"PNNXIR(7767517 +9 8 +pnnx.Input input 0 1 input #input=(1,?,%c)f32 +pnnx.Attribute op_0 0 1 weight @data #weight=(%c)f32 +pnnx.Attribute op_1 0 1 bias @data #bias=(%c)f32 +torch.mean op_2 1 1 input mean dim=(-1) keepdim=True +pnnx.Expression op_3 2 1 input mean 173 expr=sub(@0,@1) +pnnx.Expression op_4 1 1 173 174 expr=pow(@0,2.000000e+00) +torch.mean op_5 1 1 174 var dim=(-1) keepdim=True +pnnx.Expression op_6 4 1 173 var weight bias out expr=add(mul(div(@0,sqrt(add(@1,%eps))),@2),@3) +pnnx.Output output 1 0 out +)PNNXIR"; + + // *INDENT-ON* + // clang-format on + } + + const char* replace_pattern_graph() const + { + return R"PNNXIR(7767517 +3 2 +pnnx.Input input 0 1 input +nn.LayerNorm ln 1 1 input out elementwise_affine=True eps=%eps normalized_shape=(%c) @weight=%op_0.data @bias=%op_1.data +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +class fuse_layernorm_pass_2 : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + // clang-format off + // *INDENT-OFF* + + return R"PNNXIR(7767517 +7 6 +pnnx.Input input 0 1 input #input=(1,?,%c)f32 +torch.mean op_0 1 1 input mean dim=(-1) keepdim=True +pnnx.Expression op_1 2 1 input mean 25 expr=sub(@0,@1) +pnnx.Expression op_2 1 1 25 26 expr=pow(@0,2.000000e+00) +torch.mean op_3 1 1 26 var dim=(-1) keepdim=True +pnnx.Expression op_4 2 1 25 var out expr=div(@0,sqrt(add(@1,%eps))) +pnnx.Output output 1 0 out +)PNNXIR"; + + // *INDENT-ON* + // clang-format on + } + + const char* replace_pattern_graph() const + { + return R"PNNXIR(7767517 +3 2 +pnnx.Input input 0 1 input +nn.LayerNorm ln 1 1 input out elementwise_affine=False eps=%eps normalized_shape=(%c) +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +class fuse_layernorm_pass_2_1 : public fuse_layernorm_pass_2 +{ +public: + const char* match_pattern_graph() const + { + // clang-format off + // *INDENT-OFF* + + return R"PNNXIR(7767517 +7 6 +pnnx.Input input 0 1 input #input=(1,?,?,%c)f32 +torch.mean op_0 1 1 input mean dim=(-1) keepdim=True +pnnx.Expression op_1 2 1 input mean 32 expr=sub(@0,@1) +pnnx.Expression op_2 1 1 32 33 expr=pow(@0,2.000000e+00) +torch.mean op_3 1 1 33 var dim=(-1) keepdim=True +pnnx.Expression op_4 2 1 32 var out expr=div(@0,sqrt(add(@1,%eps))) +pnnx.Output output 1 0 out +)PNNXIR"; + + // *INDENT-ON* + // clang-format on + } +}; + void fuse_layernorm(Graph& graph) { fuse_layernorm_pass a; + fuse_layernorm_pass_1 b; + fuse_layernorm_pass_2 c; + fuse_layernorm_pass_2_1 c1; int opindex = 0; pnnx_graph_rewrite(graph, &a, opindex); + pnnx_graph_rewrite(graph, &b, opindex); + pnnx_graph_rewrite(graph, &c, opindex); + pnnx_graph_rewrite(graph, &c1, opindex); } } // namespace pnnx diff --git a/tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp b/tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp index d729f3957..55661366c 100644 --- a/tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp +++ b/tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp @@ -51,7 +51,7 @@ torch.unbind op_3 1 3 3 4 5 6 dim=0 pnnx.Expression op_4 1 1 4 7 expr=mul(@0,%inv_sqrt_embed_dim_per_head) Tensor.permute op_5 1 1 5 8 dims=(0,1,3,2) torch.matmul op_6 2 1 7 8 9 -F.softmax op_7 1 1 9 10 dim=-1 +F.softmax softmax 1 1 9 10 dim=%softmax_dim torch.matmul op_8 2 1 10 6 11 Tensor.permute op_9 1 1 11 12 dims=(0,2,1,3) Tensor.reshape op_10 1 1 12 13 shape=(%batch,%size,%embed_dim) @@ -70,13 +70,14 @@ pnnx.Output output 1 0 out )PNNXIR"; } - bool match(const std::map& captured_params) const + bool match(const std::map& matched_operators, const std::map& captured_params, const std::map& /*captured_attrs*/) const { const int embed_dim = captured_params.at("embed_dim").i; const int qkv_out_features = captured_params.at("qkv_out_features").i; const int num_heads = captured_params.at("num_heads").i; const int feat_per_head = captured_params.at("feat_per_head").i; const float inv_sqrt_embed_dim_per_head = captured_params.at("inv_sqrt_embed_dim_per_head").f; + const int softmax_dim = captured_params.at("softmax_dim").i; if (qkv_out_features != embed_dim * 3) return false; @@ -87,6 +88,10 @@ pnnx.Output output 1 0 out if (!NearlyEqual(inv_sqrt_embed_dim_per_head, 1.f / sqrt(feat_per_head), 0.001)) return false; + int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size(); + if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1) + return false; + return true; } @@ -157,7 +162,7 @@ Tensor.permute op_6 1 1 5 9 dims=(0,2,1,3) torch.transpose op_7 1 1 8 10 dim0=-1 dim1=-2 torch.matmul op_8 2 1 9 10 11 pnnx.Expression op_9 1 1 11 12 expr=mul(@0,%inv_sqrt_embed_dim_per_head) -nn.Softmax op_10 1 1 12 13 dim=-1 +nn.Softmax softmax 1 1 12 13 dim=%softmax_dim Tensor.permute op_11 1 1 7 14 dims=(0,2,1,3) torch.matmul op_12 2 1 13 14 15 Tensor.permute op_13 1 1 15 16 dims=(0,2,1,3) @@ -168,6 +173,129 @@ pnnx.Output output 1 0 out } }; +class fuse_multiheadattention_pass_11_0 : public fuse_multiheadattention_pass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +22 21 +pnnx.Input input 0 1 input +Tensor.permute op_1 1 1 input 13 dims=(1,0,2) +nn.Linear op_0 1 1 13 14 bias=%qkvbias in_features=%embed_dim out_features=%qkv_out_features @bias @weight +Tensor.slice op_2 1 1 14 15 dim=-1 end=%embed_dim start=0 step=1 +Tensor.slice op_3 1 1 14 16 dim=-1 end=%embed_dim2 start=%embed_dim step=1 +Tensor.slice op_4 1 1 14 17 dim=-1 end=%qkv_out_features start=%embed_dim2 step=1 +Tensor.reshape op_5 1 1 15 18 shape=(%size,%num_heads,%feat_per_head) +Tensor.permute op_6 1 1 18 19 dims=(1,0,2) +Tensor.reshape op_7 1 1 16 20 shape=(%size,%num_heads,%feat_per_head) +Tensor.reshape op_8 1 1 17 21 shape=(%size,%num_heads,%feat_per_head) +Tensor.permute op_9 1 1 21 22 dims=(1,0,2) +pnnx.Expression op_10 1 1 19 23 expr=div(@0,%sqrt_embed_dim_per_head) +Tensor.permute op_11 1 1 20 24 dims=(1,2,0) +torch.matmul op_12 2 1 23 24 25 +F.softmax softmax 1 1 25 26 dim=%softmax_dim +torch.matmul op_14 2 1 26 22 27 +Tensor.permute op_15 1 1 27 28 dims=(1,0,2) +Tensor.reshape op_16 1 1 28 29 shape=(%size,%embed_dim) +nn.Linear out_proj 1 1 29 30 bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight +Tensor.reshape op_24 1 1 30 31 shape=(%size,%batch,%embed_dim) +Tensor.permute op_25 1 1 31 out dims=(1,0,2) +pnnx.Output output 1 0 out +)PNNXIR"; + } + + bool match(const std::map& matched_operators, const std::map& captured_params, const std::map& /*captured_attrs*/) const + { + const int embed_dim = captured_params.at("embed_dim").i; + const int embed_dim2 = captured_params.at("embed_dim2").i; + const int qkv_out_features = captured_params.at("qkv_out_features").i; + const int num_heads = captured_params.at("num_heads").i; + const int feat_per_head = captured_params.at("feat_per_head").i; + const float sqrt_embed_dim_per_head = captured_params.at("sqrt_embed_dim_per_head").f; + const int softmax_dim = captured_params.at("softmax_dim").i; + + if (qkv_out_features != embed_dim * 3 || embed_dim2 != embed_dim * 2) + return false; + + if (embed_dim != num_heads * feat_per_head) + return false; + + if (!NearlyEqual(sqrt_embed_dim_per_head, sqrt(feat_per_head), 0.001)) + return false; + + int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size(); + if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1) + return false; + + return true; + } +}; + +class fuse_multiheadattention_pass_11_1 : public fuse_multiheadattention_pass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +28 27 +pnnx.Input input 0 1 input +Tensor.permute op_1 1 1 input 9 dims=(1,0,2) +nn.Linear op_0 1 1 9 10 bias=%qkvbias in_features=%embed_dim out_features=%qkv_out_features @bias @weight +Tensor.reshape op_2 1 1 10 11 shape=(1,%size,%batch,3,%embed_dim) +Tensor.permute op_3 1 1 11 12 dims=(3,1,2,0,4) +torch.squeeze op_4 1 1 12 13 dim=3 +torch.unbind op_5 1 3 13 14 15 16 dim=0 +Tensor.reshape op_6 1 1 14 17 shape=(%size,%num_heads,%feat_per_head) +Tensor.permute op_7 1 1 17 18 dims=(1,0,2) +Tensor.reshape op_8 1 1 15 19 shape=(%size,%num_heads,%feat_per_head) +Tensor.permute op_9 1 1 19 20 dims=(1,0,2) +Tensor.reshape op_10 1 1 16 21 shape=(%size,%num_heads,%feat_per_head) +Tensor.permute op_11 1 1 21 22 dims=(1,0,2) +Tensor.reshape op_12 1 1 18 23 shape=(%batch,%num_heads,%size,%feat_per_head) +Tensor.reshape op_13 1 1 20 24 shape=(%batch,%num_heads,%size,%feat_per_head) +Tensor.reshape op_14 1 1 22 25 shape=(%batch,%num_heads,%size,%feat_per_head) +Tensor.permute op_15 1 1 24 26 dims=(0,1,3,2) +pnnx.Expression op_16 1 1 23 27 expr=mul(@0,%inv_sqrt_embed_dim_per_head) +pnnx.Expression op_17 1 1 26 28 expr=mul(@0,%inv_sqrt_embed_dim_per_head) +torch.matmul op_18 2 1 27 28 29 +F.softmax softmax 1 1 29 30 dim=%softmax_dim +torch.matmul op_20 2 1 30 25 31 +Tensor.permute op_21 1 1 31 32 dims=(2,0,1,3) +Tensor.reshape op_22 1 1 32 33 shape=(%size,%embed_dim) +nn.Linear out_proj 1 1 33 34 bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight +Tensor.reshape op_24 1 1 34 35 shape=(%size,%batch,%embed_dim) +Tensor.permute op_25 1 1 35 out dims=(1,0,2) +pnnx.Output output 1 0 out +)PNNXIR"; + } + + bool match(const std::map& matched_operators, const std::map& captured_params, const std::map& /*captured_attrs*/) const + { + const int embed_dim = captured_params.at("embed_dim").i; + const int qkv_out_features = captured_params.at("qkv_out_features").i; + const int num_heads = captured_params.at("num_heads").i; + const int feat_per_head = captured_params.at("feat_per_head").i; + const float inv_sqrt_embed_dim_per_head = captured_params.at("inv_sqrt_embed_dim_per_head").f; + const int softmax_dim = captured_params.at("softmax_dim").i; + + if (qkv_out_features != embed_dim * 3) + return false; + + if (embed_dim != num_heads * feat_per_head) + return false; + + if (!NearlyEqual(inv_sqrt_embed_dim_per_head, sqrt(1.f / sqrt(feat_per_head)), 0.001)) + return false; + + int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size(); + if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1) + return false; + + return true; + } +}; + class fuse_multiheadattention_pass_sameqkv : public GraphRewriterPass { public: @@ -189,7 +317,7 @@ Tensor.permute op_9 1 1 35 40 dims=(0,2,1,3) Tensor.reshape op_10 1 1 40 41 shape=(%num_heads,%size,%feat_per_head) Tensor.permute op_11 1 1 39 42 dims=(0,2,1) torch.matmul op_12 2 1 41 42 43 -F.softmax op_13 1 1 43 44 dim=-1 +F.softmax softmax 1 1 43 44 dim=%softmax_dim Tensor.permute op_14 1 1 37 45 dims=(0,2,1,3) Tensor.reshape op_15 1 1 45 46 shape=(%num_heads,%size,%feat_per_head) torch.matmul op_16 2 1 44 46 47 @@ -211,12 +339,13 @@ pnnx.Output output 1 0 out )PNNXIR"; } - bool match(const std::map& captured_params) const + bool match(const std::map& matched_operators, const std::map& captured_params, const std::map& /*captured_attrs*/) const { const int embed_dim = captured_params.at("embed_dim").i; const int num_heads = captured_params.at("num_heads").i; const int feat_per_head = captured_params.at("feat_per_head").i; const float inv_sqrt_embed_dim_per_head = captured_params.at("inv_sqrt_embed_dim_per_head").f; + const int softmax_dim = captured_params.at("softmax_dim").i; if (embed_dim != num_heads * feat_per_head) return false; @@ -224,6 +353,10 @@ pnnx.Output output 1 0 out if (!NearlyEqual(inv_sqrt_embed_dim_per_head, 1.f / sqrt(feat_per_head), 0.001)) return false; + int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size(); + if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1) + return false; + return true; } @@ -326,7 +459,7 @@ Tensor.permute op_9 1 1 36 41 dims=(0,2,1,3) Tensor.reshape op_10 1 1 41 42 shape=(%num_heads,%qsize,%feat_per_head) Tensor.permute op_11 1 1 40 43 dims=(0,2,1) torch.matmul op_12 2 1 42 43 44 -F.softmax op_13 1 1 44 45 dim=-1 +F.softmax softmax 1 1 44 45 dim=%softmax_dim Tensor.permute op_14 1 1 38 46 dims=(0,2,1,3) Tensor.reshape op_15 1 1 46 47 shape=(%num_heads,%kvsize,%feat_per_head) torch.matmul op_16 2 1 45 47 48 @@ -461,7 +594,7 @@ Tensor.permute op_9 1 1 36 41 dims=(0,2,1,3) Tensor.reshape op_10 1 1 41 42 shape=(%num_heads,%qsize,%feat_per_head) Tensor.permute op_11 1 1 40 43 dims=(0,2,1) torch.matmul op_12 2 1 42 43 44 -F.softmax op_13 1 1 44 45 dim=-1 +F.softmax softmax 1 1 44 45 dim=%softmax_dim Tensor.permute op_14 1 1 38 46 dims=(0,2,1,3) Tensor.reshape op_15 1 1 46 47 shape=(%num_heads,%kvsize,%feat_per_head) torch.matmul op_16 2 1 45 47 48 @@ -505,7 +638,7 @@ Tensor.permute op_8 1 1 35 40 dims=(0,2,1,3) Tensor.reshape op_9 1 1 40 41 shape=(%num_heads,%size,%feat_per_head) torch.einsum op_10 2 1 41 39 42 equation=ijl,ikl->ijk pnnx.Expression op_11 1 1 42 43 expr=mul(@0,%inv_sqrt_embed_dim_per_head) -F.softmax op_12 1 1 43 44 dim=-1 +F.softmax softmax 1 1 43 44 dim=%softmax_dim Tensor.permute op_13 1 1 37 45 dims=(0,2,1,3) Tensor.reshape op_14 1 1 45 46 shape=(%num_heads,%size,%feat_per_head) torch.einsum op_15 2 1 44 46 47 equation=ijl,ilk->ijk @@ -537,7 +670,7 @@ Tensor.permute op_7 1 1 53 54 dims=(0,1,3,2) torch.transpose op_8 1 1 50 55 dim0=1 dim1=2 torch.matmul op_9 2 1 55 54 56 pnnx.Expression op_10 1 1 56 57 expr=div(@0,%sqrt_feat_per_head) -F.softmax op_11 1 1 57 58 dim=-1 +F.softmax softmax 1 1 57 58 dim=%softmax_dim torch.transpose op_12 1 1 52 59 dim0=1 dim1=2 torch.matmul op_13 2 1 58 59 60 torch.transpose op_14 1 1 60 61 dim0=1 dim1=2 @@ -547,12 +680,13 @@ pnnx.Output output 1 0 out )PNNXIR"; } - bool match(const std::map& captured_params) const + bool match(const std::map& matched_operators, const std::map& captured_params, const std::map& /*captured_attrs*/) const { const int embed_dim = captured_params.at("embed_dim").i; const int num_heads = captured_params.at("num_heads").i; const int feat_per_head = captured_params.at("feat_per_head").i; const float sqrt_feat_per_head = captured_params.at("sqrt_feat_per_head").f; + const int softmax_dim = captured_params.at("softmax_dim").i; if (embed_dim != num_heads * feat_per_head) return false; @@ -560,6 +694,10 @@ pnnx.Output output 1 0 out if (!NearlyEqual(sqrt_feat_per_head, sqrt(feat_per_head), 0.001)) return false; + int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size(); + if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1) + return false; + return true; } }; @@ -585,7 +723,7 @@ Tensor.permute op_7 1 1 53 54 dims=(0,1,3,2) torch.transpose op_8 1 1 50 55 dim0=1 dim1=2 torch.matmul op_9 2 1 55 54 56 pnnx.Expression op_10 1 1 56 57 expr=div(@0,%sqrt_feat_per_head) -F.softmax op_11 1 1 57 58 dim=-1 +F.softmax softmax 1 1 57 58 dim=%softmax_dim torch.transpose op_12 1 1 52 59 dim0=1 dim1=2 torch.matmul op_13 2 1 58 59 60 torch.transpose op_14 1 1 60 61 dim0=1 dim1=2 @@ -595,12 +733,13 @@ pnnx.Output output 1 0 out )PNNXIR"; } - bool match(const std::map& captured_params) const + bool match(const std::map& matched_operators, const std::map& captured_params, const std::map& /*captured_attrs*/) const { const int embed_dim = captured_params.at("embed_dim").i; const int num_heads = captured_params.at("num_heads").i; const int feat_per_head = captured_params.at("feat_per_head").i; const float sqrt_feat_per_head = captured_params.at("sqrt_feat_per_head").f; + const int softmax_dim = captured_params.at("softmax_dim").i; if (embed_dim != num_heads * feat_per_head) return false; @@ -608,6 +747,10 @@ pnnx.Output output 1 0 out if (!NearlyEqual(sqrt_feat_per_head, sqrt(feat_per_head), 0.001)) return false; + int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size(); + if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1) + return false; + return true; } }; @@ -634,7 +777,7 @@ Tensor.permute op_8 1 1 36 41 dims=(0,2,1,3) Tensor.reshape op_9 1 1 41 42 shape=(%num_heads,%qsize,%feat_per_head) torch.einsum op_10 2 1 42 40 43 equation=ijl,ikl->ijk pnnx.Expression op_11 1 1 43 44 expr=mul(@0,%inv_sqrt_embed_dim_per_head) -F.softmax op_12 1 1 44 45 dim=-1 +F.softmax softmax 1 1 44 45 dim=%softmax_dim Tensor.permute op_13 1 1 38 46 dims=(0,2,1,3) Tensor.reshape op_14 1 1 46 47 shape=(%num_heads,%kvsize,%feat_per_head) torch.einsum op_15 2 1 45 47 48 equation=ijl,ilk->ijk @@ -668,7 +811,7 @@ Tensor.permute op_8 1 1 36 41 dims=(0,2,1,3) Tensor.reshape op_9 1 1 41 42 shape=(%num_heads,%qsize,%feat_per_head) torch.einsum op_10 2 1 42 40 43 equation=ijl,ikl->ijk pnnx.Expression op_11 1 1 43 44 expr=mul(@0,%inv_sqrt_embed_dim_per_head) -F.softmax op_12 1 1 44 45 dim=-1 +F.softmax softmax 1 1 44 45 dim=%softmax_dim Tensor.permute op_13 1 1 38 46 dims=(0,2,1,3) Tensor.reshape op_14 1 1 46 47 shape=(%num_heads,%kvsize,%feat_per_head) torch.einsum op_15 2 1 45 47 48 equation=ijl,ilk->ijk @@ -702,7 +845,7 @@ Tensor.reshape op_9 1 1 41 42 shape=(%num_heads,%size,%feat_per_ pnnx.Attribute op_10 0 1 43 @data torch.transpose op_11 1 1 42 44 dim0=-1 dim1=-2 torch.baddbmm op_12 3 1 43 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0 -F.softmax op_13 1 1 45 46 dim=-1 +F.softmax softmax 1 1 45 46 dim=%softmax_dim Tensor.permute op_14 1 1 38 47 dims=(0,2,1,3) Tensor.reshape op_15 1 1 47 48 shape=(%num_heads,%size,%feat_per_head) torch.bmm op_16 2 1 46 48 49 @@ -739,7 +882,7 @@ pnnx.Expression op_10 2 1 40 42 43 expr=%expr_zero_shape torch.empty op_11 1 1 43 zeros torch.transpose op_12 1 1 42 44 dim0=-1 dim1=-2 torch.baddbmm op_13 3 1 zeros 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0 -F.softmax op_14 1 1 45 46 dim=-1 +F.softmax softmax 1 1 45 46 dim=%softmax_dim Tensor.permute op_15 1 1 38 47 dims=(0,2,1,3) Tensor.reshape op_16 1 1 47 48 shape=(%num_heads,%size,%feat_per_head) torch.bmm op_17 2 1 46 48 49 @@ -777,7 +920,7 @@ Tensor.reshape op_9 1 1 41 42 shape=(%num_heads,%kvsize,%feat_pe pnnx.Attribute op_10 0 1 43 @data torch.transpose op_11 1 1 42 44 dim0=-1 dim1=-2 torch.baddbmm op_12 3 1 43 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0 -F.softmax op_13 1 1 45 46 dim=-1 +F.softmax softmax 1 1 45 46 dim=%softmax_dim Tensor.permute op_14 1 1 38 47 dims=(0,2,1,3) Tensor.reshape op_15 1 1 47 48 shape=(%num_heads,%kvsize,%feat_per_head) torch.bmm op_16 2 1 46 48 49 @@ -814,7 +957,7 @@ Tensor.reshape op_9 1 1 41 42 shape=(%num_heads,%kvsize,%feat_pe pnnx.Attribute op_10 0 1 43 @data torch.transpose op_11 1 1 42 44 dim0=-1 dim1=-2 torch.baddbmm op_12 3 1 43 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0 -F.softmax op_13 1 1 45 46 dim=-1 +F.softmax softmax 1 1 45 46 dim=%softmax_dim Tensor.permute op_14 1 1 38 47 dims=(0,2,1,3) Tensor.reshape op_15 1 1 47 48 shape=(%num_heads,%kvsize,%feat_per_head) torch.bmm op_16 2 1 46 48 49 @@ -853,7 +996,7 @@ pnnx.Expression op_10 1 1 40 43 expr=%expr_zero_shape torch.empty op_11 1 1 43 zeros torch.transpose op_12 1 1 42 44 dim0=-1 dim1=-2 torch.baddbmm op_13 3 1 zeros 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0 -F.softmax op_14 1 1 45 46 dim=-1 +F.softmax softmax 1 1 45 46 dim=%softmax_dim Tensor.permute op_15 1 1 38 47 dims=(0,2,1,3) Tensor.reshape op_16 1 1 47 48 shape=(%num_heads,%kvsize,%feat_per_head) torch.bmm op_17 2 1 46 48 49 @@ -891,7 +1034,7 @@ pnnx.Expression op_10 1 1 40 43 expr=%expr_zero_shape torch.empty op_11 1 1 43 zeros torch.transpose op_12 1 1 42 44 dim0=-1 dim1=-2 torch.baddbmm op_13 3 1 zeros 40 44 45 alpha=%alpha beta=0 -F.softmax op_14 1 1 45 46 dim=-1 +F.softmax softmax 1 1 45 46 dim=%softmax_dim Tensor.permute op_15 1 1 38 47 dims=(0,2,1,3) Tensor.reshape op_16 1 1 47 48 shape=(%num_heads,%kvsize,%feat_per_head) torch.bmm op_17 2 1 46 48 49 @@ -929,7 +1072,7 @@ pnnx.Output output 1 0 out )PNNXIR"; } - bool match(const std::map& captured_params) const + bool match(const std::map& /*matched_operators*/, const std::map& captured_params, const std::map& /*captured_attrs*/) const { const int embed_dim = captured_params.at("embed_dim").i; const int num_heads = captured_params.at("num_heads").i; @@ -995,7 +1138,7 @@ pnnx.Output output 1 0 out )PNNXIR"; } - bool match(const std::map& captured_params) const + bool match(const std::map& /*matched_operators*/, const std::map& captured_params, const std::map& /*captured_attrs*/) const { const int embed_dim = captured_params.at("embed_dim").i; const int num_heads = captured_params.at("num_heads").i; @@ -1034,7 +1177,7 @@ pnnx.Output output 1 0 out )PNNXIR"; } - bool match(const std::map& captured_params) const + bool match(const std::map& /*matched_operators*/, const std::map& captured_params, const std::map& /*captured_attrs*/) const { const int embed_dim = captured_params.at("embed_dim").i; const int num_heads = captured_params.at("num_heads").i; @@ -1070,7 +1213,7 @@ Tensor.reshape op_11 1 1 10 12 shape=(%num_heads,%batch_mul_size, Tensor.reshape op_12 1 1 11 17 shape=(%num_heads,%batch_mul_size,%feat_per_head) torch.transpose op_13 1 1 12 13 dim0=1 dim1=2 torch.bmm op_14 2 1 14 13 15 -F.softmax op_15 1 1 15 16 dim=-1 +F.softmax softmax 1 1 15 16 dim=%softmax_dim torch.bmm op_16 2 1 16 17 18 Tensor.view op_17 1 1 18 19 shape=(%batch,%num_heads,%size,%feat_per_head) torch.transpose op_18 1 1 19 20 dim0=1 dim1=2 @@ -1108,7 +1251,7 @@ Tensor.view op_15 1 1 16 17 shape=(%batch,%num_heads,%size,%si pnnx.Attribute attn_mask 0 1 attn_mask @data=(1,1,%size,%size)f32 pnnx.Expression op_16 2 1 17 attn_mask 18 expr=add(@0,@1) Tensor.view op_17 1 1 18 19 shape=(%num_heads,%size,%size) -F.softmax op_18 1 1 19 20 dim=-1 +F.softmax softmax 1 1 19 20 dim=%softmax_dim torch.bmm op_19 2 1 20 21 22 Tensor.view op_20 1 1 22 23 shape=(%batch,%num_heads,%size,%feat_per_head) torch.transpose op_21 1 1 23 24 dim0=1 dim1=2 @@ -1160,7 +1303,7 @@ torch.transpose op_9 1 1 37 39 dim0=-1 dim1=-2 torch.matmul op_10 2 1 38 39 40 pnnx.Attribute attn_mask 0 1 attn_mask @data=(1,1,1,%size)f32 pnnx.Expression op_11 2 1 40 attn_mask 41 expr=add(div(@0,%sqrt_feat_per_head),@1) -F.softmax op_12 1 1 41 42 dim=-1 +F.softmax softmax 1 1 41 42 dim=%softmax_dim torch.matmul op_13 2 1 42 43 44 Tensor.permute op_14 1 1 44 45 dims=(0,2,1,3) Tensor.reshape op_15 1 1 45 46 shape=(%batch,%size,%embed_dim) @@ -1180,12 +1323,13 @@ pnnx.Output output 1 0 out )PNNXIR"; } - bool match(const std::map& captured_params) const + bool match(const std::map& matched_operators, const std::map& captured_params, const std::map& /*captured_attrs*/) const { const int embed_dim = captured_params.at("embed_dim").i; const int num_heads = captured_params.at("num_heads").i; const int feat_per_head = captured_params.at("feat_per_head").i; const float sqrt_feat_per_head = captured_params.at("sqrt_feat_per_head").f; + const int softmax_dim = captured_params.at("softmax_dim").i; if (embed_dim != num_heads * feat_per_head) return false; @@ -1193,6 +1337,10 @@ pnnx.Output output 1 0 out if (!NearlyEqual(sqrt_feat_per_head, sqrt(feat_per_head), 0.001)) return false; + int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size(); + if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1) + return false; + return true; } @@ -1237,7 +1385,7 @@ torch.transpose op_5 1 1 12 15 dim0=-2 dim1=-1 torch.matmul op_6 2 1 14 15 16 pnnx.Attribute attn_mask 0 1 attn_mask @data=(1,%num_heads,%size,%size)f32 pnnx.Expression op_7 2 1 16 attn_mask 18 expr=add(@0,@1) -F.softmax op_8 1 1 18 19 dim=-1 +F.softmax softmax 1 1 18 19 dim=%softmax_dim torch.matmul op_9 2 1 19 13 20 torch.transpose op_10 1 1 20 21 dim0=1 dim1=2 Tensor.reshape op_11 1 1 21 22 shape=(%batch,%size,%embed_dim) @@ -1281,6 +1429,33 @@ pnnx.Output output 1 0 out } }; +class fuse_multiheadattention_pass_17_1 : public fuse_multiheadattention_pass_17 +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +17 16 +pnnx.Input input_0 0 1 input +nn.Linear op_0 1 1 input 8 bias=%qkvbias in_features=%embed_dim out_features=%qkv_out_features @bias @weight +Tensor.reshape op_1 1 1 8 9 shape=(%batch,%size,3,%num_heads,%feat_per_head) +Tensor.permute op_2 1 1 9 10 dims=(2,0,3,1,4) +torch.unbind op_3 1 3 10 11 12 13 dim=0 +pnnx.Expression op_4 1 1 11 14 expr=mul(@0,%inv_sqrt_embed_dim_per_head) +Tensor.permute op_5 1 1 12 15 dims=(0,1,3,2) +torch.matmul op_6 2 1 14 15 16 +pnnx.Attribute attn_mask 0 1 attn_mask @data=(1,%num_heads,%size,%size)f32 +pnnx.Expression op_7 2 1 16 attn_mask 18 expr=add(@0,@1) +F.softmax softmax 1 1 18 19 dim=%softmax_dim +torch.matmul op_9 2 1 19 13 20 +Tensor.permute op_10 1 1 20 21 dims=(0,2,1,3) +Tensor.reshape op_11 1 1 21 22 shape=(%batch,%size,%embed_dim) +nn.Linear out_proj 1 1 22 out bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + class fuse_multiheadattention_pass_18 : public fuse_multiheadattention_pass { public: @@ -1302,7 +1477,7 @@ Tensor.view op_8 1 1 35 36 shape=(1,%batch,%num_heads,%size,% pnnx.Attribute op_9 0 1 37 @data=(1,%batch,1,%size,%size)f32 pnnx.Expression op_10 2 1 36 37 38 expr=add(@0,@1) Tensor.view op_11 1 1 38 39 shape=(%batch,%num_heads,%size,%size) -F.softmax op_12 1 1 39 40 dim=-1 +F.softmax softmax 1 1 39 40 dim=%softmax_dim torch.matmul op_13 2 1 40 30 41 torch.transpose op_14 1 1 41 42 dim0=1 dim1=2 Tensor.reshape op_15 1 1 42 43 shape=(%batch,%size,%embed_dim) @@ -1368,11 +1543,44 @@ pnnx.Output output 1 0 out } }; +class fuse_multiheadattention_pass_18_1 : public fuse_multiheadattention_pass_18 +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +20 19 +pnnx.Input input_0 0 1 input +nn.Linear op_0 1 1 input 25 bias=%qkvbias in_features=%embed_dim out_features=%qkv_out_features @bias @weight +Tensor.reshape op_1 1 1 25 26 shape=(%batch,%size,3,%num_heads,%feat_per_head) +Tensor.permute op_2 1 1 26 27 dims=(2,0,3,1,4) +torch.unbind op_3 1 3 27 28 29 30 dim=0 +pnnx.Expression op_4 1 1 28 31 expr=mul(@0,%inv_sqrt_embed_dim_per_head) +Tensor.permute op_5 1 1 29 32 dims=(0,1,3,2) +torch.matmul op_6 2 1 31 32 33 +pnnx.Attribute attn_mask 0 1 attn_mask @data=(1,%num_heads,%size,%size)f32 +pnnx.Expression op_7 2 1 33 attn_mask 35 expr=add(@0,@1) +Tensor.reshape op_8 1 1 35 36 shape=(1,%batch,%num_heads,%size,%size) +pnnx.Attribute op_9 0 1 37 @data=(1,%batch,1,%size,%size)f32 +pnnx.Expression op_10 2 1 36 37 38 expr=add(@0,@1) +Tensor.reshape op_11 1 1 38 39 shape=(%batch,%num_heads,%size,%size) +F.softmax softmax 1 1 39 40 dim=%softmax_dim +torch.matmul op_13 2 1 40 30 41 +Tensor.permute op_14 1 1 41 42 dims=(0,2,1,3) +Tensor.reshape op_15 1 1 42 43 shape=(%batch,%size,%embed_dim) +nn.Linear out_proj 1 1 43 out bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + void fuse_multiheadattention(Graph& graph) { #if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 9) fuse_multiheadattention_pass a; fuse_multiheadattention_pass_11 a1; + fuse_multiheadattention_pass_11_0 a2; + fuse_multiheadattention_pass_11_1 a3; fuse_multiheadattention_pass_sameqkv b; fuse_multiheadattention_pass_qkv c; fuse_multiheadattention_pass_q_samekv d; @@ -1395,11 +1603,15 @@ void fuse_multiheadattention(Graph& graph) fuse_multiheadattention_pass_16 o; fuse_multiheadattention_pass_16_1 o1; fuse_multiheadattention_pass_17 p; + fuse_multiheadattention_pass_17_1 p1; fuse_multiheadattention_pass_18 q; + fuse_multiheadattention_pass_18_1 q1; int opindex = 0; pnnx_graph_rewrite(graph, &a, opindex); pnnx_graph_rewrite(graph, &a1, opindex); + pnnx_graph_rewrite(graph, &a2, opindex); + pnnx_graph_rewrite(graph, &a3, opindex); pnnx_graph_rewrite(graph, &b, opindex); pnnx_graph_rewrite(graph, &c, opindex); pnnx_graph_rewrite(graph, &d, opindex); @@ -1422,7 +1634,9 @@ void fuse_multiheadattention(Graph& graph) pnnx_graph_rewrite(graph, &o, opindex); pnnx_graph_rewrite(graph, &o1, opindex); pnnx_graph_rewrite(graph, &p, opindex); + pnnx_graph_rewrite(graph, &p1, opindex); pnnx_graph_rewrite(graph, &q, opindex); + pnnx_graph_rewrite(graph, &q1, opindex); #endif } diff --git a/tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.cpp b/tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.cpp new file mode 100644 index 000000000..0495b083c --- /dev/null +++ b/tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.cpp @@ -0,0 +1,86 @@ +// 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 "fuse_slice_squeeze_to_select.h" + +#include "pass_level2.h" + +namespace pnnx { + +class fuse_slice_squeeze_to_select_pass : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input 0 1 input +Tensor.slice op_0 1 1 input a dim=%dim start=%start end=%end step=1 +torch.squeeze op_1 1 1 a out dim=%squeeze_dim +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "Tensor.select"; + } + + const char* name_str() const + { + return "select"; + } + + bool match(const std::map& captured_params) const + { + const int start = captured_params.at("start").i; + const int end = captured_params.at("end").i; + if (end != start + 1) + return false; + + const int dim = captured_params.at("dim").i; + int squeeze_dim; + if (captured_params.at("squeeze_dim").type == 2) + { + squeeze_dim = captured_params.at("squeeze_dim").i; + } + else // if (captured_params.at("squeeze_dim").type == 5) + { + if (captured_params.at("squeeze_dim").ai.size() != 1) + return false; + + squeeze_dim = captured_params.at("squeeze_dim").ai[0]; + } + if (squeeze_dim != dim) + return false; + + return true; + } + + void write(Operator* op, const std::map& captured_params) const + { + op->params["dim"] = captured_params.at("dim"); + op->params["index"] = captured_params.at("start"); + } +}; + +void fuse_slice_squeeze_to_select(Graph& graph) +{ + fuse_slice_squeeze_to_select_pass a; + int opindex = 0; + + pnnx_graph_rewrite(graph, &a, opindex); +} + +} // namespace pnnx diff --git a/tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.h b/tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.h new file mode 100644 index 000000000..562b5b052 --- /dev/null +++ b/tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.h @@ -0,0 +1,21 @@ +// Tencent is pleased to support the open source community by making ncnn available. +// +// Copyright (C) 2024 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 "ir.h" + +namespace pnnx { + +void fuse_slice_squeeze_to_select(Graph& graph); + +} // namespace pnnx diff --git a/tools/pnnx/src/pass_level5/fuse_slice_to_tensor_split.cpp b/tools/pnnx/src/pass_level5/fuse_slice_to_tensor_split.cpp index 4e767a02d..5b74636d1 100644 --- a/tools/pnnx/src/pass_level5/fuse_slice_to_tensor_split.cpp +++ b/tools/pnnx/src/pass_level5/fuse_slice_to_tensor_split.cpp @@ -35,19 +35,19 @@ void fuse_slice_to_tensor_split(Graph& graph) Operand* op_in = op->inputs[0]; - if (op->params.find("dims") == op->params.end() - || op->params.find("starts") == op->params.end() - || op->params.find("ends") == op->params.end() - || op->params.find("steps") == op->params.end()) + if ((!op->has_param("dim") && !op->has_param("dims")) + || (!op->has_param("start") && !op->has_param("starts")) + || (!op->has_param("end") && !op->has_param("ends")) + || (!op->has_param("step") && !op->has_param("steps"))) continue; - if (op->params.at("dims").ai.size() != 1) + if (!op->has_param("dim") && op->params.at("dims").ai.size() != 1) continue; - int dim = op->params.at("dims").ai[0]; - int start = op->params.at("starts").ai[0]; - int end = op->params.at("ends").ai[0]; - int step = op->params.at("steps").ai[0]; + int dim = op->has_param("dim") ? op->params.at("dim").i : op->params.at("dims").ai[0]; + int start = op->has_param("start") ? op->params.at("start").i : op->params.at("starts").ai[0]; + int end = op->has_param("end") ? op->params.at("end").i : op->params.at("ends").ai[0]; + int step = op->has_param("step") ? op->params.at("step").i : op->params.at("steps").ai[0]; if (start != 0 || step != 1) continue; @@ -74,18 +74,18 @@ void fuse_slice_to_tensor_split(Graph& graph) if (x->inputs[0] != op_in) continue; - if (x->params.find("dims") == x->params.end() - || x->params.find("starts") == x->params.end() - || x->params.find("ends") == x->params.end() - || x->params.find("steps") == x->params.end()) + if ((!x->has_param("dim") && !x->has_param("dims")) + || (!x->has_param("start") && !x->has_param("starts")) + || (!x->has_param("end") && !x->has_param("ends")) + || (!x->has_param("step") && !x->has_param("steps"))) continue; - if (x->params.at("dims").ai.size() != 1) + if (!x->has_param("dim") && x->params.at("dims").ai.size() != 1) continue; - int dim2 = x->params.at("dims").ai[0]; - int start2 = x->params.at("starts").ai[0]; - int step2 = x->params.at("steps").ai[0]; + int dim2 = x->has_param("dim") ? x->params.at("dim").i : x->params.at("dims").ai[0]; + int start2 = x->has_param("start") ? x->params.at("start").i : x->params.at("starts").ai[0]; + int step2 = x->has_param("step") ? x->params.at("step").i : x->params.at("steps").ai[0]; if (step2 != 1) continue; @@ -102,7 +102,7 @@ void fuse_slice_to_tensor_split(Graph& graph) if (std::find(graph.ops.begin(), graph.ops.end(), op2) < std::find(graph.ops.begin(), graph.ops.end(), cur)) cur = op2; - int end2 = op2->params.at("ends").ai[0]; + int end2 = op2->has_param("end") ? op2->params.at("end").i : op2->params.at("ends").ai[0]; if (end2 == INT_MAX) { slice_n_ops.push_back(op2); diff --git a/tools/pnnx/src/pass_level5/fuse_static_groupnorm.cpp b/tools/pnnx/src/pass_level5/fuse_static_groupnorm.cpp index da0d6112b..c0ab1f77c 100644 --- a/tools/pnnx/src/pass_level5/fuse_static_groupnorm.cpp +++ b/tools/pnnx/src/pass_level5/fuse_static_groupnorm.cpp @@ -41,7 +41,7 @@ pnnx.Output output 1 0 out return R"PNNXIR(7767517 3 2 pnnx.Input input 0 1 input -nn.GroupNorm group_norm 1 1 input out num_channels=%num_channels num_groups=%num_groups eps=%eps affine=True @weight=%op_weight.data @bias=%op_bias.data +nn.GroupNorm gn 1 1 input out num_channels=%num_channels num_groups=%num_groups eps=%eps affine=True @weight=%op_weight.data @bias=%op_bias.data pnnx.Output output 1 0 out )PNNXIR"; } diff --git a/tools/pnnx/src/pass_level5/fuse_static_instancenorm.cpp b/tools/pnnx/src/pass_level5/fuse_static_instancenorm.cpp index 76543b34c..86fc2c2dc 100644 --- a/tools/pnnx/src/pass_level5/fuse_static_instancenorm.cpp +++ b/tools/pnnx/src/pass_level5/fuse_static_instancenorm.cpp @@ -41,7 +41,7 @@ pnnx.Output output 1 0 out return R"PNNXIR(7767517 3 2 pnnx.Input input 0 1 input -nn.InstanceNorm1d instance_norm 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data +nn.InstanceNorm1d in 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data pnnx.Output output 1 0 out )PNNXIR"; } @@ -73,7 +73,7 @@ pnnx.Output output 1 0 out return R"PNNXIR(7767517 3 2 pnnx.Input input 0 1 input -nn.InstanceNorm2d instance_norm 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data +nn.InstanceNorm2d in 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data pnnx.Output output 1 0 out )PNNXIR"; } @@ -99,7 +99,7 @@ pnnx.Output output 1 0 out return R"PNNXIR(7767517 3 2 pnnx.Input input 0 1 input -nn.InstanceNorm3d instance_norm 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data +nn.InstanceNorm3d in 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data pnnx.Output output 1 0 out )PNNXIR"; } diff --git a/tools/pnnx/src/pass_level5/fuse_static_layernorm.cpp b/tools/pnnx/src/pass_level5/fuse_static_layernorm.cpp index 0b1f0dc41..2590bc1a0 100644 --- a/tools/pnnx/src/pass_level5/fuse_static_layernorm.cpp +++ b/tools/pnnx/src/pass_level5/fuse_static_layernorm.cpp @@ -41,7 +41,7 @@ pnnx.Output output 1 0 out return R"PNNXIR(7767517 3 2 pnnx.Input input 0 1 input -nn.LayerNorm layer_norm 1 1 input out normalized_shape=%normalized_shape eps=%eps elementwise_affine=True @weight=%op_weight.data @bias=%op_bias.data +nn.LayerNorm ln 1 1 input out normalized_shape=%normalized_shape eps=%eps elementwise_affine=True @weight=%op_weight.data @bias=%op_bias.data pnnx.Output output 1 0 out )PNNXIR"; } diff --git a/tools/pnnx/src/pass_level5/fuse_swish.cpp b/tools/pnnx/src/pass_level5/fuse_swish.cpp new file mode 100644 index 000000000..dc038de3d --- /dev/null +++ b/tools/pnnx/src/pass_level5/fuse_swish.cpp @@ -0,0 +1,71 @@ +// 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 "fuse_swish.h" + +#include "pass_level2.h" + +namespace pnnx { + +class fuse_swish_pass : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input 0 1 input +F.sigmoid op_0 1 1 input a +pnnx.Expression op_1 2 1 input a out expr=mul(@0,@1) +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* type_str() const + { + return "F.swish"; + } + + const char* name_str() const + { + return "swish"; + } +}; + +class fuse_swish_pass_1 : public fuse_swish_pass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +4 3 +pnnx.Input input 0 1 input +nn.Sigmoid op_0 1 1 input a +pnnx.Expression op_1 2 1 input a out expr=mul(@0,@1) +pnnx.Output output 1 0 out +)PNNXIR"; + } +}; + +void fuse_swish(Graph& graph) +{ + fuse_swish_pass a; + fuse_swish_pass_1 b; + int opindex = 0; + + pnnx_graph_rewrite(graph, &a, opindex); + pnnx_graph_rewrite(graph, &b, opindex); +} + +} // namespace pnnx diff --git a/tools/pnnx/src/pass_level5/fuse_swish.h b/tools/pnnx/src/pass_level5/fuse_swish.h new file mode 100644 index 000000000..7c31d6f98 --- /dev/null +++ b/tools/pnnx/src/pass_level5/fuse_swish.h @@ -0,0 +1,21 @@ +// Tencent is pleased to support the open source community by making ncnn available. +// +// Copyright (C) 2024 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 "ir.h" + +namespace pnnx { + +void fuse_swish(Graph& graph); + +} // namespace pnnx diff --git a/tools/pnnx/src/pass_onnx.cpp b/tools/pnnx/src/pass_onnx.cpp index 434a8182b..5579875bc 100644 --- a/tools/pnnx/src/pass_onnx.cpp +++ b/tools/pnnx/src/pass_onnx.cpp @@ -571,6 +571,62 @@ static void fuse_list_unpack(Graph& graph) } } +static void constant_unpooling(Graph& graph) +{ + while (1) + { + bool matched = false; + + for (size_t i = 0; i < graph.ops.size(); i++) + { + Operator* op = graph.ops[i]; + + if (op->type != "prim::Constant" && op->type != "pnnx.Expression" && op->type != "pnnx.Attribute") + continue; + + Operand* op_out = op->outputs[0]; + if (op_out->consumers.size() == 1) + continue; + + matched = true; + + // create shadow node for all consumers + for (size_t j = 1; j < op_out->consumers.size(); j++) + { + Operator* op1 = op_out->consumers[j]; + + Operator* op0 = graph.new_operator_before(op->type, op->name + "_pnnxshadow" + std::to_string(j), op1); + op0->inputnames = op->inputnames; + op0->params = op->params; + op0->attrs = op->attrs; + + Operand* op0_out = graph.new_operand(op_out->name + "_pnnxshadow" + std::to_string(j)); + op0_out->type = op_out->type; + op0_out->shape = op_out->shape; + op0_out->params = op_out->params; + + op0_out->producer = op0; + op0->outputs.push_back(op0_out); + + for (size_t k = 0; k < op1->inputs.size(); k++) + { + if (op1->inputs[k] == op_out) + op1->inputs[k] = op0_out; + } + + op0_out->consumers.push_back(op1); + } + + op_out->consumers.resize(1); + + break; + } + + if (!matched) + break; + } +} + void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph) { onnx2pnnx::OnnxModelProxy modelproxy(model); @@ -634,6 +690,11 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph) sim_op_type = "aten::split"; } + if (op_type == "Shape") + { + sim_op_type = "aten::size"; + } + // unaryop if (op_type == "Abs") sim_op_type = "aten::abs"; if (op_type == "Acos") sim_op_type = "aten::acos"; @@ -645,12 +706,14 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph) if (op_type == "Ceil") sim_op_type = "aten::ceil"; if (op_type == "Cos") sim_op_type = "aten::cos"; if (op_type == "Cosh") sim_op_type = "aten::cosh"; + if (op_type == "Erf") sim_op_type = "aten::erf"; if (op_type == "Exp") sim_op_type = "aten::exp"; if (op_type == "Floor") sim_op_type = "aten::floor"; if (op_type == "Log") sim_op_type = "aten::log"; if (op_type == "Neg") sim_op_type = "aten::neg"; if (op_type == "Reciprocal") sim_op_type = "aten::reciprocal"; if (op_type == "Round") sim_op_type = "aten::round"; + if (op_type == "Sigmoid") sim_op_type = "aten::sigmoid"; if (op_type == "Sign") sim_op_type = "aten::sign"; if (op_type == "Sin") sim_op_type = "aten::sin"; if (op_type == "Sinh") sim_op_type = "aten::sinh"; @@ -666,6 +729,9 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph) if (op_type == "Max") sim_op_type = "aten::max"; if (op_type == "Min") sim_op_type = "aten::min"; if (op_type == "Pow") sim_op_type = "aten::pow"; + + // trinaryop + if (op_type == "Clip") sim_op_type = "aten::clamp"; } else if (string_starts_with(op_type, "aten_")) { @@ -716,7 +782,9 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph) if (input.empty()) continue; - if (modelproxy.has_initializer(input)) + Operand* op_in = pnnx_graph.get_operand(input); + + if (!op_in && modelproxy.has_initializer(input)) { // skip function weight if (is_function_op) @@ -733,6 +801,9 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph) if (sim_op_type == "Reshape" && j == 1) is_attr_list = true; + if (sim_op_type == "Pad" && j == 1) + is_attr_list = true; + if (sim_op_type == "ReduceMean" && j == 1) is_attr_list = true; } @@ -893,9 +964,9 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph) op_const->attrs["data"] = tensor; } - } - Operand* op_in = pnnx_graph.get_operand(input); + op_in = pnnx_graph.get_operand(input); + } op_in->consumers.push_back(op); op->inputs.push_back(op_in); @@ -1039,6 +1110,8 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph) // post process fuse_list_unpack(pnnx_graph); + + constant_unpooling(pnnx_graph); } } // namespace pnnx diff --git a/tools/pnnx/src/pass_onnx/eliminate_noop.cpp b/tools/pnnx/src/pass_onnx/eliminate_noop.cpp index cf011f9c2..1a922480c 100644 --- a/tools/pnnx/src/pass_onnx/eliminate_noop.cpp +++ b/tools/pnnx/src/pass_onnx/eliminate_noop.cpp @@ -33,7 +33,7 @@ void eliminate_noop(onnx::ModelProto& model) const onnx::NodeProto& node = graph->node(i); const std::string& op_type = node.op_type(); - if (op_type == "Identity" || op_type == "aten_copy") + if (op_type == "Identity" || op_type == "Dropout" || op_type == "aten_copy") { const std::string& input_name = node.input(0); const std::string& output_name = node.output(0); diff --git a/tools/pnnx/src/pass_onnx/fold_constants.cpp b/tools/pnnx/src/pass_onnx/fold_constants.cpp index cb5b3a75d..5978ba066 100644 --- a/tools/pnnx/src/pass_onnx/fold_constants.cpp +++ b/tools/pnnx/src/pass_onnx/fold_constants.cpp @@ -62,8 +62,36 @@ static size_t sizeof_onnx_datatype(ONNXTensorElementDataType type) return 0; } -void fold_constants(onnx::ModelProto& model) +static ONNXTensorElementDataType get_onnx_tensor_elem_data_type(const std::string& type) { + if (type == "i8") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8; + if (type == "u8") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8; + if (type == "i16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16; + if (type == "u16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16; + if (type == "i32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; + if (type == "u32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32; + if (type == "i64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; + if (type == "u64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64; + if (type == "f16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; + if (type == "f32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; + if (type == "f64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; + if (type == "bf16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16; + if (type == "c64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64; + if (type == "c128") return ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128; + + // unknown + fprintf(stderr, "unsupported tensor elem data type %s\n", type.c_str()); + return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; +} + +void fold_constants(onnx::ModelProto& model, + const std::vector >& input_shapes, + const std::vector& input_types, + const std::vector >& input_shapes2, + const std::vector& input_types2) +{ + bool ignore_aten_size = input_shapes2.empty(); + // collect initializers std::unordered_set initializers; { @@ -114,13 +142,13 @@ void fold_constants(onnx::ModelProto& model) || op_type == "aten_ones_like" || op_type == "aten_zeros_like") { - is_outputs_foldable = true; + is_outputs_foldable = ignore_aten_size; } // TODO whitelist for static shape if (op_type == "Shape") { - is_outputs_foldable = true; + is_outputs_foldable = ignore_aten_size; } // TODO whitelist for static type @@ -263,14 +291,22 @@ void fold_constants(onnx::ModelProto& model) const onnx::ValueInfoProto& value = graph->input(i); std::vector shape; - const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape(); - for (int k = 0; k < tsp.dim_size(); k++) + ONNXTensorElementDataType datatype; + if (!input_shapes.empty()) { - // TODO has_dim_value ? - shape.push_back(tsp.dim(k).dim_value()); + shape = input_shapes[i]; + datatype = get_onnx_tensor_elem_data_type(input_types[i]); } + else + { + const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape(); + for (int k = 0; k < tsp.dim_size(); k++) + { + shape.push_back(tsp.dim(k).dim_value()); + } - ONNXTensorElementDataType datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type(); + datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type(); + } OrtValue* ort_val = 0; ort_status = ort_api->CreateTensorAsOrtValue(ort_allocator, (const int64_t*)shape.data(), shape.size(), datatype, &ort_val); diff --git a/tools/pnnx/src/pass_onnx/fold_constants.h b/tools/pnnx/src/pass_onnx/fold_constants.h index 9728a7aab..652025215 100644 --- a/tools/pnnx/src/pass_onnx/fold_constants.h +++ b/tools/pnnx/src/pass_onnx/fold_constants.h @@ -18,7 +18,11 @@ namespace pnnx { namespace onnx2pnnx { -void fold_constants(onnx::ModelProto& model); +void fold_constants(onnx::ModelProto& model, + const std::vector >& input_shapes, + const std::vector& input_types, + const std::vector >& input_shapes2, + const std::vector& input_types2); } // namespace onnx2pnnx diff --git a/tools/pnnx/src/pass_onnx/shape_inference.cpp b/tools/pnnx/src/pass_onnx/shape_inference.cpp index fb72e6b85..a2aff34a7 100644 --- a/tools/pnnx/src/pass_onnx/shape_inference.cpp +++ b/tools/pnnx/src/pass_onnx/shape_inference.cpp @@ -24,12 +24,38 @@ namespace pnnx { namespace onnx2pnnx { +static ONNXTensorElementDataType get_onnx_tensor_elem_data_type(const std::string& type) +{ + if (type == "i8") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8; + if (type == "u8") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8; + if (type == "i16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16; + if (type == "u16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16; + if (type == "i32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32; + if (type == "u32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32; + if (type == "i64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; + if (type == "u64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64; + if (type == "f16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; + if (type == "f32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; + if (type == "f64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE; + if (type == "bf16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16; + if (type == "c64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64; + if (type == "c128") return ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128; + + // unknown + fprintf(stderr, "unsupported tensor elem data type %s\n", type.c_str()); + return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; +} + static bool string_starts_with(const std::string& s, const std::string& s2) { return strncmp(s.c_str(), s2.c_str(), s2.size()) == 0; } -void shape_inference(onnx::ModelProto& model) +void shape_inference(onnx::ModelProto& model, + const std::vector >& input_shapes, + const std::vector& input_types, + const std::vector >& input_shapes2, + const std::vector& input_types2) { onnx::GraphProto* graph = model.mutable_graph(); @@ -166,14 +192,22 @@ void shape_inference(onnx::ModelProto& model) const onnx::ValueInfoProto& value = graph->input(i); std::vector shape; - const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape(); - for (int k = 0; k < tsp.dim_size(); k++) + ONNXTensorElementDataType datatype; + if (!input_shapes.empty()) { - // TODO has_dim_value ? - shape.push_back(tsp.dim(k).dim_value()); + shape = input_shapes[i]; + datatype = get_onnx_tensor_elem_data_type(input_types[i]); } + else + { + const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape(); + for (int k = 0; k < tsp.dim_size(); k++) + { + shape.push_back(tsp.dim(k).dim_value()); + } - ONNXTensorElementDataType datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type(); + datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type(); + } OrtValue* ort_val = 0; ort_status = ort_api->CreateTensorAsOrtValue(ort_allocator, (const int64_t*)shape.data(), shape.size(), datatype, &ort_val); @@ -333,6 +367,232 @@ void shape_inference(onnx::ModelProto& model) ort_api->ReleaseSessionOptions(ort_session_opt); ort_api->ReleaseEnv(ort_env); } + + // onnxrt inference for input_shapes2 + if (!input_shapes2.empty()) + { + const OrtApi* ort_api = OrtGetApiBase()->GetApi(ORT_API_VERSION); + + OrtStatus* ort_status = 0; + + OrtEnv* ort_env = 0; + ort_status = ort_api->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "pnnx", &ort_env); + if (ort_status) + { + fprintf(stderr, "ort CreateEnv failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + OrtSessionOptions* ort_session_opt = 0; + ort_status = ort_api->CreateSessionOptions(&ort_session_opt); + if (ort_status) + { + fprintf(stderr, "ort CreateSessionOptions failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + ort_status = ort_api->SetSessionGraphOptimizationLevel(ort_session_opt, ORT_DISABLE_ALL); + if (ort_status) + { + fprintf(stderr, "ort SetSessionGraphOptimizationLevel failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + // ort_status = ort_api->SetIntraOpNumThreads(ort_session_opt, 4); + // if (ort_status) + // { + // fprintf(stderr, "ort SetIntraOpNumThreads failed %s\n", ort_api->GetErrorMessage(ort_status)); + // } + // + // ort_status = ort_api->SetInterOpNumThreads(ort_session_opt, 4); + // if (ort_status) + // { + // fprintf(stderr, "ort SetInterOpNumThreads failed %s\n", ort_api->GetErrorMessage(ort_status)); + // } + + OrtSession* ort_session = 0; + ort_status = ort_api->CreateSessionFromArray(ort_env, (const void*)tmp_onnx_data.data(), tmp_onnx_data.size(), ort_session_opt, &ort_session); + if (ort_status) + { + fprintf(stderr, "ort CreateSession failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + OrtRunOptions* ort_run_opt = 0; + ort_status = ort_api->CreateRunOptions(&ort_run_opt); + if (ort_status) + { + fprintf(stderr, "ort CreateRunOptions failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + OrtAllocator* ort_allocator = 0; + ort_status = ort_api->GetAllocatorWithDefaultOptions(&ort_allocator); + if (ort_status) + { + fprintf(stderr, "ort GetAllocatorWithDefaultOptions failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + std::vector input_names; + std::vector inputs; + for (int i = 0; i < graph->input_size(); i++) + { + const onnx::ValueInfoProto& value = graph->input(i); + + std::vector shape = input_shapes2[i]; + ONNXTensorElementDataType datatype = get_onnx_tensor_elem_data_type(input_types2[i]); + + OrtValue* ort_val = 0; + ort_status = ort_api->CreateTensorAsOrtValue(ort_allocator, (const int64_t*)shape.data(), shape.size(), datatype, &ort_val); + if (ort_status) + { + fprintf(stderr, "ort CreateTensorAsOrtValue failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + input_names.push_back(value.name().c_str()); + inputs.push_back(ort_val); + } + + std::vector output_names; + std::vector outputs; + for (size_t i = 0; i < intermediates.size(); i++) + { + output_names.push_back(intermediates[i].c_str()); + outputs.push_back(0); + } + + ort_status = ort_api->Run(ort_session, ort_run_opt, + input_names.data(), inputs.data(), input_names.size(), + output_names.data(), output_names.size(), outputs.data()); + if (ort_status) + { + fprintf(stderr, "ort Run failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + // TODO get output data + + // graph->clear_output(); + + for (size_t i = 0; i < output_names.size(); i++) + { + OrtTypeInfo* type_info = 0; + ort_status = ort_api->GetTypeInfo(outputs[i], &type_info); + if (ort_status) + { + fprintf(stderr, "ort GetTypeInfo failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + ONNXType type = ONNX_TYPE_UNKNOWN; + if (type_info) + { + ort_status = ort_api->GetOnnxTypeFromTypeInfo(type_info, &type); + if (ort_status) + { + fprintf(stderr, "ort GetOnnxTypeFromTypeInfo failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + } + + if (type == ONNX_TYPE_TENSOR) + { + OrtTensorTypeAndShapeInfo* info = 0; + ort_status = ort_api->GetTensorTypeAndShape(outputs[i], &info); + if (ort_status) + { + fprintf(stderr, "ort GetTensorTypeAndShape failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + ONNXTensorElementDataType datatype = ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; + ort_status = ort_api->GetTensorElementType(info, &datatype); + if (ort_status) + { + fprintf(stderr, "ort GetTensorElementType failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + size_t out_dims = 0; + ort_status = ort_api->GetDimensionsCount(info, &out_dims); + if (ort_status) + { + fprintf(stderr, "ort GetDimensionsCount failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + // fprintf(stderr, " out_dims = %lu\n", out_dims); + + std::vector out_shape; + out_shape.resize(out_dims); + ort_status = ort_api->GetDimensions(info, out_shape.data(), out_dims); + if (ort_status) + { + fprintf(stderr, "ort GetDimensions failed %s\n", ort_api->GetErrorMessage(ort_status)); + } + + // fprintf(stderr, "%16s = ", output_names[i]); + // for (size_t j = 0; j < out_dims; j++) + // { + // fprintf(stderr, "%lu ", out_shape[j]); + // } + // fprintf(stderr, "\n"); + + // assign value info + { + onnx::ValueInfoProto* value = 0; + + // maybe output + for (size_t j = 0; j < orig_outputs.size(); j++) + { + if (orig_outputs[j] == output_names[i]) + { + value = graph->mutable_output(j); + break; + } + } + if (!value) + { + for (int j = 0; j < graph->value_info_size(); j++) + { + if (graph->mutable_value_info(j)->name() == output_names[i]) + { + value = graph->mutable_value_info(j); + break; + } + } + } + + // fprintf(stderr, "assign value info2 %s\n", value->name().c_str()); + + value->mutable_type()->mutable_tensor_type()->set_elem_type((int32_t)datatype); + + onnx::TensorShapeProto* tsp = value->mutable_type()->mutable_tensor_type()->mutable_shape(); + + // tsp->clear_dim(); + for (size_t j = 0; j < out_dims; j++) + { + if (tsp->dim(j).dim_value() == out_shape[j]) + continue; + + // dynamic dim size + tsp->mutable_dim(j)->clear_dim_value(); + } + } + + ort_api->ReleaseTensorTypeAndShapeInfo(info); + } + + if (type_info) + { + ort_api->ReleaseTypeInfo(type_info); + } + } + + for (size_t i = 0; i < input_names.size(); i++) + { + ort_api->ReleaseValue(inputs[i]); + } + + for (size_t i = 0; i < output_names.size(); i++) + { + ort_api->ReleaseValue(outputs[i]); + } + + ort_api->ReleaseRunOptions(ort_run_opt); + ort_api->ReleaseSession(ort_session); + ort_api->ReleaseSessionOptions(ort_session_opt); + ort_api->ReleaseEnv(ort_env); + } } } // namespace onnx2pnnx diff --git a/tools/pnnx/src/pass_onnx/shape_inference.h b/tools/pnnx/src/pass_onnx/shape_inference.h index ea8733345..b4cd657bb 100644 --- a/tools/pnnx/src/pass_onnx/shape_inference.h +++ b/tools/pnnx/src/pass_onnx/shape_inference.h @@ -18,7 +18,11 @@ namespace pnnx { namespace onnx2pnnx { -void shape_inference(onnx::ModelProto& model); +void shape_inference(onnx::ModelProto& model, + const std::vector >& input_shapes, + const std::vector& input_types, + const std::vector >& input_shapes2, + const std::vector& input_types2); } // namespace onnx2pnnx diff --git a/tools/pnnx/tests/onnx/CMakeLists.txt b/tools/pnnx/tests/onnx/CMakeLists.txt index b59c582a6..274fb4793 100644 --- a/tools/pnnx/tests/onnx/CMakeLists.txt +++ b/tools/pnnx/tests/onnx/CMakeLists.txt @@ -13,3 +13,5 @@ pnnx_onnx_add_test(mobilenet_v3_small) pnnx_onnx_add_test(resnet18) pnnx_onnx_add_test(shufflenet_v2_x1_0) pnnx_onnx_add_test(squeezenet1_1) +pnnx_onnx_add_test(swin_t) +pnnx_onnx_add_test(vit_b_32) diff --git a/tools/pnnx/tests/onnx/test_swin_t.py b/tools/pnnx/tests/onnx/test_swin_t.py new file mode 100644 index 000000000..be25520d0 --- /dev/null +++ b/tools/pnnx/tests/onnx/test_swin_t.py @@ -0,0 +1,65 @@ +# Tencent is pleased to support the open source community by making ncnn available. +# +# Copyright (C) 2024 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 torchvision +import torchvision.models as models +from packaging import version + +def test(): + if version.parse(torchvision.__version__) < version.parse('0.13'): + return True + + net = models.swin_t() + net.eval() + + torch.manual_seed(0) + x = torch.rand(1, 3, 224, 224) + + a = net(x) + + # export onnx + torch.onnx.export(net, (x,), "test_swin_t.onnx") + + # onnx to pnnx + import os + os.system("../../src/pnnx test_swin_t.onnx inputshape=[1,3,224,224]") + + # pnnx inference + import test_swin_t_pnnx + b = test_swin_t_pnnx.test_inference() + + if not torch.allclose(a, b, 1e-4, 1e-4): + return False + + if version.parse(torch.__version__) < version.parse('2.4'): + return True + + # export dynamo onnx + torch.onnx.dynamo_export(net, x).save("test_swin_t_dynamo.onnx") + + # onnx to pnnx + os.system("../../src/pnnx test_swin_t_dynamo.onnx inputshape=[1,3,224,224]") + + # pnnx inference + import test_swin_t_dynamo_pnnx + b = test_swin_t_dynamo_pnnx.test_inference() + + return torch.allclose(a, b, 1e-4, 1e-4) + +if __name__ == "__main__": + if test(): + exit(0) + else: + exit(1) diff --git a/tools/pnnx/tests/onnx/test_vit_b_32.py b/tools/pnnx/tests/onnx/test_vit_b_32.py new file mode 100644 index 000000000..ecb0bd350 --- /dev/null +++ b/tools/pnnx/tests/onnx/test_vit_b_32.py @@ -0,0 +1,68 @@ +# Tencent is pleased to support the open source community by making ncnn available. +# +# Copyright (C) 2024 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 torchvision +import torchvision.models as models +from packaging import version + +def test(): + if version.parse(torchvision.__version__) < version.parse('0.12'): + return True + + if version.parse(torch.__version__) >= version.parse('2.0') and version.parse(torch.__version__) < version.parse('2.1'): + return True + + net = models.vit_b_32() + net.eval() + + torch.manual_seed(0) + x = torch.rand(1, 3, 224, 224) + + a = net(x) + + # export onnx + torch.onnx.export(net, (x,), "test_vit_b_32.onnx") + + # onnx to pnnx + import os + os.system("../../src/pnnx test_vit_b_32.onnx inputshape=[1,3,224,224]") + + # pnnx inference + import test_vit_b_32_pnnx + b = test_vit_b_32_pnnx.test_inference() + + if not torch.allclose(a, b, 1e-4, 1e-4): + return False + + if version.parse(torch.__version__) < version.parse('2.4'): + return True + + # export dynamo onnx + torch.onnx.dynamo_export(net, x).save("test_vit_b_32_dynamo.onnx") + + # onnx to pnnx + os.system("../../src/pnnx test_vit_b_32_dynamo.onnx inputshape=[1,3,224,224]") + + # pnnx inference + import test_vit_b_32_dynamo_pnnx + b = test_vit_b_32_dynamo_pnnx.test_inference() + + return torch.allclose(a, b, 1e-4, 1e-4) + +if __name__ == "__main__": + if test(): + exit(0) + else: + exit(1)