From 9be93dd7a5d63fe8249e974184c09cbaec45756a Mon Sep 17 00:00:00 2001 From: kkannan Date: Wed, 11 Dec 2024 15:25:44 +0000 Subject: [PATCH] Add support for the reflect padding mode in Conv2D operation --- python/tvm/relay/op/contrib/forge/forge_passes.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/op/contrib/forge/forge_passes.py b/python/tvm/relay/op/contrib/forge/forge_passes.py index 45f811aa5..fd6533460 100644 --- a/python/tvm/relay/op/contrib/forge/forge_passes.py +++ b/python/tvm/relay/op/contrib/forge/forge_passes.py @@ -156,9 +156,16 @@ def callback(self, pre, post, node_map): pad_width = pad.attrs.pad_width top_pad, bottom_pad = pad_width[-2] left_pad, right_pad = pad_width[-1] + + pad_mode = pad.attrs.pad_mode + + if pad_mode == "constant": + padding = [top_pad, left_pad, bottom_pad, right_pad] - padding = [top_pad, left_pad, bottom_pad, right_pad] - + if pad_mode == "reflect": + act = tvm.relay.op.nn.pad(act, pad_width, pad_mode="reflect") + padding = [0,0,0,0] + op_attrs = {**conv_pool.attrs} op_attrs["padding"] = padding