From a90e6db985ee3a7f708b530d440b47b104fd8441 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Mon, 15 Jun 2020 10:58:35 -0700 Subject: [PATCH] fix(pooling): fix the tests and the 1D pooling cases Signed-off-by: Naren Dasan Signed-off-by: Naren Dasan --- core/conversion/converters/impl/pooling.cpp | 93 +++++++-- tests/core/converters/test_pooling.cpp | 208 ++++++++++++++++++-- 2 files changed, 271 insertions(+), 30 deletions(-) diff --git a/core/conversion/converters/impl/pooling.cpp b/core/conversion/converters/impl/pooling.cpp index 141677fe2b..5ecc9443af 100644 --- a/core/conversion/converters/impl/pooling.cpp +++ b/core/conversion/converters/impl/pooling.cpp @@ -23,23 +23,23 @@ bool MaxPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& ar } - auto kernel_size = util::toDimsHW(args[1].unwrapToIntList()); + auto kernel_size = util::toDims(args[1].unwrapToIntList()); LOG_DEBUG("kernel_size: " << kernel_size); - auto padding = util::toDimsHW(args[3].unwrapToIntList()); + auto padding = util::toDims(args[3].unwrapToIntList()); LOG_DEBUG("padding: " << padding); auto stride = util::toDims(args[2].unwrapToIntList()); LOG_DEBUG("stride: " << stride); auto dilation = util::toDims(args[4].unwrapToIntList()); - TRTORCH_ASSERT(dilation == util::toDims(std::vector({1,1})), "Pooling dilation is not supported in TensorRT"); + TRTORCH_ASSERT(dilation == util::toDims(std::vector(dilation.nbDims, 1)), "Pooling dilation is not supported in TensorRT"); LOG_DEBUG("dilation: " << dilation); LOG_WARNING("Dilation not used in max pooling converter"); bool ceil_mode = args[5].unwrapToBool(); auto new_layer = ctx->net->addPoolingNd(*in, nvinfer1::PoolingType::kMAX, kernel_size); - TRTORCH_CHECK(new_layer, "Unable to create Max Pool 2D layer from node: " << *n); + TRTORCH_CHECK(new_layer, "Unable to create Max Pooling layer from node: " << *n); new_layer->setName(util::node_info(n).c_str()); new_layer->setPaddingNd(padding); @@ -77,9 +77,9 @@ bool AvgPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& ar } - auto kernel_size = util::toDimsHW(args[1].unwrapToIntList()); + auto kernel_size = util::toDims(args[1].unwrapToIntList()); LOG_DEBUG("kernel_size: " << kernel_size); - auto padding = util::toDimsHW(args[3].unwrapToIntList()); + auto padding = util::toDims(args[3].unwrapToIntList()); LOG_DEBUG("padding: " << padding); auto stride = util::toDims(args[2].unwrapToIntList()); LOG_DEBUG("stride: " << stride); @@ -88,7 +88,7 @@ bool AvgPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& ar bool count_inlcude_pad = args[5].unwrapToBool(); auto new_layer = ctx->net->addPoolingNd(*in, nvinfer1::PoolingType::kAVERAGE, kernel_size); - TRTORCH_CHECK(new_layer, "Unable to create Avg Pool 2D layer from node: " << *n); + TRTORCH_CHECK(new_layer, "Unable to create Avg Pooling layer from node: " << *n); new_layer->setName(util::node_info(n).c_str()); new_layer->setPaddingNd(padding); @@ -118,12 +118,67 @@ bool AvgPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& ar auto pooling_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns() .pattern({ - "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)", + "aten::max_pool1d(Tensor self, int[1] kernel_size, int[1] stride=[], int[1] padding=[], int[1] dilation=[], bool ceil_mode=False) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - return MaxPoolingConverter(ctx, n, args); + auto in = args[0].ITensor(); + auto shape = util::toVec(in->getDimensions()); + + // Max Pool needs at least 4D input + if (shape.size() < 4) { + auto new_shape = util::toDimsPad(shape, 4); + LOG_DEBUG("Input shape is less than 4D got: " << util::toDims(shape) << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape); + auto shuffle = ctx->net->addShuffle(*in); + shuffle->setReshapeDimensions(new_shape); + shuffle->setName((util::node_info(n) + " [Reshape to " + util::toStr(new_shape) + ']').c_str()); + in = shuffle->getOutput(0); + } + + auto kernel_vec = args[1].unwrapToIntList().vec(); + kernel_vec.insert(kernel_vec.begin(), 1); + auto kernel_size = util::toDims(kernel_vec); + LOG_DEBUG("kernel_size: " << kernel_size); + auto stride_vec = args[2].unwrapToIntList().vec(); + stride_vec.insert(stride_vec.begin(), 1); + auto stride = util::toDims(stride_vec); + LOG_DEBUG("stride: " << stride); + auto padding_vec = args[3].unwrapToIntList().vec(); + padding_vec.insert(padding_vec.begin(), 0); + auto padding = util::toDims(padding_vec); + LOG_DEBUG("padding: " << padding); + + auto dilation = util::toDims(args[4].unwrapToIntList()); + + TRTORCH_ASSERT(dilation == util::toDims(std::vector(dilation.nbDims, 1)), "Pooling dilation is not supported in TensorRT"); + + LOG_DEBUG("dilation: " << dilation); + LOG_WARNING("Dilation not used in max pooling converter"); + bool ceil_mode = args[5].unwrapToBool(); + + auto new_layer = ctx->net->addPoolingNd(*in, nvinfer1::PoolingType::kMAX, kernel_size); + TRTORCH_CHECK(new_layer, "Unable to create Max Pooling layer from node: " << *n); + + new_layer->setName(util::node_info(n).c_str()); + new_layer->setPaddingNd(padding); + if (stride.nbDims != 2 && ctx->settings.device == nvinfer1::DeviceType::kDLA) { + if (!ctx->settings.allow_gpu_fallback) { + TRTORCH_THROW_ERROR("DLA Pooling stride is limited to 2D, allow GPU fallback"); + } else { + LOG_WARNING("DLA Pooling stride is limited to 2D, will run on GPU"); + } + } + new_layer->setStrideNd(stride); + + auto padding_mode = ceil_mode ? nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP : nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN; + new_layer->setPaddingMode(padding_mode); + + new_layer->setName(util::node_info(n).c_str()); + auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions()); + return true; } }).pattern({ - "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)", + "aten::avg_pool1d(Tensor self, int[1] kernel_size, int[1] stride=[], int[1] padding=0, bool ceil_mode=False, bool count_include_pad=True) -> Tensor", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { auto in = args[0].ITensor(); auto shape = util::toVec(in->getDimensions()); @@ -139,12 +194,18 @@ auto pooling_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns() } - auto kernel_size = util::toDimsHW(args[1].unwrapToIntList()); + auto kernel_vec = args[1].unwrapToIntList().vec(); + kernel_vec.insert(kernel_vec.begin(), 1); + auto kernel_size = util::toDims(kernel_vec); LOG_DEBUG("kernel_size: " << kernel_size); - auto padding = util::toDimsHW(args[3].unwrapToIntList()); - LOG_DEBUG("padding: " << padding); - auto stride = util::toDims(args[2].unwrapToIntList()); + auto stride_vec = args[2].unwrapToIntList().vec(); + stride_vec.insert(stride_vec.begin(), 1); + auto stride = util::toDims(stride_vec); LOG_DEBUG("stride: " << stride); + auto padding_vec = args[3].unwrapToIntList().vec(); + padding_vec.insert(padding_vec.begin(), 0); + auto padding = util::toDims(padding_vec); + LOG_DEBUG("padding: " << padding); bool ceil_mode = args[4].unwrapToBool(); bool count_inlcude_pad = args[5].unwrapToBool(); @@ -187,12 +248,12 @@ auto pooling_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns() return AvgPoolingConverter(ctx, n, args); } }).pattern({ - "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)", + "aten::max_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=[], int[3] dilation=[], bool ceil_mode=False) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { return MaxPoolingConverter(ctx, n, args); } }).pattern({ - "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)", + "aten::avg_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=[], bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { return AvgPoolingConverter(ctx, n, args); } diff --git a/tests/core/converters/test_pooling.cpp b/tests/core/converters/test_pooling.cpp index 939695ae46..5c85c88145 100644 --- a/tests/core/converters/test_pooling.cpp +++ b/tests/core/converters/test_pooling.cpp @@ -11,16 +11,16 @@ TEST(Converters, ATenMaxPool1DConvertsCorrectly) { %2 : int = prim::Constant[value=1]() %3 : int = prim::Constant[value=2]() %5 : bool = prim::Constant[value=0]() - %6 : int[] = prim::ListConstruct(%1, %1) - %7 : int[] = prim::ListConstruct(%2, %2) - %8 : int[] = prim::ListConstruct(%3, %3) - %10 : Tensor = aten::max_pool2d(%0, %8, %7, %6, %7, %5) + %6 : int[] = prim::ListConstruct(%1) + %7 : int[] = prim::ListConstruct(%2) + %8 : int[] = prim::ListConstruct(%3) + %10 : Tensor = aten::max_pool1d(%0, %8, %7, %6, %7, %5) return (%10))IR"; auto g = std::make_shared(); torch::jit::parseIR(graph, &*g); - auto in = at::randint(-5, 5, {1, 4, 4}, at::kCUDA); + auto in = at::randint(-5, 5, {1, 1, 10}, at::kCUDA); auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); @@ -48,7 +48,7 @@ TEST(Converters, ATenMaxPool2DConvertsCorrectly) { torch::jit::parseIR(graph, &*g); //PyTorch MaxPool needs a 3D input - auto in = at::randint(-5, 5, {1, 4, 4}, at::kCUDA); + auto in = at::randint(-5, 5, {1, 10, 10}, at::kCUDA); auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); @@ -66,17 +66,106 @@ TEST(Converters, ATenMaxPool3DConvertsCorrectly) { %2 : int = prim::Constant[value=1]() %3 : int = prim::Constant[value=2]() %5 : bool = prim::Constant[value=0]() - %6 : int[] = prim::ListConstruct(%1, %1) - %7 : int[] = prim::ListConstruct(%2, %2) - %8 : int[] = prim::ListConstruct(%3, %3) - %10 : Tensor = aten::max_pool2d(%0, %8, %7, %6, %7, %5) + %6 : int[] = prim::ListConstruct(%1, %1, %1) + %7 : int[] = prim::ListConstruct(%2, %2, %2) + %8 : int[] = prim::ListConstruct(%3, %3, %3) + %10 : Tensor = aten::max_pool3d(%0, %8, %7, %6, %7, %5) return (%10))IR"; auto g = std::make_shared(); torch::jit::parseIR(graph, &*g); //PyTorch MaxPool needs a 3D input - auto in = at::randint(-5, 5, {1, 4, 4, 4}, at::kCUDA); + auto in = at::randint(-5, 5, {1, 3, 10, 10, 10}, at::kCUDA); + auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in}); + + ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, ATenAvgPool1DConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %1 : int = prim::Constant[value=0]() + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=2]() + %4 : bool = prim::Constant[value=0]() + %5 : bool = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%1) + %7 : int[] = prim::ListConstruct(%2) + %8 : int[] = prim::ListConstruct(%3) + %9 : None = prim::Constant() + %10 : Tensor = aten::avg_pool1d(%0, %8, %7, %6, %4, %5) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + //PyTorch AvgPool needs a 3D input + auto in = at::randint(-5, 5, {1, 1, 10}, at::kCUDA); + auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in}); + + ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + + +TEST(Converters, ATenAvgPool1DCeilConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %1 : int = prim::Constant[value=0]() + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=2]() + %4 : bool = prim::Constant[value=0]() + %5 : bool = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%1) + %7 : int[] = prim::ListConstruct(%2) + %8 : int[] = prim::ListConstruct(%3) + %10 : Tensor = aten::avg_pool1d(%0, %8, %7, %6, %5, %5) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + //PyTorch AvgPool needs a 3D input + auto in = at::randint(-5, 5, {1, 1, 4}, at::kCUDA); + auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in}); + + ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, ATenAvgPool1DNoCountPadConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %1 : int = prim::Constant[value=0]() + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=2]() + %4 : bool = prim::Constant[value=0]() + %5 : bool = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%1) + %7 : int[] = prim::ListConstruct(%2) + %8 : int[] = prim::ListConstruct(%3) + %10 : Tensor = aten::avg_pool1d(%0, %8, %7, %6, %4, %4) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + //PyTorch AvgPool needs a 3D input + auto in = at::randint(-5, 5, {1, 1, 4}, at::kCUDA); auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); @@ -105,7 +194,7 @@ TEST(Converters, ATenAvgPool2DConvertsCorrectly) { auto g = std::make_shared(); torch::jit::parseIR(graph, &*g); - //PyTorch MaxPool needs a 3D input + //PyTorch AvgPool needs a 3D input auto in = at::randint(-5, 5, {1, 4, 4}, at::kCUDA); auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); @@ -136,7 +225,7 @@ TEST(Converters, ATenAvgPool2DCeilConvertsCorrectly) { auto g = std::make_shared(); torch::jit::parseIR(graph, &*g); - //PyTorch MaxPool needs a 3D input + //PyTorch AvgPool needs a 3D input auto in = at::randint(-5, 5, {1, 4, 4}, at::kCUDA); auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); @@ -166,7 +255,7 @@ TEST(Converters, ATenAvgPool2DNoCountPadConvertsCorrectly) { auto g = std::make_shared(); torch::jit::parseIR(graph, &*g); - //PyTorch MaxPool needs a 3D input + //PyTorch AvgPool needs a 3D input auto in = at::randint(-5, 5, {1, 4, 4}, at::kCUDA); auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); @@ -178,6 +267,97 @@ TEST(Converters, ATenAvgPool2DNoCountPadConvertsCorrectly) { ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); } +TEST(Converters, ATenAvgPool3DConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %1 : int = prim::Constant[value=0]() + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=2]() + %4 : bool = prim::Constant[value=0]() + %5 : bool = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%1, %1, %1) + %7 : int[] = prim::ListConstruct(%2, %2, %2) + %8 : int[] = prim::ListConstruct(%3, %3, %3) + %9 : None = prim::Constant() + %10 : Tensor = aten::avg_pool3d(%0, %8, %7, %6, %4, %5, %9) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + //PyTorch AvgPool needs a 3D input + auto in = at::randint(-5, 5, {1, 3, 4, 4, 4}, at::kCUDA); + auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in}); + + ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + + +TEST(Converters, ATenAvgPool3DCeilConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %1 : int = prim::Constant[value=0]() + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=2]() + %4 : bool = prim::Constant[value=0]() + %5 : bool = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%1, %1, %1) + %7 : int[] = prim::ListConstruct(%2, %2, %2) + %8 : int[] = prim::ListConstruct(%3, %3, %3) + %9 : None = prim::Constant() + %10 : Tensor = aten::avg_pool3d(%0, %8, %7, %6, %5, %5, %9) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + //PyTorch AvgPool needs a 3D input + auto in = at::randint(-5, 5, {1, 3, 4, 4, 4}, at::kCUDA); + auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in}); + + ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + +TEST(Converters, ATenAvgPool3DNoCountPadConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %1 : int = prim::Constant[value=0]() + %2 : int = prim::Constant[value=1]() + %3 : int = prim::Constant[value=2]() + %4 : bool = prim::Constant[value=0]() + %5 : bool = prim::Constant[value=1]() + %6 : int[] = prim::ListConstruct(%1, %1, %1) + %7 : int[] = prim::ListConstruct(%2, %2, %2) + %8 : int[] = prim::ListConstruct(%3, %3, %3) + %9 : None = prim::Constant() + %10 : Tensor = aten::avg_pool3d(%0, %8, %7, %6, %4, %4, %9) + return (%10))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + //PyTorch AvgPool needs a 3D input + auto in = at::randint(-5, 5, {1, 3, 4, 4, 4}, at::kCUDA); + auto params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto jit_results = trtorch::tests::util::RunGraph(g, params, {in}); + + in = at::clone(in); + params = trtorch::core::conversion::get_named_params(g->inputs(), {}); + auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in}); + + ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6)); +} + TEST(Converters, ATenAdaptiveAvgPool2DConvertsCorrectly) { const auto graph = R"IR( graph(%0 : Tensor):