feat: Adding automatic casting to compare layers

narendasan · narendasan · commit 90af26eab94b · 2021-08-02T21:30:28.000-07:00
BERT converts but produces NaNs

Signed-off-by: Naren Dasan &lt;naren@narendasan.com&gt;
Signed-off-by: Naren Dasan &lt;narens@nvidia.com&gt;
diff --git a/core/conversion/converters/Weights.cpp b/core/conversion/converters/Weights.cpp
@@ -114,6 +114,7 @@ Weights::Weights(ConversionCtx* ctx, at::Tensor t) {
 // clang-format off
 std::ostream& operator<<(std::ostream& os, const Weights& w) {
   os << "Weights: " << w.shape
+     << "\n    Data Type: " << w.data.type
      << "\n    Number of input maps: " << w.num_input_maps
      << "\n    Number of output maps: " << w.num_output_maps
      << "\n    Element shape: [";
diff --git a/core/conversion/converters/impl/element_wise.cpp b/core/conversion/converters/impl/element_wise.cpp
@@ -25,6 +25,27 @@ nvinfer1::ITensor* clamp_util(
   return clamp_layer_out;
 }
 
+nvinfer1::ITensor* scalar_to_tensor(ConversionCtx* ctx, at::Scalar s) {
+  nvinfer1::ITensor* out;
+  if (s.isIntegral(false)) {
+    auto s_int = s.to<int64_t>();
+    auto s_t = torch::tensor({s_int}).to(at::kInt);
+    out = tensor_to_const(ctx, s_t);
+  } else if (s.isBoolean()) {
+    auto s_bool = s.to<bool>();
+    auto s_t = torch::tensor({s_bool}).to(at::kBool);
+    out = tensor_to_const(ctx, s_t);
+  } else if (s.isFloatingPoint()) {
+    auto other_float = s.to<float>();
+    auto s_t = torch::tensor({other_float});
+    out = tensor_to_const(ctx, s_t);
+  } else {
+    out = nullptr;
+    TRTORCH_THROW_ERROR("Unsupported data type for scalar. Found: (" << s.type() << ")");
+  }
+  return out;
+}
+
 auto element_wise_registrations TRTORCH_UNUSED =
     RegisterNodeConversionPatterns()
         .pattern({"aten::add.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> "
@@ -557,8 +578,10 @@ auto element_wise_registrations TRTORCH_UNUSED =
         .pattern({"aten::gt.Scalar(Tensor self, Scalar other) -> (Tensor)",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto self = args[0].ITensorOrFreeze(ctx);
-                    auto otherScalar = args[1].unwrapToScalar().to<float>();
-                    auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+                    auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
+                    if (self->getType() != other->getType()) {
+                      other = castITensor(ctx, other, self->getType());
+                    }
                     auto gt =
                         add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n));
                     TRTORCH_CHECK(gt, "Unable to create greater layer from node: " << *n);
@@ -584,8 +607,10 @@ auto element_wise_registrations TRTORCH_UNUSED =
         .pattern({"aten::lt.Scalar(Tensor self, Scalar other) -> (Tensor)",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto self = args[0].ITensorOrFreeze(ctx);
-                    auto otherScalar = args[1].unwrapToScalar().to<float>();
-                    auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+                    auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
+                    if (self->getType() != other->getType()) {
+                      other = castITensor(ctx, other, self->getType());
+                    }
                     auto lt =
                         add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n));
                     TRTORCH_CHECK(lt, "Unable to create less layer from node: " << *n);
@@ -613,6 +638,18 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     auto self = args[0].ITensorOrFreeze(ctx);
                     auto otherScalar = args[1].unwrapToScalar().to<float>();
                     auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+                    if (self->getType() == nvinfer1::DataType::kBOOL) {
+                      if (otherScalar == 0 || otherScalar == 1) {
+                        LOG_DEBUG("Since input tensor is type bool, casting input tensor and scalar to int32");
+                        other = castITensor(ctx, other, nvinfer1::DataType::kINT32);
+                        self = castITensor(ctx, self, nvinfer1::DataType::kINT32);
+                      } else {
+                        LOG_WARNING("Input Tensor has type bool, but scalar is not 0 or 1. Found: " << otherScalar);
+                      }
+                    }
+                    if (self->getType() != other->getType()) {
+                      other = castITensor(ctx, other, self->getType());
+                    }
                     auto eq =
                         add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n));
                     TRTORCH_CHECK(eq, "Unable to create equal layer from node: " << *n);
@@ -648,8 +685,10 @@ auto element_wise_registrations TRTORCH_UNUSED =
         .pattern({"aten::ge.Scalar(Tensor self, Scalar other) -> (Tensor)",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto self = args[0].ITensorOrFreeze(ctx);
-                    auto otherScalar = args[1].unwrapToScalar().to<float>();
-                    auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+                    auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
+                    if (self->getType() != other->getType()) {
+                      other = castITensor(ctx, other, self->getType());
+                    }
 
                     auto greater = add_elementwise(
                         ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater");
@@ -695,8 +734,10 @@ auto element_wise_registrations TRTORCH_UNUSED =
         .pattern({"aten::le.Scalar(Tensor self, Scalar other) -> (Tensor)",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto self = args[0].ITensorOrFreeze(ctx);
-                    auto otherScalar = args[1].unwrapToScalar().to<float>();
-                    auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+                    auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
+                    if (self->getType() != other->getType()) {
+                      other = castITensor(ctx, other, self->getType());
+                    }
 
                     auto less = add_elementwise(
                         ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less");
diff --git a/core/conversion/converters/impl/reduce.cpp b/core/conversion/converters/impl/reduce.cpp
@@ -144,8 +144,11 @@ auto reduce_registrations TRTORCH_UNUSED =
         .pattern({"aten::prod.dim_int(Tensor self, int dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto in_tensor = args[0].ITensorOrFreeze(ctx);
+                    auto in_dims = in_tensor->getDimensions();
                     auto dim = args[1].unwrapToInt();
-                    LOG_DEBUG("Dim to reduce: " << dim); // Some abuse of toDim but just for debug info
+                    LOG_DEBUG("Dim to reduce (original): " << dim);
+                    dim = dim < 0 ? (in_dims.nbDims + dim) : dim;
+                    LOG_DEBUG("Dim to reduce (converted): " << dim);
 
                     uint32_t axis_mask = 1 << dim;
                     LOG_DEBUG("Axis Mask: " << std::bitset<32>(axis_mask));
diff --git a/core/conversion/evaluators/aten.cpp b/core/conversion/evaluators/aten.cpp
@@ -187,9 +187,14 @@ auto aten_registrations TRTORCH_UNUSED =
                         if (tensor_var.isITensor()) {
                           auto tensor = tensor_var.ITensor();
                           return util::toVec(tensor->getDimensions());
-                        } else {
+                        } else if (tensor_var.IValue()->isTensor()) {
                           auto tensor = tensor_var.unwrapToTensor();
                           return tensor.sizes();
+                        } else if (tensor_var.IValue()->isCustomClass()) {
+                          auto tensor = tensor_var.IValue()->toCustomClass<TensorContainer>()->tensor();
+                          return util::toVec(tensor->getDimensions());
+                        } else {
+                          TRTORCH_THROW_ERROR("IValue is not some class of Tensor. Found: " << tensor_var.IValue()->type());
                         }
                       } else {
                         auto dim = args.at(n->input(1)).unwrapToInt();
@@ -201,13 +206,23 @@ auto aten_registrations TRTORCH_UNUSED =
                             dim += nbDims;
                           }
                           return dims[dim];
-                        } else {
+                        } else if (tensor_var.IValue()->isTensor()) {
                           auto tensor = tensor_var.unwrapToTensor();
                           auto nbDims = tensor.sizes().size();
                           if (dim < 0) {
                             dim += nbDims;
                           }
                           return tensor.sizes()[dim];
+                        } else if (tensor_var.IValue()->isCustomClass()) {
+                          auto tensor = tensor_var.IValue()->toCustomClass<TensorContainer>()->tensor();
+                          auto dims = util::toVec(tensor->getDimensions());
+                          auto nbDims = tensor->getDimensions().nbDims;
+                          if (dim < 0) {
+                            dim += nbDims;
+                          }
+                          return dims[dim];
+                        } else {
+                          TRTORCH_THROW_ERROR("IValue is not some class of Tensor. Found: " << tensor_var.IValue()->type());
                         }
                       }
                     },
diff --git a/core/conversion/var/Var.cpp b/core/conversion/var/Var.cpp
@@ -132,7 +132,8 @@ nvinfer1::ITensor* Var::ITensorOrFreeze(ConversionCtx* ctx) {
     out = ptr_.tensor;
   }
 
-  LOG_DEBUG("Frozen tensor shape: " << out->getDimensions());
+  LOG_DEBUG("ITensor shape: " << out->getDimensions());
+  LOG_DEBUG("ITensor type: " << out->getType());
   return out;
 }
 
diff --git a/cpp/trtorchc/main.cpp b/cpp/trtorchc/main.cpp
@@ -25,6 +25,34 @@
 #include "trtorch/ptq.h"
 #include "trtorch/trtorch.h"
 
+
+at::ScalarType to_torch_dtype(trtorch::CompileSpec::DataType dtype) {
+  switch (dtype) {
+    case trtorch::CompileSpec::DataType::kHalf:
+      return at::kHalf;
+    case trtorch::CompileSpec::DataType::kChar:
+      return at::kChar;
+    case trtorch::CompileSpec::DataType::kInt:
+      return at::kInt;
+    case trtorch::CompileSpec::DataType::kBool:
+      return at::kBool;
+    case trtorch::CompileSpec::DataType::kFloat:
+    default:
+      return at::kFloat;
+  }
+}
+
+const std::unordered_map<nvinfer1::DataType, at::ScalarType>& get_trt_at_type_map() {
+  static const std::unordered_map<nvinfer1::DataType, at::ScalarType> trt_at_type_map = {
+      {nvinfer1::DataType::kFLOAT, at::kFloat},
+      {nvinfer1::DataType::kHALF, at::kHalf},
+      {nvinfer1::DataType::kINT32, at::kInt},
+      {nvinfer1::DataType::kINT8, at::kChar},
+      {nvinfer1::DataType::kBOOL, at::kBool},
+  };
+  return trt_at_type_map;
+}
+
 bool checkRtol(const at::Tensor& diff, const std::vector<at::Tensor> inputs, float threshold) {
   double maxValue = 0.0;
   for (auto& tensor : inputs) {
@@ -238,6 +266,9 @@ int main(int argc, char** argv) {
       "Maximum acceptable numerical deviation from standard torchscript output (default 2e-5)",
       {'t', "threshold"});
 
+  args::Flag no_threshold_check(parser, "no-threshold-check", "Skip checking threshold compliance", {"no-threshold-check", "no-threshold-check"});
+  args::Flag truncate_long_and_double(parser, "truncate-long-double", "Truncate weights that are provided in 64bit to 32bit (Long, Double to Int, Float)", {"truncate", "truncate-long-double", "truncate-64bit"});
+
   args::Flag save_engine(
       parser,
       "save_engine",
@@ -481,6 +512,10 @@ int main(int argc, char** argv) {
     compile_settings.max_batch_size = args::get(max_batch_size);
   }
 
+  if (truncate_long_and_double) {
+    compile_settings.truncate_long_and_double = true;
+  }
+
   auto real_input_path = resolve_path(args::get(input_path));
   auto real_output_path = resolve_path(args::get(output_path));
 
@@ -507,9 +542,9 @@ int main(int argc, char** argv) {
   } else {
     auto trt_mod = trtorch::CompileGraph(mod, compile_settings);
 
-    if (compile_settings.enabled_precisions.size() == 1 &&
+    if (!no_threshold_check && (compile_settings.enabled_precisions.size() == 1 &&
         compile_settings.enabled_precisions.find(trtorch::CompileSpec::DataType::kFloat) !=
-            compile_settings.enabled_precisions.end()) {
+            compile_settings.enabled_precisions.end())) {
       double threshold_val = 2e-5;
       if (threshold) {
         threshold_val = args::get(threshold);
@@ -520,10 +555,12 @@ int main(int argc, char** argv) {
 
       for (auto i : ranges) {
         auto in = at::randn(i.opt_shape, {at::kCUDA});
+        in = in.to(to_torch_dtype(i.dtype));
         jit_inputs_ivalues.push_back(in.clone());
         trt_inputs_ivalues.push_back(in.clone());
       }
 
+      mod.to({at::kCUDA});
       torch::jit::IValue jit_results_ivalues = mod.forward(jit_inputs_ivalues);
       std::vector<at::Tensor> jit_results;
       if (jit_results_ivalues.isTensor()) {
@@ -557,9 +594,15 @@ int main(int argc, char** argv) {
         }
       }
     } else {
-      trtorch::logging::log(
+      if (no_threshold_check) {
+        trtorch::logging::log(
+          trtorch::logging::Level::kWARNING,
+          "Threshold check skipped, numerical precision is not checked");
+      } else {
+        trtorch::logging::log(
           trtorch::logging::Level::kWARNING,
           "Due to change in operating data type, numerical precision is not checked");
+      }
     }
 
     trt_mod.save(real_output_path);

Original file line number	Diff line number	Diff line change
`@@ -132,7 +132,8 @@ nvinfer1::ITensor* Var::ITensorOrFreeze(ConversionCtx* ctx) {`
`132`	`132`	`out = ptr_.tensor;`
`133`	`133`	`}`
`134`	`134`
`135`		`- LOG_DEBUG("Frozen tensor shape: " << out->getDimensions());`
	`135`	`+ LOG_DEBUG("ITensor shape: " << out->getDimensions());`
	`136`	`+ LOG_DEBUG("ITensor type: " << out->getType());`
`136`	`137`	`return out;`
`137`	`138`	`}`
`138`	`139`