diff --git a/tools/onnx/onnx2ncnn.cpp b/tools/onnx/onnx2ncnn.cpp index d47a2eb80..77e89db1a 100644 --- a/tools/onnx/onnx2ncnn.cpp +++ b/tools/onnx/onnx2ncnn.cpp @@ -495,9 +495,6 @@ static void fuse_hardswish(onnx::GraphProto* mutable_graph, std::mapop_type() != "Mul") continue; - if (node_reference.find(node2->output(0)) == node_reference.end() || node_reference[node2->output(0)] != 1) - continue; - if (node2->input(0) != node->input(0) || node2->input(1) != node->output(0)) continue; @@ -635,9 +632,6 @@ static void fuse_batchnorm1d_squeeze_unsqueeze(onnx::GraphProto* mutable_graph, if (node_reference.find(node2->output(0)) == node_reference.end() || node_reference[node2->output(0)] != 1) continue; - if (node_reference.find(node3->output(0)) == node_reference.end() || node_reference[node3->output(0)] != 1) - continue; - if (node2->input(0) != node->output(0) || node3->input(0) != node2->output(0)) continue; @@ -659,6 +653,59 @@ static void fuse_batchnorm1d_squeeze_unsqueeze(onnx::GraphProto* mutable_graph, } } +static void fuse_unsqueeze_prelu(onnx::GraphProto* mutable_graph, std::map& weights, std::map& binaryop_weights, std::map& node_reference, std::set& blob_names, int& reduced_node_count, std::vector& reduced_binaryop_weights) +{ + int node_count = mutable_graph->node_size(); + for (int i=0; imutable_node(i); + + // PReLU <= Unsqueeze - PReLU + if (node->op_type() == "Unsqueeze") + { + // check weight + if (weights.find(node->input(0)) == weights.end()) + continue; + + onnx::TensorProto& B = weights[node->input(0)]; + if (B.dims_size() != 1) + continue; + + if (node_reference.find(node->output(0)) == node_reference.end() || node_reference[node->output(0)] != 1) + continue; + + // axes = (1, 2) + std::vector axes = get_node_attr_ai(*node, "axes"); + if (axes.size() != 2) + continue; + if (axes[0] != 1 || axes[1] != 2) + continue; + + if (i+1 >= node_count) + continue; + + onnx::NodeProto* node2 = mutable_graph->mutable_node(i+1); + + if (node2->op_type() != "PRelu") + continue; + + if (node2->input(1) != node->output(0)) + continue; + + // reduce + node->set_op_type("noop_reducedncnn"); + + node_reference.erase(node_reference.find(node->output(0))); + blob_names.erase(node->output(0)); + + node2->set_input(1, node->input(0)); + + reduced_node_count += 1; + i += 1; + } + } +} + int main(int argc, char** argv) { const char* onnxpb = argv[1]; @@ -852,6 +899,7 @@ int main(int argc, char** argv) fuse_hardsigmoid (mutable_graph, weights, binaryop_weights, node_reference, blob_names, reduced_node_count, reduced_binaryop_weights); fuse_hardswish (mutable_graph, weights, binaryop_weights, node_reference, blob_names, reduced_node_count, reduced_binaryop_weights); fuse_batchnorm1d_squeeze_unsqueeze(mutable_graph, weights, binaryop_weights, node_reference, blob_names, reduced_node_count, reduced_binaryop_weights); + fuse_unsqueeze_prelu(mutable_graph, weights, binaryop_weights, node_reference, blob_names, reduced_node_count, reduced_binaryop_weights); // remove node_reference entry with reference equals to one int splitncnn_blob_count = 0;