diff --git a/README.md b/README.md index 78789a9451..9cd32c9533 100644 --- a/README.md +++ b/README.md @@ -80,7 +80,8 @@ trt_ts_module = torch_tensorrt.compile(torch_script_module, # For static size shape=[1, 3, 224, 224] dtype=torch.half) # Datatype of input tensor. Allowed options torch.(float|half|int8|int32|bool) ], - enabled_precisions = {torch.half}, # Run with FP16) + enabled_precisions = {torch.half}, # Run with FP16 +) result = trt_ts_module(input_data) # run inference torch.jit.save(trt_ts_module, "trt_torchscript_module.ts") # save the TRT embedded Torchscript diff --git a/core/conversion/converters/impl/element_wise.cpp b/core/conversion/converters/impl/element_wise.cpp index 0ad4c128f6..32c7050289 100644 --- a/core/conversion/converters/impl/element_wise.cpp +++ b/core/conversion/converters/impl/element_wise.cpp @@ -25,8 +25,6 @@ nvinfer1::ITensor* clamp_util( return clamp_layer_out; } - - auto element_wise_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() .pattern( diff --git a/core/conversion/converters/impl/select.cpp b/core/conversion/converters/impl/select.cpp index 20a03f6f5e..8025a1086b 100644 --- a/core/conversion/converters/impl/select.cpp +++ b/core/conversion/converters/impl/select.cpp @@ -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); } @@ -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)", @@ -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); @@ -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; }}) @@ -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()); @@ -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()); diff --git a/core/conversion/converters/impl/shuffle.cpp b/core/conversion/converters/impl/shuffle.cpp index c1f45707f0..2df7e653ef 100644 --- a/core/conversion/converters/impl/shuffle.cpp +++ b/core/conversion/converters/impl/shuffle.cpp @@ -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 out_shape; - if (ctx->input_is_dynamic && in_shape[0] != -1) { - out_shape = std::vector({in_shape[0], -1}); - } else if (ctx->input_is_dynamic && in_shape[0] == -1) { - out_shape = std::vector( - {-1, - -1 * std::accumulate(std::begin(in_shape), std::end(in_shape), 1, std::multiplies())}); + 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()); + 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(); } @@ -45,7 +71,16 @@ static auto shuffle_registrations TORCHTRT_UNUSED = auto in_shape = util::toVec(in->getDimensions()); std::vector 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(); }