From c5f608f08861e2b4e869e0db0f68def7bcf65338 Mon Sep 17 00:00:00 2001 From: Leyuan Wang Date: Fri, 5 Mar 2021 04:47:07 -0800 Subject: [PATCH] [BYOC][TRT]Fix groups cannot divide output channel count error for deconv when groups>1 (#7595) * trt num_outputs * asdf * fix lint Co-authored-by: Leyuan Wang --- src/runtime/contrib/tensorrt/tensorrt_ops.cc | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/runtime/contrib/tensorrt/tensorrt_ops.cc b/src/runtime/contrib/tensorrt/tensorrt_ops.cc index 824178eaa619..04b1e838ee8e 100644 --- a/src/runtime/contrib/tensorrt/tensorrt_ops.cc +++ b/src/runtime/contrib/tensorrt/tensorrt_ops.cc @@ -309,8 +309,8 @@ class Conv3DOpConverter : public TensorRTOpConverter { bool use_asymmetric_padding; GetPadding3D(str_padding, &use_asymmetric_padding, &prepadding, &postpadding); - // Could use attrs->channels.as()->value - const int num_outputs = weight_shape[0]; + const int num_outputs = + std::stoi(params->node.GetAttr>("channels")[0]); const auto kernel_size = nvinfer1::Dims3(weight_shape[2], weight_shape[3], weight_shape[4]); nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, nullptr, 0}; auto conv_layer = params->network->addConvolutionNd(*input_tensor, num_outputs, kernel_size, @@ -788,8 +788,8 @@ class Conv2DTransposeOpConverter : public TensorRTOpConverter { } #endif - // Could use conv2d_attr->channels.as()->value - const int num_outputs = weight_shape[1]; + const int num_outputs = + std::stoi(params->node.GetAttr>("channels")[0]); const auto kernel_size = nvinfer1::DimsHW(weight_shape[2], weight_shape[3]); nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, nullptr, 0}; auto deconv_layer = params->network->addDeconvolution(*input_tensor, num_outputs, kernel_size, @@ -846,8 +846,8 @@ class Conv3DTransposeOpConverter : public TensorRTOpConverter { bool use_asymmetric_padding; GetPadding3D(str_padding, &use_asymmetric_padding, &prepadding, &postpadding); - // Could use attrs->channels.as()->value - const int num_outputs = weight_shape[1]; + const int num_outputs = + std::stoi(params->node.GetAttr>("channels")[0]); const auto kernel_size = nvinfer1::Dims3(weight_shape[2], weight_shape[3], weight_shape[4]); nvinfer1::Weights bias{nvinfer1::DataType::kFLOAT, nullptr, 0}; auto deconv_layer = params->network->addDeconvolutionNd(*input_tensor, num_outputs, kernel_size,