Skip to content

Commit

Permalink
feat: support aten::adaptive_max_pool1d, aten::adaptive_avg_pool3d an…
Browse files Browse the repository at this point in the history 
…d aten::adaptive_max_pool3d operators

Signed-off-by: Ruoqian Guo <[email protected]>
  • Loading branch information
@ruoqianguo
ruoqianguo committed Dec 31, 2021
1 parent deb9f74 commit e554dbd
Show file tree
Hide file tree
Showing 3 changed files with 190 additions and 4 deletions.
22 changes: 19 additions & 3 deletions core/conversion/converters/impl/pooling.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());

Expand Down Expand Up @@ -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
Expand Down
8 changes: 7 additions & 1 deletion core/plugins/impl/interpolate_plugin.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]}));
}
}

Expand Down
164 changes: 164 additions & 0 deletions tests/core/conversion/converters/test_pooling.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::Graph>();
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::Graph>();
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):
Expand Down Expand Up @@ -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::Graph>();
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::Graph>();
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::Graph>();
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::Graph>();
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));
}

0 comments on commit e554dbd

Please sign in to comment.