Skip to content

Commit

Permalink
[BYOC][TRT]Fix groups cannot divide output channel count error for de…
Browse files Browse the repository at this point in the history
…conv when groups>1 (#7595)

* trt num_outputs

* asdf

* fix lint

Co-authored-by: Leyuan Wang <leyuan.wang@bytedance.com>
  • Loading branch information
Laurawly and Leyuan Wang committed Mar 5, 2021
1 parent b9adce2 commit c5f608f
Showing 1 changed file with 6 additions and 6 deletions.
12 changes: 6 additions & 6 deletions src/runtime/contrib/tensorrt/tensorrt_ops.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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<IntImmNode>()->value
const int num_outputs = weight_shape[0];
const int num_outputs =
std::stoi(params->node.GetAttr<std::vector<std::string>>("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,
Expand Down Expand Up @@ -788,8 +788,8 @@ class Conv2DTransposeOpConverter : public TensorRTOpConverter {
}
#endif

// Could use conv2d_attr->channels.as<IntImmNode>()->value
const int num_outputs = weight_shape[1];
const int num_outputs =
std::stoi(params->node.GetAttr<std::vector<std::string>>("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,
Expand Down Expand Up @@ -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<IntImmNode>()->value
const int num_outputs = weight_shape[1];
const int num_outputs =
std::stoi(params->node.GetAttr<std::vector<std::string>>("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,
Expand Down

0 comments on commit c5f608f

Please sign in to comment.