From f6f5e3efe5799725c231ad3fe5319b5967b5e65b Mon Sep 17 00:00:00 2001 From: Ruoqian Guo Date: Thu, 30 Dec 2021 08:36:30 +0000 Subject: [PATCH 1/5] fix: Update "reduceAxes" variable in GlobalPoolingConverter function and add corresponding uTests Signed-off-by: Ruoqian Guo --- core/conversion/converters/impl/pooling.cpp | 3 +- .../conversion/converters/test_pooling.cpp | 52 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/core/conversion/converters/impl/pooling.cpp b/core/conversion/converters/impl/pooling.cpp index 48e3a70ff4..2f591ab797 100644 --- a/core/conversion/converters/impl/pooling.cpp +++ b/core/conversion/converters/impl/pooling.cpp @@ -16,8 +16,9 @@ bool GlobalPoolingConverter( nvinfer1::PoolingType pool_type) { auto in = args[0].ITensorOrFreeze(ctx); nvinfer1::Dims dims = in->getDimensions(); + auto out_size = util::toDims(args[1].unwrapToIntList()); // Generate a bitmask of all 1s except the last 2 bits (N and C axes) - uint32_t reduceAxes = ((1 << dims.nbDims) - 1) & ~0b11; + uint32_t reduceAxes = ((1 << dims.nbDims) - 1) ^ ((1 << (dims.nbDims - out_size.nbDims)) - 1); auto* new_layer = ctx->net->addReduce( *in, pool_type == nvinfer1::PoolingType::kMAX ? nvinfer1::ReduceOperation::kMAX : nvinfer1::ReduceOperation::kAVG, diff --git a/tests/core/conversion/converters/test_pooling.cpp b/tests/core/conversion/converters/test_pooling.cpp index 6b49eb7baf..a8c1cad760 100644 --- a/tests/core/conversion/converters/test_pooling.cpp +++ b/tests/core/conversion/converters/test_pooling.cpp @@ -436,6 +436,32 @@ TEST(Converters, ATenAdaptiveAvgPool2DConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); } +TEST(Converters, ATenAdaptiveAvgPool2DGlobalPoolingConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%2, %3) + %10 : Tensor = aten::adaptive_avg_pool2d(%0, %6) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch PyTorch adaptive_avg_pool2d needs a 4D input or a 3D input + auto in = at::randint(-5, 5, {64, 16, 32, 32}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + TEST(Converters, ATenAdaptiveAvgPool2DConvertsCorrectlyWithDynamicInput) { const auto graph = R"IR( graph(%0 : Tensor): @@ -488,6 +514,32 @@ TEST(Converters, ATenAdaptiveAvgPool1DConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 1.0)); } +TEST(Converters, ATenAdaptiveAvgPool1DGlobalPoolingConvertsCorrectly) { + const auto graph = + R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%2) + %10 : Tensor = aten::adaptive_avg_pool1d(%0, %6) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch adaptive_avg_pool1d needs a 3D input or a 2D input + auto in = at::randint(-5, 5, {3, 16}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + TEST(Converters, ATenAdaptiveMaxPool2DConvertsCorrectly) { const auto graph = R"IR( graph(%0 : Tensor): From deb9f743e765faf9334870b56670f440b228f836 Mon Sep 17 00:00:00 2001 From: Ruoqian Guo Date: Thu, 30 Dec 2021 09:47:23 +0000 Subject: [PATCH 2/5] fix: add at::adaptive_avg_pool1d in interpolate plugin and fix #791 Signed-off-by: Ruoqian Guo --- core/conversion/converters/impl/pooling.cpp | 24 +++++------------ core/plugins/impl/interpolate_plugin.cpp | 2 ++ .../conversion/converters/test_pooling.cpp | 26 +++++++++++++++++++ 3 files changed, 34 insertions(+), 18 deletions(-) diff --git a/core/conversion/converters/impl/pooling.cpp b/core/conversion/converters/impl/pooling.cpp index 2f591ab797..f18ef91c1b 100644 --- a/core/conversion/converters/impl/pooling.cpp +++ b/core/conversion/converters/impl/pooling.cpp @@ -37,7 +37,7 @@ bool AdaptivePoolingConverter( ConversionCtx* ctx, const torch::jit::Node* n, args& args, - nvinfer1::PoolingType pool_type) { + nvinfer1::PoolingType pool_type, const std::string& mode) { auto in = args[0].ITensorOrFreeze(ctx); auto out_size = util::toDims(args[1].unwrapToIntList()); @@ -48,15 +48,7 @@ bool AdaptivePoolingConverter( } auto orig_dims = in->getDimensions(); - bool expandDims = (orig_dims.nbDims < 4); - TORCHTRT_CHECK(orig_dims.nbDims > 2, "Unable to create pooling layer from node: " << *n); - if (expandDims) { - in = addPadding(ctx, n, in, 4, false, false); - } - - if (out_size.nbDims == 1) { - out_size = util::unsqueezeDims(out_size, 0, 1); - } + TORCHTRT_CHECK(orig_dims.nbDims > 1, "Unable to create pooling layer from node: " << *n); auto in_shape = util::toVec(in->getDimensions()); nvinfer1::ILayer* new_layer = nullptr; @@ -90,10 +82,6 @@ bool AdaptivePoolingConverter( int32_t use_scales_casted = 0; f.emplace_back(nvinfer1::PluginField("use_scales", &use_scales_casted, nvinfer1::PluginFieldType::kINT32, 1)); - std::string mode = "adaptive_avg_pool2d"; - if (pool_type == nvinfer1::PoolingType::kMAX) { - mode = "adaptive_max_pool2d"; - } f.emplace_back(nvinfer1::PluginField("mode", &mode, nvinfer1::PluginFieldType::kCHAR, 1)); fc.nbFields = f.size(); @@ -110,7 +98,7 @@ bool AdaptivePoolingConverter( TORCHTRT_CHECK(new_layer, "Unable to create pooling (interpolation) plugin from node" << *n); new_layer->setName(util::node_info(n).c_str()); - auto layer_output = addUnpadding(ctx, n, new_layer->getOutput(0), orig_dims.nbDims, false, false); + auto layer_output = new_layer->getOutput(0); ctx->AssociateValueAndTensor(n->outputs()[0], layer_output); LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions()); @@ -238,15 +226,15 @@ auto pooling_registrations TORCHTRT_UNUSED = }}) .pattern({"aten::adaptive_avg_pool1d(Tensor self, int[1] output_size) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kAVERAGE); + return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kAVERAGE, "adaptive_avg_pool1d"); }}) .pattern({"aten::adaptive_avg_pool2d(Tensor self, int[2] output_size) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kAVERAGE); + return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kAVERAGE, "adaptive_avg_pool2d"); }}) .pattern({"aten::adaptive_max_pool2d(Tensor self, int[2] output_size) -> (Tensor, Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kMAX); + return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kMAX, "adaptive_max_pool2d"); }}); } // namespace } // namespace impl diff --git a/core/plugins/impl/interpolate_plugin.cpp b/core/plugins/impl/interpolate_plugin.cpp index ee1d77df1c..3568a7481e 100644 --- a/core/plugins/impl/interpolate_plugin.cpp +++ b/core/plugins/impl/interpolate_plugin.cpp @@ -289,6 +289,8 @@ int InterpolatePlugin::enqueue( out = at::upsample_bilinear2d(input, {size_[0], size_[1]}, align_corners_); } else if (mode_ == "trilinear") { out = at::upsample_trilinear3d(input, {size_[0], size_[1], size_[2]}, align_corners_); + } else if(mode_ == "adaptive_avg_pool1d"){ + out = at::adaptive_avg_pool1d(input, {size_[0]}); } else if (mode_ == "adaptive_avg_pool2d") { out = at::adaptive_avg_pool2d(input, {size_[0], size_[1]}); } else if (mode_ == "adaptive_max_pool2d") { diff --git a/tests/core/conversion/converters/test_pooling.cpp b/tests/core/conversion/converters/test_pooling.cpp index a8c1cad760..55bf88506b 100644 --- a/tests/core/conversion/converters/test_pooling.cpp +++ b/tests/core/conversion/converters/test_pooling.cpp @@ -540,6 +540,32 @@ TEST(Converters, ATenAdaptiveAvgPool1DGlobalPoolingConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); } +TEST(Converters, ATenAdaptiveAvgPool1DUsingPluginConvertsCorrectly) { + const auto graph = + R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=3]() + %6 : int[] = prim::ListConstruct(%2) + %10 : Tensor = aten::adaptive_avg_pool1d(%0, %6) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch adaptive_avg_pool1d needs a 3D input or a 2D input + auto in = at::randint(-5, 5, {1, 3, 16}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + TEST(Converters, ATenAdaptiveMaxPool2DConvertsCorrectly) { const auto graph = R"IR( graph(%0 : Tensor): From e554dbd118b9fc8e5a07bfef28490e8bd81bf1e0 Mon Sep 17 00:00:00 2001 From: Ruoqian Guo Date: Fri, 31 Dec 2021 02:18:52 +0000 Subject: [PATCH 3/5] feat: support aten::adaptive_max_pool1d, aten::adaptive_avg_pool3d and aten::adaptive_max_pool3d operators Signed-off-by: Ruoqian Guo --- core/conversion/converters/impl/pooling.cpp | 22 ++- core/plugins/impl/interpolate_plugin.cpp | 8 +- .../conversion/converters/test_pooling.cpp | 164 ++++++++++++++++++ 3 files changed, 190 insertions(+), 4 deletions(-) diff --git a/core/conversion/converters/impl/pooling.cpp b/core/conversion/converters/impl/pooling.cpp index f18ef91c1b..58c5004ee0 100644 --- a/core/conversion/converters/impl/pooling.cpp +++ b/core/conversion/converters/impl/pooling.cpp @@ -37,7 +37,8 @@ bool AdaptivePoolingConverter( ConversionCtx* ctx, const torch::jit::Node* n, args& args, - nvinfer1::PoolingType pool_type, const std::string& mode) { + nvinfer1::PoolingType pool_type, + const std::string& mode) { auto in = args[0].ITensorOrFreeze(ctx); auto out_size = util::toDims(args[1].unwrapToIntList()); @@ -226,15 +227,30 @@ auto pooling_registrations TORCHTRT_UNUSED = }}) .pattern({"aten::adaptive_avg_pool1d(Tensor self, int[1] output_size) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kAVERAGE, "adaptive_avg_pool1d"); + return AdaptivePoolingConverter( + ctx, n, args, nvinfer1::PoolingType::kAVERAGE, "adaptive_avg_pool1d"); + }}) + .pattern({"aten::adaptive_max_pool1d(Tensor self, int[2] output_size) -> (Tensor, Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kMAX, "adaptive_max_pool1d"); }}) .pattern({"aten::adaptive_avg_pool2d(Tensor self, int[2] output_size) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kAVERAGE, "adaptive_avg_pool2d"); + return AdaptivePoolingConverter( + ctx, n, args, nvinfer1::PoolingType::kAVERAGE, "adaptive_avg_pool2d"); }}) .pattern({"aten::adaptive_max_pool2d(Tensor self, int[2] output_size) -> (Tensor, Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kMAX, "adaptive_max_pool2d"); + }}) + .pattern({"aten::adaptive_avg_pool3d(Tensor self, int[3] output_size) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + return AdaptivePoolingConverter( + ctx, n, args, nvinfer1::PoolingType::kAVERAGE, "adaptive_avg_pool3d"); + }}) + .pattern({"aten::adaptive_max_pool3d(Tensor self, int[3] output_size) -> (Tensor, Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + return AdaptivePoolingConverter(ctx, n, args, nvinfer1::PoolingType::kMAX, "adaptive_max_pool3d"); }}); } // namespace } // namespace impl diff --git a/core/plugins/impl/interpolate_plugin.cpp b/core/plugins/impl/interpolate_plugin.cpp index 3568a7481e..0eb89c9011 100644 --- a/core/plugins/impl/interpolate_plugin.cpp +++ b/core/plugins/impl/interpolate_plugin.cpp @@ -289,12 +289,18 @@ int InterpolatePlugin::enqueue( out = at::upsample_bilinear2d(input, {size_[0], size_[1]}, align_corners_); } else if (mode_ == "trilinear") { out = at::upsample_trilinear3d(input, {size_[0], size_[1], size_[2]}, align_corners_); - } else if(mode_ == "adaptive_avg_pool1d"){ + } else if (mode_ == "adaptive_avg_pool1d") { out = at::adaptive_avg_pool1d(input, {size_[0]}); + } else if (mode_ == "adaptive_max_pool1d") { + out = std::get<0>(at::adaptive_max_pool1d(input, {size_[0]})); } else if (mode_ == "adaptive_avg_pool2d") { out = at::adaptive_avg_pool2d(input, {size_[0], size_[1]}); } else if (mode_ == "adaptive_max_pool2d") { out = std::get<0>(at::adaptive_max_pool2d(input, {size_[0], size_[1]})); + } else if (mode_ == "adaptive_avg_pool3d") { + out = at::adaptive_avg_pool3d(input, {size_[0], size_[1], size_[2]}); + } else if (mode_ == "adaptive_max_pool3d") { + out = std::get<0>(at::adaptive_max_pool3d(input, {size_[0], size_[1], size_[2]})); } } diff --git a/tests/core/conversion/converters/test_pooling.cpp b/tests/core/conversion/converters/test_pooling.cpp index 55bf88506b..743f786c0c 100644 --- a/tests/core/conversion/converters/test_pooling.cpp +++ b/tests/core/conversion/converters/test_pooling.cpp @@ -566,6 +566,58 @@ TEST(Converters, ATenAdaptiveAvgPool1DUsingPluginConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); } +TEST(Converters, ATenAdaptiveMaxPool1DGlobalPoolingConvertsCorrectly) { + const auto graph = + R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%2) + %10 : Tensor, %11 : Tensor = aten::adaptive_max_pool1d(%0, %6) + return (%10, %11))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch adaptive_max_pool1d needs a 3D input or a 2D input + auto in = at::randint(-5, 5, {1, 3, 16}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, ATenAdaptiveMaxPool1DUsingPluginConvertsCorrectly) { + const auto graph = + R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=3]() + %6 : int[] = prim::ListConstruct(%2) + %10 : Tensor, %11 : Tensor = aten::adaptive_max_pool1d(%0, %6) + return (%10, %11))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch adaptive_max_pool1d needs a 3D input or a 2D input + auto in = at::randint(-5, 5, {1, 3, 16}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + TEST(Converters, ATenAdaptiveMaxPool2DConvertsCorrectly) { const auto graph = R"IR( graph(%0 : Tensor): @@ -617,3 +669,115 @@ TEST(Converters, ATenAdaptiveMaxPool2DConvertsCorrectlyWithDynamicInput) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); } + +TEST(Converters, ATenAdaptiveAvgPool3DGlobalPoolingConvertsCorrectly) { + const auto graph = + R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=1]() + %4 : int = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%2, %3, %4) + %10 : Tensor = aten::adaptive_avg_pool3d(%0, %6) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch adaptive_avg_pool3d needs a 5D input or a 4D input + auto in = at::randint(-5, 5, {4, 5, 3, 15, 16}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, ATenAdaptiveAvgPool3DUsingPluginConvertsCorrectly) { + const auto graph = + R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=7]() + %3 : int = prim::Constant[value=6]() + %4 : int = prim::Constant[value=5]() + %6 : int[] = prim::ListConstruct(%2, %3, %4) + %10 : Tensor = aten::adaptive_avg_pool3d(%0, %6) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch adaptive_avg_pool3d needs a 5D input or a 4D input + auto in = at::randint(-5, 5, {4, 5, 3, 15, 16}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, ATenAdaptiveMaxPool3DGlobalPoolingConvertsCorrectly) { + const auto graph = + R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=1]() + %4 : int = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%2, %3, %4) + %10 : Tensor, %11 : Tensor = aten::adaptive_max_pool3d(%0, %6) + return (%10, %11))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch adaptive_max_pool3d needs a 5D input or a 4D input + auto in = at::randint(-5, 5, {5, 3, 15, 16}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, ATenAdaptiveMaxPool3DUsingPluginConvertsCorrectly) { + const auto graph = + R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=7]() + %3 : int = prim::Constant[value=8]() + %4 : int = prim::Constant[value=9]() + %6 : int[] = prim::ListConstruct(%2, %3, %4) + %10 : Tensor, %11 : Tensor = aten::adaptive_max_pool3d(%0, %6) + return (%10, %11))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, g.get()); + + // PyTorch adaptive_max_pool3d needs a 5D input or a 4D input + auto in = at::randint(-5, 5, {4, 5, 3, 15, 16}, at::kCUDA); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} From 15aa0985d279d93aa4cd22e0093d964ee8ee2f1a Mon Sep 17 00:00:00 2001 From: Ruoqian Guo Date: Fri, 31 Dec 2021 03:45:59 +0000 Subject: [PATCH 4/5] chore: update reduceAxes variable in GlobalPoolingConverter Signed-off-by: Ruoqian Guo --- core/conversion/converters/impl/pooling.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/core/conversion/converters/impl/pooling.cpp b/core/conversion/converters/impl/pooling.cpp index 58c5004ee0..c8514c4555 100644 --- a/core/conversion/converters/impl/pooling.cpp +++ b/core/conversion/converters/impl/pooling.cpp @@ -16,9 +16,13 @@ bool GlobalPoolingConverter( nvinfer1::PoolingType pool_type) { auto in = args[0].ITensorOrFreeze(ctx); nvinfer1::Dims dims = in->getDimensions(); - auto out_size = util::toDims(args[1].unwrapToIntList()); - // Generate a bitmask of all 1s except the last 2 bits (N and C axes) - uint32_t reduceAxes = ((1 << dims.nbDims) - 1) ^ ((1 << (dims.nbDims - out_size.nbDims)) - 1); + // Generate a bitmask of all 1s except the last 2 bits (N and C axes) when dims.nbDims >= 2 + uint32_t reduceAxes = ((1 << dims.nbDims) - 1) & ~0b11; + // Generate a bitmask of all 1s except the last 1 bits (N axes) when dims.nbDims == 2. `aten::adaptive_avg_pool1d`'s + // input can be (N, C, L) or (C, L). + if (dims.nbDims == 2) { + reduceAxes = ((1 << dims.nbDims) - 1) & ~0b1; + } auto* new_layer = ctx->net->addReduce( *in, pool_type == nvinfer1::PoolingType::kMAX ? nvinfer1::ReduceOperation::kMAX : nvinfer1::ReduceOperation::kAVG, From 143fc3b8fdef9af07c1935f11d97e166199e63ed Mon Sep 17 00:00:00 2001 From: Ruoqian Guo Date: Fri, 31 Dec 2021 03:49:13 +0000 Subject: [PATCH 5/5] chore: update annotation in line 19 of pooling.cpp Signed-off-by: Ruoqian Guo --- core/conversion/converters/impl/pooling.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/conversion/converters/impl/pooling.cpp b/core/conversion/converters/impl/pooling.cpp index c8514c4555..0c950a5d3a 100644 --- a/core/conversion/converters/impl/pooling.cpp +++ b/core/conversion/converters/impl/pooling.cpp @@ -16,7 +16,7 @@ bool GlobalPoolingConverter( nvinfer1::PoolingType pool_type) { auto in = args[0].ITensorOrFreeze(ctx); nvinfer1::Dims dims = in->getDimensions(); - // Generate a bitmask of all 1s except the last 2 bits (N and C axes) when dims.nbDims >= 2 + // Generate a bitmask of all 1s except the last 2 bits (N and C axes) when dims.nbDims > 2 uint32_t reduceAxes = ((1 << dims.nbDims) - 1) & ~0b11; // Generate a bitmask of all 1s except the last 1 bits (N axes) when dims.nbDims == 2. `aten::adaptive_avg_pool1d`'s // input can be (N, C, L) or (C, L).