From 402001128d0ba13f8ebd7c590924e2b2d97c5676 Mon Sep 17 00:00:00 2001 From: nihui Date: Mon, 25 Jul 2022 14:04:41 +0800 Subject: [PATCH] eliminate math expression with torch.zeros/ones input (#4083) --- tools/pnnx/src/pass_level3/eliminate_noop_math.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tools/pnnx/src/pass_level3/eliminate_noop_math.cpp b/tools/pnnx/src/pass_level3/eliminate_noop_math.cpp index 75bccda89..310fbda32 100644 --- a/tools/pnnx/src/pass_level3/eliminate_noop_math.cpp +++ b/tools/pnnx/src/pass_level3/eliminate_noop_math.cpp @@ -130,6 +130,12 @@ static bool operator_is_all_constant(const Operator* op, float vf, int vi) if (op->type == "prim::Constant") return constant_is_all_constant(op, vf, vi); + if (op->type == "torch.zeros" || op->type == "torch.zeros_like") + return (vf == 0.f && vi == 0); + + if (op->type == "torch.ones" || op->type == "torch.ones_like") + return (vf == 1.f && vi == 1); + return false; } @@ -163,7 +169,7 @@ void eliminate_noop_math(Graph& graph) need_eliminate = true; identity_input_id = 0; } - else if (operator_is_all_constant(op0, 0.f, 0) && operator_is_all_constant(op0, 1.f, 1)) + else if (operator_is_all_constant(op0, 0.f, 0) && operator_is_all_constant(op2, 1.f, 1)) { // x <= 0 + b * 1 need_eliminate = true;