Skip to content

Commit 7dc4af4

Browse files
narendasannarendasan
committed
feat(//core/conversion/converters/impl): Round out pooling
implementations Implements: - aten::max_pool1d - aten::max_pool3d - aten::avg_pool1d - aten::avg_pool3d Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
1 parent 0c39519 commit 7dc4af4

File tree

2 files changed

+197
-53
lines changed

2 files changed

+197
-53
lines changed

Diff for: ‎core/conversion/converters/impl/pooling.cpp

+142-53
Original file line numberDiff line numberDiff line change
@@ -8,62 +8,127 @@ namespace converters {
88
namespace impl {
99
namespace {
1010

11-
auto pooling_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
12-
.pattern({
13-
"aten::max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], int[2] dilation=[1, 1], bool ceil_mode=False) -> (Tensor)",
14-
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
15-
auto in = args[0].ITensor();
16-
auto shape = util::toVec(in->getDimensions());
17-
18-
// Max Pool needs at least 4D input
19-
if (shape.size() < 4) {
20-
auto new_shape = util::toDimsPad(shape, 4);
21-
LOG_DEBUG("Input shape is less than 4D got: " << util::toDims(shape) << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape);
22-
auto shuffle = ctx->net->addShuffle(*in);
23-
shuffle->setReshapeDimensions(new_shape);
24-
shuffle->setName((util::node_info(n) + " [Reshape to " + util::toStr(new_shape) + ']').c_str());
25-
in = shuffle->getOutput(0);
26-
}
27-
28-
29-
auto kernel_size = util::toDimsHW(args[1].unwrapToIntList());
30-
LOG_DEBUG("kernel_size: " << kernel_size);
31-
auto padding = util::toDimsHW(args[3].unwrapToIntList());
32-
LOG_DEBUG("padding: " << padding);
33-
auto dilation = util::toDims(args[4].unwrapToIntList());
34-
35-
TRTORCH_ASSERT(dilation == util::toDims(std::vector<int64_t>({1,1})), "Pooling dilation is not supported in TensorRT");
36-
37-
LOG_DEBUG("dilation: " << dilation);
38-
LOG_WARNING("Dilation not used in max pooling converter");
39-
bool ceil_mode = args[5].unwrapToBool();
11+
bool MaxPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& args) {
12+
auto in = args[0].ITensor();
13+
auto shape = util::toVec(in->getDimensions());
14+
15+
// Max Pool needs at least 4D input
16+
if (shape.size() < 4) {
17+
auto new_shape = util::toDimsPad(shape, 4);
18+
LOG_DEBUG("Input shape is less than 4D got: " << util::toDims(shape) << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape);
19+
auto shuffle = ctx->net->addShuffle(*in);
20+
shuffle->setReshapeDimensions(new_shape);
21+
shuffle->setName((util::node_info(n) + " [Reshape to " + util::toStr(new_shape) + ']').c_str());
22+
in = shuffle->getOutput(0);
23+
}
24+
25+
26+
auto kernel_size = util::toDimsHW(args[1].unwrapToIntList());
27+
LOG_DEBUG("kernel_size: " << kernel_size);
28+
auto padding = util::toDimsHW(args[3].unwrapToIntList());
29+
LOG_DEBUG("padding: " << padding);
30+
auto stride = util::toDims(args[2].unwrapToIntList());
31+
LOG_DEBUG("stride: " << stride);
32+
33+
auto dilation = util::toDims(args[4].unwrapToIntList());
34+
35+
TRTORCH_ASSERT(dilation == util::toDims(std::vector<int64_t>({1,1})), "Pooling dilation is not supported in TensorRT");
36+
37+
LOG_DEBUG("dilation: " << dilation);
38+
LOG_WARNING("Dilation not used in max pooling converter");
39+
bool ceil_mode = args[5].unwrapToBool();
40+
41+
auto new_layer = ctx->net->addPoolingNd(*in, nvinfer1::PoolingType::kMAX, kernel_size);
42+
TRTORCH_CHECK(new_layer, "Unable to create Max Pool 2D layer from node: " << *n);
43+
44+
new_layer->setName(util::node_info(n).c_str());
45+
new_layer->setPaddingNd(padding);
46+
if (stride.nbDims != 2 && ctx->settings.device == nvinfer1::DeviceType::kDLA) {
47+
if (!ctx->settings.allow_gpu_fallback) {
48+
TRTORCH_THROW_ERROR("DLA Pooling stride is limited to 2D, allow GPU fallback");
49+
} else {
50+
LOG_WARNING("DLA Pooling stride is limited to 2D, will run on GPU");
51+
}
52+
}
53+
new_layer->setStrideNd(stride);
54+
55+
auto padding_mode = ceil_mode ? nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP : nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN;
56+
new_layer->setPaddingMode(padding_mode);
57+
58+
new_layer->setName(util::node_info(n).c_str());
59+
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
60+
61+
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
62+
return true;
63+
}
64+
65+
bool AvgPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& args) {
66+
auto in = args[0].ITensor();
67+
auto shape = util::toVec(in->getDimensions());
68+
69+
// Avg Pool needs at least 4D input
70+
if (shape.size() < 4) {
71+
auto new_shape = util::toDimsPad(shape, 4);
72+
LOG_DEBUG("Input shape is less than 4D got: " << util::toDims(shape) << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape);
73+
auto shuffle = ctx->net->addShuffle(*in);
74+
shuffle->setReshapeDimensions(new_shape);
75+
shuffle->setName((util::node_info(n) + " [Reshape to " + util::toStr(new_shape) + ']').c_str());
76+
in = shuffle->getOutput(0);
77+
}
78+
79+
80+
auto kernel_size = util::toDimsHW(args[1].unwrapToIntList());
81+
LOG_DEBUG("kernel_size: " << kernel_size);
82+
auto padding = util::toDimsHW(args[3].unwrapToIntList());
83+
LOG_DEBUG("padding: " << padding);
84+
auto stride = util::toDims(args[2].unwrapToIntList());
85+
LOG_DEBUG("stride: " << stride);
86+
87+
bool ceil_mode = args[4].unwrapToBool();
88+
bool count_inlcude_pad = args[5].unwrapToBool();
89+
90+
auto new_layer = ctx->net->addPoolingNd(*in, nvinfer1::PoolingType::kAVERAGE, kernel_size);
91+
TRTORCH_CHECK(new_layer, "Unable to create Avg Pool 2D layer from node: " << *n);
92+
93+
new_layer->setName(util::node_info(n).c_str());
94+
new_layer->setPaddingNd(padding);
95+
if (stride.nbDims != 2 && ctx->settings.device == nvinfer1::DeviceType::kDLA) {
96+
if (!ctx->settings.allow_gpu_fallback) {
97+
TRTORCH_THROW_ERROR("DLA Pooling stride is limited to 2D, allow GPU fallback");
98+
} else {
99+
LOG_WARNING("DLA Pooling stride is limited to 2D, will run on GPU");
100+
}
101+
}
102+
new_layer->setStrideNd(stride);
40103

41-
auto new_layer = ctx->net->addPoolingNd(*in, nvinfer1::PoolingType::kMAX, kernel_size);
42-
TRTORCH_CHECK(new_layer, "Unable to create Max Pool 2D layer from node: " << *n);
104+
auto padding_mode = ceil_mode ? nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP : nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN;
105+
new_layer->setPaddingMode(padding_mode);
106+
new_layer->setAverageCountExcludesPadding(!count_inlcude_pad);
43107

44-
new_layer->setName(util::node_info(n).c_str());
45-
new_layer->setPaddingNd(padding);
46-
if (args[2].unwrapToIntList().size() == 2) {
47-
auto stride = util::toDims(args[2].unwrapToIntList());
48-
new_layer->setStrideNd(stride);
49-
}
108+
if (!(args[6].IValue()->isNone())) {
109+
LOG_WARNING("Divisor override is now handled by Avg Pooling Converter");
110+
}
50111

51-
auto padding_mode = ceil_mode ? nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP : nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN;
52-
new_layer->setPaddingMode(padding_mode);
112+
new_layer->setName(util::node_info(n).c_str());
113+
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
53114

54-
new_layer->setName(util::node_info(n).c_str());
55-
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
115+
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
116+
return true;
117+
}
56118

57-
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
58-
return true;
119+
auto pooling_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
120+
.pattern({
121+
"aten::max_pool1d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], int[2] dilation=[1, 1], bool ceil_mode=False) -> (Tensor)",
122+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
123+
return MaxPoolingConverter(ctx, n, args);
59124
}
60125
}).pattern({
61-
"aten::avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> (Tensor)",
126+
"aten::avg_pool1d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], bool ceil_mode=False, bool count_include_pad=True) -> (Tensor)",
62127
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
63128
auto in = args[0].ITensor();
64129
auto shape = util::toVec(in->getDimensions());
65130

66-
// Abg Pool needs at least 4D input
131+
// Avg Pool needs at least 4D input
67132
if (shape.size() < 4) {
68133
auto new_shape = util::toDimsPad(shape, 4);
69134
LOG_DEBUG("Input shape is less than 4D got: " << util::toDims(shape) << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape);
@@ -78,6 +143,8 @@ auto pooling_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
78143
LOG_DEBUG("kernel_size: " << kernel_size);
79144
auto padding = util::toDimsHW(args[3].unwrapToIntList());
80145
LOG_DEBUG("padding: " << padding);
146+
auto stride = util::toDims(args[2].unwrapToIntList());
147+
LOG_DEBUG("stride: " << stride);
81148

82149
bool ceil_mode = args[4].unwrapToBool();
83150
bool count_inlcude_pad = args[5].unwrapToBool();
@@ -87,26 +154,48 @@ auto pooling_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
87154

88155
new_layer->setName(util::node_info(n).c_str());
89156
new_layer->setPaddingNd(padding);
90-
if (args[2].unwrapToIntList().size() == 2) {
91-
auto stride = util::toDims(args[2].unwrapToIntList());
92-
LOG_DEBUG("stride: " << stride);
93-
new_layer->setStrideNd(stride);
157+
158+
if (stride.nbDims != 2 && ctx->settings.device == nvinfer1::DeviceType::kDLA) {
159+
if (!ctx->settings.allow_gpu_fallback) {
160+
TRTORCH_THROW_ERROR("DLA Pooling stride is limited to 2D, allow GPU fallback");
161+
} else {
162+
LOG_WARNING("DLA Pooling stride is limited to 2D, will run on GPU");
163+
}
94164
}
95165

166+
new_layer->setStrideNd(stride);
167+
96168
auto padding_mode = ceil_mode ? nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP : nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN;
97169
new_layer->setPaddingMode(padding_mode);
98170
new_layer->setAverageCountExcludesPadding(!count_inlcude_pad);
99171

100-
if (!(args[6].IValue()->isNone())) {
101-
LOG_WARNING("Divisor override is now handled by Avg Pooling Converter");
102-
}
103-
104172
new_layer->setName(util::node_info(n).c_str());
105173
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
106174

107175
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
108176
return true;
109177
}
178+
})
179+
.pattern({
180+
"aten::max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], int[2] dilation=[1, 1], bool ceil_mode=False) -> (Tensor)",
181+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
182+
return MaxPoolingConverter(ctx, n, args);
183+
}
184+
}).pattern({
185+
"aten::avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> (Tensor)",
186+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
187+
return AvgPoolingConverter(ctx, n, args);
188+
}
189+
}).pattern({
190+
"aten::max_pool3d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], int[2] dilation=[1, 1], bool ceil_mode=False) -> (Tensor)",
191+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
192+
return MaxPoolingConverter(ctx, n, args);
193+
}
194+
}).pattern({
195+
"aten::avg_pool3d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=[0, 0], bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> (Tensor)",
196+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
197+
return AvgPoolingConverter(ctx, n, args);
198+
}
110199
}).pattern({
111200
"aten::adaptive_avg_pool2d(Tensor self, int[2] output_size) -> (Tensor)",
112201
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {

Diff for: ‎tests/core/converters/test_pooling.cpp

+55
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,33 @@
44
#include "tests/util/util.h"
55
#include "core/compiler.h"
66

7+
TEST(Converters, ATenMaxPool1DConvertsCorrectly) {
8+
const auto graph = R"IR(
9+
graph(%0 : Tensor):
10+
%1 : int = prim::Constant[value=0]()
11+
%2 : int = prim::Constant[value=1]()
12+
%3 : int = prim::Constant[value=2]()
13+
%5 : bool = prim::Constant[value=0]()
14+
%6 : int[] = prim::ListConstruct(%1, %1)
15+
%7 : int[] = prim::ListConstruct(%2, %2)
16+
%8 : int[] = prim::ListConstruct(%3, %3)
17+
%10 : Tensor = aten::max_pool2d(%0, %8, %7, %6, %7, %5)
18+
return (%10))IR";
19+
20+
auto g = std::make_shared<torch::jit::Graph>();
21+
torch::jit::parseIR(graph, &*g);
22+
23+
auto in = at::randint(-5, 5, {1, 4, 4}, at::kCUDA);
24+
auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
25+
auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
26+
27+
in = at::clone(in);
28+
params = trtorch::core::conversion::get_named_params(g->inputs(), {});
29+
auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
30+
31+
ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
32+
}
33+
734
TEST(Converters, ATenMaxPool2DConvertsCorrectly) {
835
const auto graph = R"IR(
936
graph(%0 : Tensor):
@@ -32,6 +59,34 @@ TEST(Converters, ATenMaxPool2DConvertsCorrectly) {
3259
ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
3360
}
3461

62+
TEST(Converters, ATenMaxPool3DConvertsCorrectly) {
63+
const auto graph = R"IR(
64+
graph(%0 : Tensor):
65+
%1 : int = prim::Constant[value=0]()
66+
%2 : int = prim::Constant[value=1]()
67+
%3 : int = prim::Constant[value=2]()
68+
%5 : bool = prim::Constant[value=0]()
69+
%6 : int[] = prim::ListConstruct(%1, %1)
70+
%7 : int[] = prim::ListConstruct(%2, %2)
71+
%8 : int[] = prim::ListConstruct(%3, %3)
72+
%10 : Tensor = aten::max_pool2d(%0, %8, %7, %6, %7, %5)
73+
return (%10))IR";
74+
75+
auto g = std::make_shared<torch::jit::Graph>();
76+
torch::jit::parseIR(graph, &*g);
77+
78+
//PyTorch MaxPool needs a 3D input
79+
auto in = at::randint(-5, 5, {1, 4, 4, 4}, at::kCUDA);
80+
auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
81+
auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
82+
83+
in = at::clone(in);
84+
params = trtorch::core::conversion::get_named_params(g->inputs(), {});
85+
auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
86+
87+
ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
88+
}
89+
3590
TEST(Converters, ATenAvgPool2DConvertsCorrectly) {
3691
const auto graph = R"IR(
3792
graph(%0 : Tensor):

0 commit comments

Comments
 (0)