Skip to content

Add support for aten::squeeze without a dim #1393

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 7, 2022
Merged

Add support for aten::squeeze without a dim #1393

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 43 additions & 21 deletions core/conversion/converters/impl/squeeze.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,35 +14,57 @@ namespace converters {
namespace impl {
namespace {

auto squeeze_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
{"aten::squeeze.dim(Tensor(a) self, int dim) -> (Tensor(a))",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
auto dim = args[1].unwrapToInt();
auto squeeze_registrations TORCHTRT_UNUSED =
RegisterNodeConversionPatterns()
.pattern(
{"aten::squeeze.dim(Tensor(a) self, int dim) -> (Tensor(a))",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
auto dim = args[1].unwrapToInt();

auto selfDim = util::toVec(self->getDimensions());
if (dim < 0) {
dim = selfDim.size() + dim;
}
auto selfDim = util::toVec(self->getDimensions());
if (dim < 0) {
dim = selfDim.size() + dim;
}

if (selfDim[dim] != 1) {
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], self);
if (selfDim[dim] != 1) {
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], self);

LOG_DEBUG("Output tensor shape: " << out->getDimensions());
LOG_DEBUG("Output tensor shape: " << out->getDimensions());

return true;
}
return true;
}

auto shuffle_layer = ctx->net->addShuffle(*self);
TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
shuffle_layer->setReshapeDimensions(util::squeezeDims(self->getDimensions(), dim));
auto shuffle_layer = ctx->net->addShuffle(*self);
TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
shuffle_layer->setReshapeDimensions(util::squeezeDims(self->getDimensions(), dim));

auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_layer->getOutput(0));
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_layer->getOutput(0));

LOG_DEBUG("Output tensor shape: " << out->getDimensions());
LOG_DEBUG("Output tensor shape: " << out->getDimensions());

return true;
}});
return true;
}})
.pattern(
{"aten::squeeze(Tensor(a) self) -> (Tensor(a))",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
auto self_dims = self->getDimensions();
auto out = self;
auto squeeze_dims = util::squeezeAllDims(self_dims);
if (squeeze_dims != self_dims) {
auto shuffle_layer = ctx->net->addShuffle(*self);
TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
shuffle_layer->setReshapeDimensions(squeeze_dims);
out = shuffle_layer->getOutput(0);
}

auto trt_out = ctx->AssociateValueAndTensor(n->outputs()[0], out);

LOG_DEBUG("Output tensor shape: " << trt_out->getDimensions());

return true;
}});

} // namespace
} // namespace impl
Expand Down
13 changes: 13 additions & 0 deletions core/util/trt_util.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,19 @@ nvinfer1::Dims squeezeDims(const nvinfer1::Dims& d, int pos, bool use_zeros) {
return dims;
}

nvinfer1::Dims squeezeAllDims(const nvinfer1::Dims& d, bool use_zeros_for_unknown_dims) {
nvinfer1::Dims dims;
int j = 0;
for (int i = 0; i < d.nbDims; i++) {
if (d.d[i] != 1) {
dims.d[j++] = (use_zeros_for_unknown_dims && d.d[i] == -1) ? 0 : d.d[i];
}
}
dims.nbDims = j;

return dims;
}

std::vector<int64_t> toVec(nvinfer1::Dims d) {
std::vector<int64_t> dims;
for (int i = 0; i < d.nbDims; i++) {
Expand Down
1 change: 1 addition & 0 deletions core/util/trt_util.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ nvinfer1::Dims toDimsTailPad(c10::List<int64_t> l, uint64_t pad_to);
nvinfer1::Dims unpadDims(const nvinfer1::Dims& d);
nvinfer1::Dims unsqueezeDims(const nvinfer1::Dims& d, int pos, int val = 1, bool use_zeros = true);
nvinfer1::Dims squeezeDims(const nvinfer1::Dims& d, int pos, bool use_zeros = true);
nvinfer1::Dims squeezeAllDims(const nvinfer1::Dims& d, bool use_zeros_for_unknown_dims = true);
nvinfer1::Dims toDims(c10::IntArrayRef l);
nvinfer1::Dims toDims(c10::List<int64_t> l);
nvinfer1::DimsHW toDimsHW(c10::List<int64_t> l);
Expand Down
26 changes: 26 additions & 0 deletions tests/core/conversion/converters/test_squeeze.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,3 +56,29 @@ TEST(Converters, ATenSqueezeDontNeedSqueezeConvertsCorrectly) {
ASSERT_TRUE(
torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
}

TEST(Converters, ATenSqueezeNoDimConvertsCorrectly) {
const auto graph = R"IR(
graph(%0 : Tensor):
%1 : Tensor = aten::squeeze(%0)
return (%1))IR";

auto g = std::make_shared<torch::jit::Graph>();
torch::jit::parseIR(graph, g.get());

auto validate_squeeze_with_input = [&g](const at::Tensor& in) {
auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in});

params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in});
ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
};

validate_squeeze_with_input(at::randint(1, 10, {2, 1, 3, 3}, {at::kCUDA}));
validate_squeeze_with_input(at::randint(1, 10, {1, 1, 1, 3}, {at::kCUDA}));
validate_squeeze_with_input(at::randint(1, 10, {1, 10, 1, 3}, {at::kCUDA}));
validate_squeeze_with_input(at::randint(1, 10, {2, 10, 3, 3}, {at::kCUDA}));
validate_squeeze_with_input(at::randint(1, 10, {1, 1}, {at::kCUDA}));
validate_squeeze_with_input(at::randint(1, 10, {1}, {at::kCUDA}));
}