@@ -10,6 +10,74 @@ namespace converters {
1010namespace impl {
1111namespace {
1212
13+ void add_output_padding (nvinfer1::Dims& padding, nvinfer1::Dims& out_padding, bool & has_output_padding) {
14+ int nbSpatialDims = out_padding.nbDims ;
15+ // When there is out_padding, if padding is larger than out_padding, just adjust padding Or reduce out_padding as
16+ // minimum as possible.
17+ for (int i = 0 ; i < nbSpatialDims; ++i) {
18+ if (padding.d [i] - out_padding.d [i] >= 0 ) {
19+ padding.d [i] -= out_padding.d [i];
20+ out_padding.d [i] = 0 ;
21+ } else {
22+ // Reduce out_padding as possible.
23+ out_padding.d [i] -= padding.d [i];
24+ padding.d [i] = 0 ;
25+ has_output_padding = true ;
26+ }
27+ }
28+ }
29+
30+ nvinfer1::ILayer* add_bias_layer (
31+ ConversionCtx* ctx,
32+ nvinfer1::ITensor* input_tensor,
33+ nvinfer1::Dims& input_dims,
34+ nvinfer1::Dims& output_padding,
35+ Weights& bias) {
36+ nvinfer1::ITensor* input_shape = ctx->net ->addShape (*input_tensor)->getOutput (0 );
37+ // Add padding layer
38+ nvinfer1::ITensor* start;
39+ nvinfer1::ITensor* totalPadding;
40+ auto in_nbDims = input_dims.nbDims ;
41+ std::vector<int32_t > startVec (in_nbDims, 0 );
42+ std::vector<int32_t > totalPaddingVec (in_nbDims, 0 );
43+ int32_t diff = in_nbDims - output_padding.nbDims ;
44+ for (int32_t i = diff; i < in_nbDims; i++) {
45+ int32_t idx = i - diff;
46+ startVec[i] = 0 ; // Don't need begin padding, only post padding
47+ totalPaddingVec[i] = output_padding.d [idx];
48+ }
49+ start = tensor_to_const (ctx, torch::tensor (startVec, torch::kInt32 ));
50+ totalPadding = tensor_to_const (ctx, torch::tensor (totalPaddingVec, torch::kInt32 ));
51+
52+ const auto size =
53+ ctx->net ->addElementWise (*input_shape, *totalPadding, nvinfer1::ElementWiseOperation::kSUM )->getOutput (0 );
54+
55+ nvinfer1::Dims stride;
56+ stride.nbDims = in_nbDims;
57+ for (int64_t i = 0 ; i < in_nbDims; i++) {
58+ stride.d [i] = 1 ;
59+ }
60+ const auto & dummy = stride;
61+ auto * sliceLayer = ctx->net ->addSlice (*input_tensor, dummy, dummy, stride);
62+ sliceLayer->setInput (1 , *start);
63+ sliceLayer->setInput (2 , *size);
64+ sliceLayer->setMode (nvinfer1::SliceMode::kFILL );
65+ nvinfer1::ITensor* slice_output = sliceLayer->getOutput (0 );
66+
67+ nvinfer1::Dims constantDims;
68+ constantDims.nbDims = in_nbDims;
69+ for (int64_t i = 0 ; i < in_nbDims; i++) {
70+ constantDims.d [i] = 1 ;
71+ }
72+ constantDims.d [diff - 1 ] =
73+ bias.shape .d [0 ]; // Set C dimension to bias dim and other dimensions to 1 to enable broadcast
74+ auto const_layer = ctx->net ->addConstant (constantDims, bias.data );
75+ auto bias_layer =
76+ ctx->net ->addElementWise (*slice_output, *const_layer->getOutput (0 ), nvinfer1::ElementWiseOperation::kSUM );
77+
78+ return bias_layer;
79+ }
80+
1381bool add_conv_deconv (ConversionCtx* ctx, const torch::jit::Node* n, args& args) {
1482 // Input to conv/deconv
1583 auto in = args[0 ].ITensor ();
@@ -76,16 +144,29 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
76144
77145 nvinfer1::ILayer* layer = nullptr ;
78146 if (transposed) {
79- nvinfer1::IDeconvolutionLayer* deconvLayer =
80- ctx->net ->addDeconvolutionNd (*in, kernel_dims.d [0 ], filter_dim, kernel_weights, bias.data );
147+ // Fix padding based on output_padding provided
148+ nvinfer1::Dims begPadding = padding;
149+ bool hasOutputPadding = false ;
150+ add_output_padding (padding, out_padding, hasOutputPadding);
151+
152+ nvinfer1::IDeconvolutionLayer* deconvLayer = ctx->net ->addDeconvolutionNd (
153+ *in, kernel_dims.d [0 ], filter_dim, kernel_weights, hasOutputPadding ? nvinfer1::Weights{} : bias.data );
81154 deconvLayer->setStrideNd (stride);
82155 deconvLayer->setDilationNd (dilation);
83156 deconvLayer->setNbGroups (groups);
84- deconvLayer->setPaddingNd (padding);
157+ deconvLayer->setPrePadding (begPadding);
158+ deconvLayer->setPostPadding (padding);
159+
85160 // Set deconv kernel weights
86161 deconvLayer->setInput (1 , *kernel);
87162 TORCHTRT_CHECK (deconvLayer, " Unable to create deconv layer with non-const weights from node: "
88163 layer = deconvLayer;
164+ if (hasOutputPadding) {
165+ LOG_DEBUG (" Padding output deconvolution tensor with:"
166+ nvinfer1::ITensor* tensorPtr = deconvLayer->getOutput (0 );
167+ auto dims = in->getDimensions ();
168+ layer = add_bias_layer (ctx, tensorPtr, dims, out_padding, bias);
169+ }
89170 } else {
90171 nvinfer1::IConvolutionLayer* convLayer =
91172 ctx->net ->addConvolutionNd (*in, kernel_dims.d [0 ], filter_dim, kernel_weights, bias.data );
@@ -155,20 +236,7 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
155236 // https://github.com/onnx/onnx-tensorrt/blob/c3cfcbc8248c6bd007e6630af2085df5e4834b42/builtin_op_importers.cpp#L734
156237 nvinfer1::Dims begPadding = padding;
157238 bool hasOutputPadding = false ;
158- int nbSpatialDims = out_padding.nbDims ;
159- // When there is out_padding, if padding is larger than out_padding, just adjust padding Or reduce out_padding as
160- // minimum as possible.
161- for (int i = 0 ; i < nbSpatialDims; ++i) {
162- if (padding.d [i] - out_padding.d [i] >= 0 ) {
163- padding.d [i] -= out_padding.d [i];
164- out_padding.d [i] = 0 ;
165- } else {
166- // Reduce out_padding as possible.
167- out_padding.d [i] -= padding.d [i];
168- padding.d [i] = 0 ;
169- hasOutputPadding = true ;
170- }
171- }
239+ add_output_padding (padding, out_padding, hasOutputPadding);
172240
173241 // shape of deconvolution's weight: [in, out/groups, ...]
174242 // If there is still output padding, remove the bias. Bias will be added below.
@@ -190,51 +258,8 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
190258#endif
191259 if (hasOutputPadding) {
192260 LOG_DEBUG (" Padding output deconvolution tensor with:"
193-
194- // Add padding layer
195- nvinfer1::ITensor* start;
196- nvinfer1::ITensor* totalPadding;
197- auto in_nbDims = orig_dims.nbDims ;
198- std::vector<int32_t > startVec (in_nbDims, 0 );
199- std::vector<int32_t > totalPaddingVec (in_nbDims, 0 );
200- int32_t diff = in_nbDims - out_padding.nbDims ;
201- for (int32_t i = diff; i < in_nbDims; i++) {
202- int32_t idx = i - diff;
203- startVec[i] = 0 ; // Don't need begin padding, only post padding
204- totalPaddingVec[i] = out_padding.d [idx];
205- }
206- start = tensor_to_const (ctx, torch::tensor (startVec, torch::kInt32 ));
207- totalPadding = tensor_to_const (ctx, torch::tensor (totalPaddingVec, torch::kInt32 ));
208-
209261 nvinfer1::ITensor* tensorPtr = deconv->getOutput (0 );
210- nvinfer1::ITensor* deconvOutShape = ctx->net ->addShape (*tensorPtr)->getOutput (0 );
211- const auto size =
212- ctx->net ->addElementWise (*deconvOutShape, *totalPadding, nvinfer1::ElementWiseOperation::kSUM )->getOutput (0 );
213-
214- nvinfer1::Dims stride;
215- stride.nbDims = in_nbDims;
216- for (int64_t i = 0 ; i < in_nbDims; i++) {
217- stride.d [i] = 1 ;
218- }
219- const auto & dummy = stride;
220- auto * sliceLayer = ctx->net ->addSlice (*tensorPtr, dummy, dummy, stride);
221- sliceLayer->setInput (1 , *start);
222- sliceLayer->setInput (2 , *size);
223- sliceLayer->setMode (nvinfer1::SliceMode::kFILL );
224- tensorPtr = sliceLayer->getOutput (0 );
225-
226- nvinfer1::Dims constantDims;
227- constantDims.nbDims = in_nbDims;
228- for (int64_t i = 0 ; i < in_nbDims; i++) {
229- constantDims.d [i] = 1 ;
230- }
231- constantDims.d [diff - 1 ] =
232- bias.shape .d [0 ]; // Set C dimension to bias dim and other dimensions to 1 to enable broadcast
233- auto const_layer = ctx->net ->addConstant (constantDims, bias.data );
234- auto add_bias_layer =
235- ctx->net ->addElementWise (*tensorPtr, *const_layer->getOutput (0 ), nvinfer1::ElementWiseOperation::kSUM );
236-
237- new_layer = add_bias_layer;
262+ new_layer = add_bias_layer (ctx, tensorPtr, orig_dims, out_padding, bias);
238263 } else {
239264 new_layer = deconv;
240265 }
0 commit comments