refactor(//core/conversion): relax restrictions on input tensor types

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>

Author: narendasan

core/conversion/InterfaceTypes.cpp

@@ -10,10 +10,9 @@ namespace conversion {
 GraphParams get_named_params(c10::ArrayRef<torch::jit::Value*> inputs,
                              std::vector<at::Tensor> params) {
     GraphParams named_params;
-    auto type_lut = torch::jit::script::string_to_type_lut();
     auto param_it = params.begin();
     for (auto in : inputs) {
-        if (in->type() != type_lut["Tensor"]                            \
+        if (in->type() != c10::TensorType::get()                     \
             && in->isCompleteTensor() && param_it != params.end()) {
             named_params[in] = *param_it;
             ++param_it;
@@ -35,7 +34,7 @@ InputRange::InputRange(std::vector<int64_t> d) {
     min = util::toDims(d);
     max = util::toDims(d);
     input_shape = util::toDims(d);
-    
+
 }
 
 
@@ -48,14 +47,14 @@ InputRange::InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_
     sizes.insert(min_shape.size());
     sizes.insert(opt_shape.size());
     sizes.insert(max_shape.size());
-    
+
     if (sizes.size() != 1) {
         LOG_ERROR("Expected all input sizes have the same dimensions, but found dimensions: min(" \
                   << min_shape.size() << "), opt("
                   << opt_shape.size() << "), max("
                   << max_shape.size() << ")");
     }
-    
+
     min = util::toDimsPad(min_shape, 4);
     opt = util::toDimsPad(opt_shape, 4);
     max = util::toDimsPad(max_shape, 4);
@@ -72,9 +71,9 @@ InputRange::InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_
             dyn_shape.push_back(opt_shape[i]);
         }
     }
-    
+
     input_shape = util::toDimsPad(dyn_shape, 4);
-    
+
 }
 
 } // namespace conversion

core/conversion/conversion.cpp

@@ -117,16 +117,14 @@ void AddInputs(ConversionCtx* ctx,
                 at::ArrayRef<const torch::jit::Value*> inputs,
                 std::vector<InputRange>& input_dims) {
 
-    auto type_lut = torch::jit::script::string_to_type_lut();
     std::vector<const torch::jit::Value*> input_tensors;
     for (auto in : inputs) {
         // Disregarding inputs that are not tensors
         //
         // Ex.
         // self.1:__torch__.alexnet -> ignored
         // input.1:Tensor -> used
-        auto pt_type = in->type();
-        if (pt_type == type_lut["Tensor"]) {
+        if (in->type()->isSubtypeOf(c10::TensorType::get()) && ctx->evaluated_value_map.find(in) == ctx->evaluated_value_map.end()) {
             input_tensors.push_back(in);
         }
     }

core/conversion/conversion.h

@@ -6,14 +6,6 @@
 #include "torch/csrc/jit/ir/ir.h"
 #include "core/conversion/conversionctx/ConversionCtx.h"
 
-namespace torch {
-namespace jit {
-namespace script {
-const std::unordered_map<std::string, c10::TypePtr>& string_to_type_lut();
-}
-}
-}
-
 namespace trtorch {
 namespace core {
 namespace conversion {

core/conversion/conversion_blacklist.cpp

@@ -19,6 +19,7 @@ const std::unordered_set<std::string>& get_non_convertable_nodes() {
         "prim::device",
         "prim::GetAttr",
         "prim::CallMethod",
+        "prim::Drop",
         "aten:dropout",
     };
     return nonconvertable_nodes;

core/conversion/string_to_type_lut.cpp

@@ -22,7 +22,7 @@ std::vector<at::Tensor> RunCudaEngine(nvinfer1::IExecutionContext* ctx, std::pai
         auto dims = core::util::toDimsPad(inputs[i].sizes(), 1);
         auto shape = core::util::toVec(dims);
         contig_inputs.push_back(inputs[i].view(shape).contiguous());
-        LOG_DEBUG("In shape:" << shape);
+        LOG_DEBUG("In shape: " << shape);
         ctx->setBindingDimensions(i, dims);
         gpu_handles.push_back(contig_inputs.back().data_ptr());
     }
@@ -32,7 +32,7 @@ std::vector<at::Tensor> RunCudaEngine(nvinfer1::IExecutionContext* ctx, std::pai
     std::vector<at::Tensor> outputs;
     for (uint64_t o = inputs.size(); o < (io.first + io.second); o++) {
         auto out_shape = ctx->getBindingDimensions(o);
-        //LOG_DEBUG("Output: " << engine->getBindingName(o) << " out shape: " << out_shape);
+        LOG_DEBUG("Output shape: " << out_shape);
         auto dims = core::util::toVec(out_shape);
         auto type = util::toATenDType(ctx->getEngine().getBindingDataType(o));
         outputs.push_back(at::empty(dims, {at::kCUDA}).to(type).contiguous());