Support fuse add into ConvTranspose.

onnx · Nov 22, 2022 · ded269d · ded269d
1 parent 74fdf9c
commit ded269d
Showing 1 changed file with 15 additions and 3 deletions.
diff --git a/onnxoptimizer/passes/fuse_add_bias_into_conv.h b/onnxoptimizer/passes/fuse_add_bias_into_conv.h
@@ -32,10 +32,21 @@ struct FuseAddBiasIntoConv final : public PredicateBasedPass {
   std::string getPassName() const override {
     return "fuse_add_bias_into_conv";
   }
-  bool patternMatchPredicate(Node *node) override {
+
+  inline bool matchConvAdd(Node *node) {
     return node->kind() == kAdd && node->inputs()[0]->node()->kind() == kConv &&
            node->inputs()[0]->node()->inputs().size() == 2;
   }
+
+  inline bool matchConvTransposeAdd(Node *node) {
+    return node->kind() == kAdd && node->inputs()[0]->node()->kind() == kConvTranspose &&
+           node->inputs()[0]->node()->inputs().size() == 2;
+  }
+
+  bool patternMatchPredicate(Node *node) override {
+    return matchConvAdd(node) || matchConvTransposeAdd(node);
+  }
+
   static Node *makeSqueezeOrUnsqueeze(Graph &graph, std::vector<int64_t> &axes,
                                       Value *input, Node *target_node,
                                       BuiltinSymbol k) {
@@ -61,6 +72,7 @@ struct FuseAddBiasIntoConv final : public PredicateBasedPass {
                     NodeDestroyType &destroy_current) override {
     // due to current broadcasting's constraint, Conv has to be the first
     // operand
+    const bool is_conv = matchConvAdd(n);
     destroy_current = NodeDestroyType::DestroyZero;
     auto orig_conv = n->inputs()[0];
     auto orig_bias = n->inputs()[1];
@@ -85,8 +97,8 @@ struct FuseAddBiasIntoConv final : public PredicateBasedPass {
     }
     // try to get feature M and rank from weight_shape
     if (weight_shape.size() > 0 && weight_shape[0].is_int) {
-      ONNX_ASSERT(M == -1 || M == weight_shape[0].dim);
-      M = weight_shape[0].dim;
+      ONNX_ASSERT(M == -1 || M == weight_shape[0].dim || M == weight_shape[1].dim);
+      M = is_conv ? weight_shape[0].dim : weight_shape[1].dim;
       ONNX_ASSERT(rank == -1 ||
                   rank == static_cast<int64_t>(weight_shape.size()));
       rank = weight_shape.size();