【SCU】【Paddle TensorRT No.18】Add pd_op.conv3d and pd_op.conv3d_transpose converter (#70647)

PolaKuma · web-flow · commit 7756eed82a27 · 2025-02-24T16:29:26.000+08:00
* add_conv3d

* fix codestyle

* fix codestyle

* fix codestyle

* delete fp16

* update test

* update

* Update test_converter_conv.py

* Update test_converter_conv.py

* Update test_converter_conv.py

* Update test_converter_conv.py

* fix codestyle

* update

* fix codestyle

* Update test_converter_conv.py

* Update trt_op_marker_pass.cc
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -596,6 +596,91 @@ class DepthwiseConv2dTransposeOpPattern
   }
 };
 
+class Conv3dTransposeOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::Conv3dTransposeOp> {
+ public:
+  using pir::OpRewritePattern<
+      paddle::dialect::Conv3dTransposeOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(paddle::dialect::Conv3dTransposeOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    auto paddings_attr = op->attribute<pir::ArrayAttribute>("paddings");
+    std::vector<int32_t> paddings;
+    for (const auto &attr : paddings_attr.AsVector()) {
+      paddings.push_back(attr.dyn_cast<pir::Int32Attribute>().data());
+    }
+    if (paddings.size() > 3) {
+      VLOG(3) << "In conv3d_transpose, paddings size must be less than or "
+                 "equal to 3";
+      return false;
+    }
+    if (!op->HasAttribute("dilations")) {
+      VLOG(3) << "In conv3d_transpose, dilations attribute does not exist";
+      return false;
+    } else {
+      auto dilation_attr = op->attribute<pir::ArrayAttribute>("dilations");
+      std::vector<int32_t> dilations;
+      for (const auto &attr : dilation_attr.AsVector()) {
+        dilations.push_back(attr.dyn_cast<pir::Int32Attribute>().data());
+      }
+      if (dilations[0] != 1 || dilations[1] != 1 || dilations[2] != 1) {
+        VLOG(3) << "In conv3d_transpose, Dilations must be (1, 1, 1) for "
+                   "tensorRT, but given ("
+                << dilations[0] << ", " << dilations[1] << ", " << dilations[2]
+                << ")";
+        return false;
+      }
+    }
+    pir::Value filter = op.operand_source(1);
+    auto filter_type =
+        filter.type().dyn_cast<paddle::dialect::DenseTensorType>();
+    auto filter_shape = filter_type.dims();
+    if (filter_shape.size() != 5) {
+      VLOG(3) << "The conv3d filter's dims size should be 5";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
+class Conv3dOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::Conv3dOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::Conv3dOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(paddle::dialect::Conv3dOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    auto paddings_attr = op->attribute<pir::ArrayAttribute>("paddings");
+    std::vector<int32_t> paddings;
+    for (const auto &attr : paddings_attr.AsVector()) {
+      paddings.push_back(attr.dyn_cast<pir::Int32Attribute>().data());
+    }
+    if (paddings.size() > 3) {
+      VLOG(3) << "In conv3d, paddings size must be less than or equal to 3";
+      return false;
+    }
+    pir::Value filter = op.operand_source(1);
+    auto filter_type =
+        filter.type().dyn_cast<paddle::dialect::DenseTensorType>();
+    auto filter_shape = filter_type.dims();
+    if (filter_shape.size() != 5) {
+      VLOG(3) << "The conv3d filter's dims size should be 5";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class DeformableConvOpPattern
     : public pir::OpRewritePattern<paddle::dialect::DeformableConvOp> {
  public:
@@ -2971,6 +3056,8 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<NotEqualOpPattern>(context));
     ps.Add(std::make_unique<LogicalXorOpPattern>(context));
     ps.Add(std::make_unique<CeluOpPattern>(context));
+    ps.Add(std::make_unique<Conv3dOpPattern>(context));
+    ps.Add(std::make_unique<Conv3dTransposeOpPattern>(context));
     ps.Add(std::make_unique<OneHotOpPattern>(context));
     ps.Add(std::make_unique<PadOpPattern>(context));
     ps.Add(std::make_unique<TemporalShiftOpPattern>(context));
diff --git a/python/paddle/tensorrt/converter_utils.py b/python/paddle/tensorrt/converter_utils.py
@@ -563,6 +563,74 @@ def convert_conv2d(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
+def convert_conv3d(network, paddle_op, inputs):
+    input_tensor, filter = inputs
+    filter_shape = paddle_op.operands()[1].source().shape
+
+    n_output = filter_shape[0]
+    n_input = filter_shape[1]
+    filter_d = filter_shape[2]
+    filter_h = filter_shape[3]
+    filter_w = filter_shape[4]
+
+    groups = paddle_op.attrs().get("groups", 1)
+    dilations = paddle_op.attrs().get("dilations", [1, 1, 1])
+    strides = paddle_op.attrs().get("strides", [1, 1, 1])
+    paddings = paddle_op.attrs().get("paddings", [0, 0, 0])
+    padding_algorithm = paddle_op.attrs().get("padding_algorithm", "EXPLICIT")
+    # for conv3d_transpose
+    output_padding = paddle_op.attrs().get("output_padding", [])
+
+    nv_ksize = trt.Dims3(filter_d, filter_h, filter_w)
+    nv_dilations = trt.Dims3(dilations[0], dilations[1], dilations[2])
+    nv_strides = trt.Dims3(strides[0], strides[1], strides[2])
+    nv_pre_paddings = trt.Dims3(paddings[0], paddings[1], paddings[2])
+
+    if paddle_op.name() == "pd_op.conv3d":
+        layer = network.add_convolution_nd(
+            input=input_tensor,
+            num_output_maps=n_output,
+            kernel_shape=nv_ksize,
+            kernel=filter,
+            bias=None,
+        )
+    elif paddle_op.name() == "pd_op.conv3d_transpose":
+        layer = network.add_deconvolution_nd(
+            input=input_tensor,
+            num_output_maps=n_input * groups,
+            kernel_shape=nv_ksize,
+            kernel=filter,
+            bias=None,
+        )
+    layer.stride_nd = nv_strides
+    layer.pre_padding = nv_pre_paddings
+
+    nv_post_paddings = trt.Dims3(paddings[0], paddings[1], paddings[2])
+    if output_padding:
+        nv_post_paddings[0] -= output_padding[0]
+        nv_post_paddings[1] -= output_padding[1]
+        nv_post_paddings[2] -= output_padding[2]
+
+        if (
+            nv_post_paddings[0] < 0
+            or nv_post_paddings[1] < 0
+            or nv_post_paddings[2] < 0
+        ):
+            raise ValueError(
+                "The value in conv3d_transpose's PostPadding should be >= 0."
+            )
+
+    layer.post_padding = nv_post_paddings
+    layer.num_groups = groups
+
+    if padding_algorithm == "SAME":
+        layer.padding_mode = trt.PaddingMode.SAME_UPPER
+
+    layer.dilation_nd = nv_dilations
+
+    return layer.get_output(0)
+
+
 def get_input_constant_value(paddle_op, inputs, input_index):
     input_op = paddle_op.operands()[input_index].source().get_defining_op()
     constant_manager = TensorRTConstantManager()
diff --git a/python/paddle/tensorrt/impls/conv.py b/python/paddle/tensorrt/impls/conv.py
@@ -13,7 +13,10 @@
 # limitations under the License.
 
 
-from paddle.tensorrt.converter_utils import convert_conv2d
+from paddle.tensorrt.converter_utils import (
+    convert_conv2d,
+    convert_conv3d,
+)
 from paddle.tensorrt.register import converter_registry
 
 
@@ -28,3 +31,9 @@
 )
 def conv2d_converter(network, paddle_op, inputs):
     return convert_conv2d(network, paddle_op, inputs)
+
+
+@converter_registry.register("pd_op.conv3d_transpose", trt_version="8.x")
+@converter_registry.register("pd_op.conv3d", trt_version="8.x")
+def conv3d_converter(network, paddle_op, inputs):
+    return convert_conv3d(network, paddle_op, inputs)
diff --git a/python/paddle/tensorrt/util.py b/python/paddle/tensorrt/util.py
@@ -293,6 +293,8 @@ def weight_to_tensor(network, paddle_value, trt_tensor, use_op_name):
         "pd_op.depthwise_conv2d",
         "pd_op.conv2d",
         "pd_op.conv2d_transpose",
+        "pd_op.conv3d",
+        "pd_op.conv3d_transpose",
         "pd_op.batch_norm",
         "pd_op.batch_norm_",
         "pd_op.layer_norm",
diff --git a/test/tensorrt/test_converter_conv.py b/test/tensorrt/test_converter_conv.py
@@ -316,6 +316,169 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+def conv3d_wrapper(x):
+    conv = paddle.nn.Conv3D(3, 3, (3, 3, 3))
+    return conv(x)
+
+
+def conv3d_python_api(x, padding="SAME", stride=(1, 1, 1)):
+    conv = paddle.nn.Conv3D(3, 3, (3, 3, 3), padding=padding, stride=stride)
+    return conv(x)
+
+
+class TestConv3dTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = conv3d_wrapper
+        self.api_args = {
+            "x": np.random.random([2, 3, 8, 8, 8]).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 8, 8, 8]}
+        self.opt_shape = {"x": [1, 3, 8, 8, 8]}
+        self.max_shape = {"x": [10, 3, 8, 8, 8]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(rtol=1e-3, atol=1e-3, precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestConv3dPaddingAlgorithmTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = conv3d_python_api
+        self.api_args = {
+            "x": np.random.random([2, 3, 8, 8, 8]).astype("float32"),
+            "paddings": "SAME",
+            "stride": (1, 1, 1),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 8, 8, 8]}
+        self.opt_shape = {"x": [1, 3, 8, 8, 8]}
+        self.max_shape = {"x": [10, 3, 8, 8, 8]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(rtol=1e-3, atol=1e-3, precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+def depthwise_conv3d_transpose_wrapper(x):
+    conv = paddle.nn.Conv3DTranspose(
+        in_channels=2, out_channels=2, kernel_size=(3, 3, 3)
+    )
+    return conv(x)
+
+
+def depthwise_conv3d_transpose_python_api(
+    x, padding="SAME", stride=(1, 1, 1), dilation=(1, 1, 1)
+):
+    conv = paddle.nn.Conv3DTranspose(
+        in_channels=2,
+        out_channels=2,
+        kernel_size=(3, 3, 3),
+        stride=stride,
+        padding=padding,
+        dilation=dilation,
+    )
+    return conv(x)
+
+
+def depthwise_conv3d_transpose_wrapper_outpadding(x, output_padding):
+    conv = paddle.nn.Conv3DTranspose(
+        in_channels=3,
+        out_channels=3,
+        kernel_size=(3, 3, 3),
+        stride=2,
+        output_padding=output_padding,
+    )
+    return conv(x)
+
+
+def conv3d_transpose_with_algorithm(x, algorithm):
+    conv = paddle.nn.Conv3DTranspose(
+        in_channels=3,
+        out_channels=3,
+        kernel_size=(3, 3, 3),
+        padding=algorithm,
+    )
+    return conv(x)
+
+
+class TestDepthwiseConv3dTransposeTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = depthwise_conv3d_transpose_wrapper
+        self.api_args = {
+            "x": np.random.random([3, 2, 8, 8, 8]).astype("float32")
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 8, 8, 8]}
+        self.opt_shape = {"x": [1, 2, 8, 8, 8]}
+        self.max_shape = {"x": [10, 2, 8, 8, 8]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(rtol=1e-3, atol=1e-3, precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestDepthwiseConv3dTransposeSameTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = conv3d_transpose_with_algorithm
+        self.api_args = {
+            "x": np.random.random([2, 3, 8, 8, 8]).astype("float32"),
+            "padding_algorithm": "SAME",
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 8, 8, 8]}
+        self.opt_shape = {"x": [1, 3, 8, 8, 8]}
+        self.max_shape = {"x": [10, 3, 8, 8, 8]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(rtol=1e-3, atol=1e-3, precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestDepthwiseConv3dTransposeOutputPaddingTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = depthwise_conv3d_transpose_wrapper_outpadding
+        self.api_args = {
+            "x": np.random.random([2, 3, 8, 8, 8]).astype("float32"),
+            "output_padding": [1, 1, 1],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 8, 8, 8]}
+        self.opt_shape = {"x": [1, 3, 8, 8, 8]}
+        self.max_shape = {"x": [10, 3, 8, 8, 8]}
+
+    def test_trt_result(self):
+        with self.assertRaises(ValueError) as context:
+            self.check_trt_result()
+
+
+class TestDepthwiseConv3dTransposeOutputPadding2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = depthwise_conv3d_transpose_wrapper_outpadding
+        self.api_args = {
+            "x": np.random.random([2, 3, 8, 8, 8]).astype("float32"),
+            "output_padding": [0, 0, 0],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 8, 8, 8]}
+        self.opt_shape = {"x": [1, 3, 8, 8, 8]}
+        self.max_shape = {"x": [10, 3, 8, 8, 8]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(rtol=1e-3, atol=1e-3, precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
 class TestFusedConv2dAddActTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = conv2d_wrapper
