From f42562b33e328fe24d3266c0f6b15dd4dd2a5b19 Mon Sep 17 00:00:00 2001 From: Abhiram Iyer Date: Fri, 5 Jun 2020 14:03:54 -0700 Subject: [PATCH] fix(//core/conversion/converters/impl): 1d case not working Signed-off-by: Abhiram Iyer Signed-off-by: Abhiram Iyer --- .../converters/impl/interpolate.cpp | 38 ++++++++++++++++--- 1 file changed, 33 insertions(+), 5 deletions(-) diff --git a/core/conversion/converters/impl/interpolate.cpp b/core/conversion/converters/impl/interpolate.cpp index 3adcd42a56..5c0f4b6413 100755 --- a/core/conversion/converters/impl/interpolate.cpp +++ b/core/conversion/converters/impl/interpolate.cpp @@ -15,10 +15,19 @@ auto interpolate_registrations = RegisterNodeConversionPatterns() .pattern({ "aten::upsample_nearest1d(Tensor self, int[1] output_size, float? scales=None) -> (Tensor)", [](ConversionCtx* ctx, const torch::jit::Node*n, args& args) -> bool { - std::cout << "GOT IN HERE!!!!!!" << std::endl; - auto in = args[0].ITensor(); auto in_shape = util::toVec(in->getDimensions()); + + // remove padding that TensorRt adds automatically + if (in_shape.size() >= 4) { + // remove first dimension + in_shape.erase(in_shape.begin()); + + auto shuffle = ctx->net->addShuffle(*in); + shuffle->setReshapeDimensions(util::toDims(in_shape)); + shuffle->setName( (util::node_info(n) + " [Reshape to " + util::toStr(util::toDims(in_shape)) + "]").c_str() ); + in = shuffle->getOutput(0); + } // Case 1: user uses output size and not scales if (!args[1].IValue()->isNone() && args[2].IValue()->isNone()) { @@ -28,9 +37,6 @@ auto interpolate_registrations = RegisterNodeConversionPatterns() auto out_shape = in_shape; std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size())); - - // remove padding that TensorRT adds automatically - // out_shape.erase(out_shape.begin(), out_shape.begin()+1); auto resize_layer = ctx->net->addResize(*in); TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n); @@ -39,8 +45,30 @@ auto interpolate_registrations = RegisterNodeConversionPatterns() resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST); resize_layer->setName(util::node_info(n).c_str()); + // auto out_tensor = resize_layer->getOutput(0); + // out_shape.erase(out_shape.begin()); + // auto shuffle = ctx->net->addShuffle(*out_tensor); + // shuffle->setReshapeDimensions(util::toDims(out_shape)); + // shuffle->setName( (util::node_info(n) + " [Reshape to " + util::toStr(util::toDims(out_shape)) + "]").c_str() ); + // auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0)); + // LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions()); + + // std::cout << "PRINTING STUFF AT THE END!" << std::endl; + // auto final = util::toVec(shuffle->getOutput(0)->getDimensions()); + // for (auto iter = final.begin(); iter != final.end(); iter++) { + // std::cout << *iter << std::endl; + // } + + //std::raise(SIGABRT); + auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0)); LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions()); + + // std::cout << "PRINTING STUFF AT THE END!" << std::endl; + // auto final = util::toVec(resize_layer->getOutput(0)->getDimensions()); + // for (auto iter = final.begin(); iter != final.end(); iter++) { + // std::cout << *iter << std::endl; + // } } else { LOG_DEBUG("scale factor parameter not supported yet."); }