Skip to content

Support swin/bert with dynamic batch #1270

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
Aug 16, 2022
Merged

Support swin/bert with dynamic batch #1270

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
2 changes: 0 additions & 2 deletions core/conversion/converters/impl/element_wise.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,6 @@ nvinfer1::ITensor* clamp_util(
return clamp_layer_out;
}



auto element_wise_registrations TORCHTRT_UNUSED =
RegisterNodeConversionPatterns()
.pattern(
Expand Down
51 changes: 28 additions & 23 deletions core/conversion/converters/impl/select.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,8 @@ auto select_registrations TORCHTRT_UNUSED =
// IShuffleLayer removes redundant dimensions
auto shuffle_layer = ctx->net->addShuffle(*out);
TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
shuffle_layer->setReshapeDimensions(util::squeezeDims(out->getDimensions(), dim));
shuffle_layer->setReshapeDimensions(
util::squeezeDims(out->getDimensions(), dim, !ctx->input_is_dynamic));
shuffle_layer->setName(util::node_info(n).c_str());
out = shuffle_layer->getOutput(0);
}
Expand Down Expand Up @@ -249,21 +250,19 @@ auto select_registrations TORCHTRT_UNUSED =
auto dims = args[2].unwrapToIntList().vec();

TORCHTRT_CHECK(dims.size() == shifts.size(), "dims.size() should be equal to shifts.size()");
if (ctx->input_is_dynamic) {
TORCHTRT_THROW_ERROR("aten::roll is currently not support in dynamic input shape compilation");
} else {
auto in_shape = util::toVec(in->getDimensions());
for (size_t i = 0; i < dims.size(); i++) {
auto dim = dims[i] < 0 ? (in_shape.size() + dims[i]) : dims[i];
TORCHTRT_CHECK(dim < in_shape.size(), "Dimension out of range");
in = roll(ctx, in, shifts[i], dim, in_shape);
}
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
auto in_shape = util::toVec(in->getDimensions());
for (size_t i = 0; i < dims.size(); i++) {
auto dim = dims[i] < 0 ? (in_shape.size() + dims[i]) : dims[i];
TORCHTRT_CHECK(dim < in_shape.size(), "Dimension out of range");
TORCHTRT_CHECK(
in_shape[dim] != -1, "aten::roll is not supported when the targeted dimension is dynamic");
in = roll(ctx, in, shifts[i], dim, in_shape);
}
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);

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

return true;
}
return true;
}})
.pattern(
{"aten::index.Tensor(Tensor self, Tensor?[] indices) -> (Tensor)",
Expand Down Expand Up @@ -360,9 +359,15 @@ auto select_registrations TORCHTRT_UNUSED =
stride_.d[i] = 1;
}
}
auto slice_layer = ctx->net->addSlice(*in, start_, size_, stride_);

if (dynamic_shape) { // dynamic shape
if (!dynamic_shape) {
auto slice_layer = ctx->net->addSlice(*in, start_, size_, stride_);
LOG_DEBUG("start_:" << start_);
LOG_DEBUG("size_:" << size_);
LOG_DEBUG("stride_:" << stride_);
auto slice_out = slice_layer->getOutput(0);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], slice_out);
LOG_DEBUG("Slice layer output shape: " << out->getDimensions());
} else { // dynamic shape
LOG_DEBUG("Using dynamic version of slice");
// start tensor
at::Tensor start_tensor = torch::zeros({nbdims}).to(torch::kI32);
Expand Down Expand Up @@ -398,13 +403,13 @@ auto select_registrations TORCHTRT_UNUSED =
auto size_itensor = get_slice_size(ctx, out_start, out_end, stride_itensor, nbdims, node_name);

// update slice layer
auto slice_layer = ctx->net->addSlice(*in, start_, size_, stride_);
slice_layer->setInput(1, *out_start); // start
slice_layer->setInput(2, *size_itensor); // size, must be set if input is dynamic
auto slice_out = slice_layer->getOutput(0);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], slice_out);
LOG_DEBUG("Slice layer output shape: " << out->getDimensions());
}
auto slice_out = slice_layer->getOutput(0);

auto out = ctx->AssociateValueAndTensor(n->outputs()[0], slice_out);
LOG_DEBUG("Slice layer output shape: " << out->getDimensions());

return true;
}})
Expand Down Expand Up @@ -484,7 +489,7 @@ auto select_registrations TORCHTRT_UNUSED =

auto layer = ctx->net->addScatter(*self, *index, *value_tensor, nvinfer1::ScatterMode::kELEMENT);
layer->setAxis(dim);

TORCHTRT_CHECK(layer, "Unable to create layer for aten::scatter.value");

layer->setName(util::node_info(n).c_str());
Expand All @@ -503,7 +508,7 @@ auto select_registrations TORCHTRT_UNUSED =

auto layer = ctx->net->addScatter(*self, *index, *src, nvinfer1::ScatterMode::kELEMENT);
layer->setAxis(dim);

TORCHTRT_CHECK(layer, "Unable to create layer for aten::scatter.src");

layer->setName(util::node_info(n).c_str());
Expand Down
49 changes: 42 additions & 7 deletions core/conversion/converters/impl/shuffle.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,38 @@ static auto shuffle_registrations TORCHTRT_UNUSED =
auto end_dim = args[2].unwrapToInt();
auto in_shape = util::toVec(in->getDimensions());
std::vector<int64_t> out_shape;
if (ctx->input_is_dynamic && in_shape[0] != -1) {
out_shape = std::vector<int64_t>({in_shape[0], -1});
} else if (ctx->input_is_dynamic && in_shape[0] == -1) {
out_shape = std::vector<int64_t>(
{-1,
-1 * std::accumulate(std::begin(in_shape), std::end(in_shape), 1, std::multiplies<int64_t>())});
if (ctx->input_is_dynamic) {
end_dim = (end_dim == -1) ? in_shape.size() - 1 : end_dim;
int nbDynamicFlattenedDims = 0;
int nbDynamicUnflattenedDims = 0;
for (int i = 0; i < (int)in_shape.size(); i++) {
if (in_shape[i] == -1) {
if (i >= start_dim && i <= end_dim)
nbDynamicFlattenedDims++;
else
nbDynamicUnflattenedDims++;
}
}
if (nbDynamicFlattenedDims > 0 && nbDynamicUnflattenedDims > 0) {
TORCHTRT_THROW_ERROR(
"Flatten is currently not supported when target shape contains more than one dynamic dimension");
}
if (nbDynamicUnflattenedDims > 1) {
TORCHTRT_THROW_ERROR(
"Flatten is currently not supported when target shape contains more than one dynamic dimension");
}
out_shape = in_shape;
out_shape.erase(std::begin(out_shape) + start_dim, std::begin(out_shape) + end_dim + 1);
if (nbDynamicFlattenedDims == 0) {
auto flattened_dim = std::accumulate(
std::begin(in_shape) + start_dim,
std::begin(in_shape) + end_dim + 1,
1,
std::multiplies<int64_t>());
out_shape.insert(std::begin(out_shape) + start_dim, flattened_dim);
} else {
out_shape.insert(std::begin(out_shape) + start_dim, -1);
}
} else {
out_shape = torch::flatten(torch::rand(in_shape), start_dim, end_dim).sizes().vec();
}
Expand All @@ -45,7 +71,16 @@ static auto shuffle_registrations TORCHTRT_UNUSED =
auto in_shape = util::toVec(in->getDimensions());
std::vector<int64_t> new_shape;
if (ctx->input_is_dynamic) {
TORCHTRT_THROW_ERROR("Resize is currently not support in dynamic input shape compilation");
new_shape = util::toVec(args[1].unwrapToIntList().vec());
int nbDynamicDims = 0;
for (size_t i = 0; i < new_shape.size(); i++) {
if (in_shape[i] == -1)
nbDynamicDims++;
}
if (nbDynamicDims > 1) {
TORCHTRT_THROW_ERROR(
"Resize is currently not supported when target shape contains more than one dynamic dimension");
}
} else {
new_shape = torch::reshape(torch::rand(in_shape), args[1].unwrapToIntList().vec()).sizes().vec();
}
Expand Down