【BUPT】【Paddle TensorRT No.18】Add pd_op.conv3d and pd_op.conv3d_transpose converter

paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc

@@ -373,6 +373,116 @@ class DepthwiseConv2dTransposeOpPattern
     return true;
   }
 };
+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;
+    }
+
+    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;
+      }
+    }
+
+    if (op->HasAttribute("output_padding")) {
+      auto output_padding_attr =
+          op->attribute<pir::ArrayAttribute>("output_padding");
+      std::vector<int32_t> output_padding;
+      for (const auto &attr : output_padding_attr.AsVector()) {
+        output_padding.push_back(attr.dyn_cast<pir::Int32Attribute>().data());
+      }
+      if (!output_padding.empty()) {
+        int max_padding =
+            *std::max_element(output_padding.begin(), output_padding.end());
+        if (max_padding > 0) {
+          VLOG(3)
+              << "TensorRT doesn't support output_padding when version < 8406";
+          return false;
+        }
+      }
+    }
+
+    if (op->HasAttribute("padding_algorithm")) {
+      auto padding_algorithm =
+          op->attribute<pir::StrAttribute>("padding_algorithm").AsString();
+      if (padding_algorithm == "SAME") {
+        VLOG(3)
+            << "TensorRT error is raised if conv3d_transpose and SAME padding";
+        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;
+    }
+
+    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;
+    }
+
+    if (op->HasAttribute("padding_algorithm")) {
+      auto padding_algorithm =
+          op->attribute<pir::StrAttribute>("padding_algorithm").AsString();
+      if (padding_algorithm == "SAME") {
+        VLOG(3) << "TensorRT error is raised if conv3d and SAME padding";
+        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;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
 
 class DeformableConvOpPattern
     : public pir::OpRewritePattern<paddle::dialect::DeformableConvOp> {
@@ -399,6 +509,7 @@ class DeformableConvOpPattern
               << input_shape.size() << "-D Tensor";
       return false;
     }
+
     pir::Value filter = op.operand_source(2);
     auto filter_type =
         filter.type().dyn_cast<paddle::dialect::DenseTensorType>();
@@ -2204,6 +2315,8 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<OneHotOpPattern>(context));
     ps.Add(std::make_unique<AssignValueOpPattern>(context));
     ps.Add(std::make_unique<AssignValue_OpPattern>(context));
+    ps.Add(std::make_unique<Conv3dOpPattern>(context));
+    ps.Add(std::make_unique<Conv3dTransposeOpPattern>(context));
     return ps;
   }
 };

python/paddle/tensorrt/converter_utils.py

@@ -510,6 +510,89 @@ def convert_conv2d(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
+def convert_conv3d(network, paddle_op, inputs):
+    if paddle_op.name() == "pd_op.conv3d":
+        input_tensor, filter = inputs
+    elif paddle_op.name() == "pd_op.conv3d_transpose":
+        if len(inputs) == 3:
+            input_tensor, filter, output_size = inputs
+        elif len(inputs) == 2:
+            input_tensor, filter = inputs
+            output_size = None
+        else:
+            raise ValueError("Invalid number of inputs for conv3d_transpose")
+
+    filter_shape = paddle_op.operands()[1].source().shape
+
+    if len(filter_shape) != 5:
+        raise ValueError(
+            f"The conv3d filter's dims size should be 5, but got {len(filter_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")
+    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.d[0] -= output_padding[0]
+        nv_post_paddings.d[1] -= output_padding[1]
+        nv_post_paddings.d[2] -= output_padding[2]
+
+        if (
+            nv_post_paddings.d[0] < 0
+            or nv_post_paddings.d[1] < 0
+            or nv_post_paddings.d[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 add_reduce_layer(network, paddle_op, inputs, op_type):
     input_tensor = inputs[0]
     axis = paddle_op.operands()[1].source().get_defining_op().attrs()["value"]

python/paddle/tensorrt/impls/conv.py

@@ -13,7 +13,7 @@
 # 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
 
 
@@ -25,3 +25,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)

test/tensorrt/test_converter_conv.py

@@ -298,5 +298,127 @@ 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.max_shape = {"x": [10, 3, 8, 8, 8]}
+        self.enable_fp16 = True
+
+    def test_trt_result(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"),
+            "padding": "SAME",
+            "stride": (1, 1, 1),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 8, 8, 8]}
+        self.max_shape = {"x": [10, 3, 8, 8, 8]}
+        self.enable_fp16 = True
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+def conv3dtranspose_wrapper(
+    x,
+    stride=1,
+    padding=0,
+    output_padding=(0, 0, 0),
+    output_size=[],
+    padding_algorithm="EXPLICIT",
+    groups=1,
+    dilation=1,
+    data_format="NCDHW",
+):
+
+    fc = x.shape[1]
+    weight = paddle.static.create_parameter(
+        name="weight",
+        shape=[fc, 6, 3, 3, 3],
+        dtype="float32",
+        default_initializer=paddle.nn.initializer.Normal(mean=0.0, std=1.0),
+    )
+
+    return _C_ops.conv3d_transpose(
+        x,
+        weight,
+        stride,
+        padding,
+        output_padding,
+        output_size,
+        padding_algorithm,
+        groups,
+        dilation,
+        data_format,
+    )
+
+
+class TestConv3dTransposeTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = conv3dtranspose_wrapper
+        self.api_args = {
+            "x": np.random.random([1, 2, 5, 5, 5]).astype("float32"),
+            "stride": [1, 1, 1],
+            "padding": [1, 0, 1, 2, 1, 1],
+            "output_padding": [0, 0, 0],
+            "output_size": [],
+            "padding_algorithm": "EXPLICIT",
+            "groups": 1,
+            "dilation": [1, 1, 1],
+            "data_format": "NCDHW",
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 5, 5, 5]}
+        self.max_shape = {"x": [4, 2, 5, 5, 5]}
+        self.enable_fp16 = True
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestConv3dTransposePaddingAlgorithmTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = conv3dtranspose_wrapper
+        self.api_args = {
+            "x": np.random.random([1, 2, 5, 5, 5]).astype("float32"),
+            "stride": [1, 1, 1],
+            "padding": [1, 0, 1, 2, 1, 1],
+            "output_padding": [0, 0, 0],
+            "output_size": [],
+            "padding_algorithm": "SAME",
+            "groups": 1,
+            "dilation": [1, 1, 1],
+            "data_format": "NCDHW",
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 5, 5, 5]}
+        self.max_shape = {"x": [4, 2, 5, 5, 5]}
+        self.enable_fp16 = True
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()