Merge pull request #307 from NVIDIA/new_evaluators

Adding aten::is_floating_point evaluator

Author: narendasan

core/conversion/evaluators/aten.cpp

@@ -585,6 +585,22 @@ auto aten_registrations TRTORCH_UNUSED =
                       return {};
                     },
                     EvalOptions()})
+        .evaluator({c10::Symbol::fromQualString("aten::is_floating_point"),
+                    [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+                      auto tensor_var = args.at(n->input(0));
+                      if (tensor_var.isITensor()) {
+                        auto tensor = tensor_var.ITensor();
+                        auto t = tensor->getType();
+                        return (t == nvinfer1::DataType::kFLOAT || t == nvinfer1::DataType::kHALF);
+                      } else {
+                        auto tensor = tensor_var.unwrapToTensor();
+                        auto t = tensor.scalar_type();
+                        return at::isFloatingType(t);
+                      }
+                    },
+                    EvalOptions().validSchemas({
+                        "aten::is_floating_point(Tensor self) -> (bool)",
+                    })})
         .evaluator(
             {c10::Symbol::fromQualString("aten::tensor"),
              [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {

tests/accuracy/datasets/cifar10.cpp

@@ -51,7 +51,8 @@ std::pair<torch::Tensor, torch::Tensor> read_batch(const std::string& path) {
     labels.push_back(label);
     auto image_tensor =
         torch::from_blob(image.data(), {kImageChannels, kImageDim, kImageDim}, torch::TensorOptions().dtype(torch::kU8))
-            .to(torch::kF32).div(255);
+            .to(torch::kF32)
+            .div(255);
     images.push_back(image_tensor);
   }
 

tests/core/conversion/evaluators/test_aten_evaluators.cpp

@@ -471,3 +471,39 @@ TEST(Evaluators, IntFloatEvaluatesCorrectly) {
 
   ASSERT_TRUE(jit_results[0] == trt_results[0]);
 }
+
+TEST(Evaluators, ATenIsFloatingPointEvaluatesTrueCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : bool = aten::is_floating_point(%0)
+        return (%1))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {1, 3, 3, 3}, {at::kCUDA}).to(torch::kF32);
+  auto in_trt = in.clone();
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {in});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {in_trt});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}
+
+TEST(Evaluators, ATenIsFloatingPointEvaluatesFalseCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : bool = aten::is_floating_point(%0)
+        return (%1))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {1, 3, 3, 3}, {at::kCUDA}).to(torch::kI8);
+  auto in_trt = in.clone();
+
+  auto jit_results = trtorch::tests::util::EvaluateGraphJIT(g, {in});
+  auto trt_results = trtorch::tests::util::EvaluateGraph(g->block(), {in_trt});
+
+  ASSERT_TRUE(jit_results[0] == trt_results[0]);
+}