feat: Add aten::full converter, quantization ops testcases

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

Author: peri044

core/conversion/conversionctx/ConversionCtx.cpp

@@ -71,8 +71,9 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
       }
       input_type = nvinfer1::DataType::kFLOAT;
       // Networks trained with Quantization aware training approach don't need a calibrator as they have Q/DQ nodes.
-      if (!settings.calibrator){
-        LOG_WARNING("Int8 precision has been enabled but no calibrator provided. This assumes the network has Q/DQ nodes obtained from Quantization aware training. For more details, refer to https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#work-with-qat-networks");
+      if (!settings.calibrator) {
+        LOG_WARNING(
+            "Int8 precision has been enabled but no calibrator provided. This assumes the network has Q/DQ nodes obtained from Quantization aware training. For more details, refer to https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#work-with-qat-networks");
       }
       break;
     case nvinfer1::DataType::kFLOAT:

core/conversion/converters/impl/constant.cpp

@@ -1,3 +1,4 @@
+#include <torch/torch.h>
 #include "core/conversion/converters/converters.h"
 #include "core/util/prelude.h"
 
@@ -25,6 +26,18 @@ auto constant_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
 
               LOG_DEBUG("Output tensor shape: " << const_out->getDimensions());
 
+              return true;
+            }})
+  .pattern({"aten::full(int[] size, Scalar fill_value, *, int? dtype=None, int? layout=None, Device? device=None, bool? pin_memory=None) -> (Tensor)",
+            [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+              auto size = args[0].unwrapToIntList();
+              auto scalar = args[1].unwrapToScalar().to<float>();
+              auto scalar_tensor = torch::full({5}, scalar);
+              auto full_tensor = tensor_to_const(ctx, scalar_tensor);
+              auto output = ctx->AssociateValueAndTensor(n->outputs()[0], full_tensor);
+
+              LOG_DEBUG("Output tensor shape: " << output->getDimensions());
+
               return true;
             }});
 // clang-format on

core/conversion/converters/impl/quantization.cpp

@@ -13,18 +13,18 @@ namespace {
 auto quantization_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
   .pattern({"aten::fake_quantize_per_tensor_affine(Tensor self, float scale, int zero_point, int quant_min, int quant_max) -> (Tensor)",
             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+              // This aten operator is generated from torch.fake_quantize_per_tensor_affine op in Pytorch python API.
+              // Example usage: https://github.com/pytorch/pytorch/blob/master/torch/quantization/fake_quantize.py#L145
               auto input = args[0].ITensorOrFreeze(ctx);
               auto scale = args[1].unwrapToScalar().to<float>();
               auto scaleTensor = tensor_to_const(ctx, torch::tensor({scale}));
-
               // Add and configure a QuantizeLayer.
               nvinfer1::IQuantizeLayer* quantize_layer = ctx->net->addQuantize(*input, *scaleTensor);
-              // Set an invalid axis
-              quantize_layer->setAxis(1);
+              quantize_layer->setAxis(0);
 
-              // Add and configure DequantizeLayer
+              // Add and configure DequantizeLayer following a QuantizeLayer
               nvinfer1::IDequantizeLayer* dequantize_layer = ctx->net->addDequantize(*quantize_layer->getOutput(0), *scaleTensor);
-              dequantize_layer->setAxis(1);
+              dequantize_layer->setAxis(0);
 
               auto qdq_out = ctx->AssociateValueAndTensor(n->outputs()[0], dequantize_layer->getOutput(0));
               LOG_DEBUG("[fake_quantize_per_tensor_affine] Output tensor shape: " << qdq_out->getDimensions());
@@ -33,17 +33,19 @@ auto quantization_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns(
             }})
   .pattern({"aten::fake_quantize_per_channel_affine(Tensor self, Tensor scale, Tensor zero_point, int axis, int quant_min, int quant_max) -> (Tensor)",
             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+              // This aten operator is generated from torch.fake_quantize_per_channel_affine op in Pytorch python API.
+              // Example usage: https://github.com/pytorch/pytorch/blob/master/torch/quantization/fake_quantize.py#L141
               auto input = args[0].ITensorOrFreeze(ctx);
               auto scale = args[1].ITensorOrFreeze(ctx);
-
+              int64_t axis = args[3].unwrapToScalar().to<int64_t>();
               // Add and configure a QuantizeLayer.
               nvinfer1::IQuantizeLayer* quantize_layer = ctx->net->addQuantize(*input, *scale);
-              // Set a channel axis=0 which represents output channels
-              quantize_layer->setAxis(0);
+              // Set a channel axis which represents output channels
+              quantize_layer->setAxis(axis);
 
               // Add and configure a DequantizeLayer.
               nvinfer1::IDequantizeLayer* dequantize_layer = ctx->net->addDequantize(*quantize_layer->getOutput(0), *scale);
-              dequantize_layer->setAxis(0);
+              dequantize_layer->setAxis(axis);
               auto qdq_out = ctx->AssociateValueAndTensor(n->outputs()[0], dequantize_layer->getOutput(0));
 
               LOG_DEBUG("[fake_quantize_per_channel_affine] Ouput tensor shape: " << qdq_out->getDimensions());

core/plugins/impl/interpolate_plugin.cpp

@@ -206,7 +206,6 @@ bool InterpolatePlugin::supportsFormatCombination(
     const nvinfer1::PluginTensorDesc* inOut,
     int nbInputs,
     int nbOutputs) noexcept {
-
   if (nbInputs != 1) {
     LOG_ERROR("Expected a single tensor as input to interpolate plugin");
   }

core/util/jit_util.h

@@ -13,7 +13,8 @@ inline std::string node_info(const torch::jit::Node* n) {
   std::stringstream ss;
   ss << *n;
   std::string node_info = ss.str();
-  // Nodes in torchscript graph have file name and line numbers commented for every node. Remove that when returning a node name for easier readability. 
+  // Nodes in torchscript graph have file name and line numbers commented for every node. Remove that when returning a
+  // node name for easier readability.
   node_info = node_info.substr(0, node_info.find("#", 0));
   node_info.erase(std::remove(node_info.begin(), node_info.end(), '\n'), node_info.end());
   return node_info;

cpp/trtorchexec/main.cpp

@@ -56,54 +56,59 @@ int main(int argc, const char* argv[]) {
   }
 
   auto compile_spec = trtorch::CompileSpec(dims);
+  // compile_spec.torch_fallback = trtorch::CompileSpec::TorchFallback(true);
   compile_spec.workspace_size = 1 << 24;
-
-  std::cout << "Checking operator support" << std::endl;
-  if (!trtorch::CheckMethodOperatorSupport(mod, "forward")) {
-    std::cerr << "Method is not currently supported by TRTorch" << std::endl;
-    return -1;
-  }
-
-  std::cout << "Compiling graph to save as TRT engine (/tmp/engine_converted_from_jit.trt)" << std::endl;
+  compile_spec.op_precision = torch::kChar;
+  // compile_spec.input_dtypes = {torch::kInt32, torch::kInt32};
+  // std::cout << "===Compile Spec: " << compile_spec << std::endl;
+  // compile_spec.torch_fallback = trtorch::CompileSpec::TorchFallback(true);
+  // compile_spec.torch_fallback.min_block_size = 1;
+  // std::cout << "Checking operator support" << std::endl;
+  // if (!trtorch::CheckMethodOperatorSupport(mod, "forward")) {
+  //   std::cerr << "Method is not currently supported by TRTorch" << std::endl;
+  //   return -1;
+  // }
+  //
+  // std::cout << "Compiling graph to save as TRT engine (/tmp/engine_converted_from_jit.trt)" << std::endl;
   auto engine = trtorch::ConvertGraphToTRTEngine(mod, "forward", compile_spec);
   std::ofstream out("/tmp/engine_converted_from_jit.trt");
   out << engine;
   out.close();
 
-  std::vector<torch::jit::IValue> jit_inputs_ivalues;
-  std::vector<torch::jit::IValue> trt_inputs_ivalues;
-  auto in = at::randint(5, dims[0], {at::kCUDA});
-  jit_inputs_ivalues.push_back(in.clone());
-  trt_inputs_ivalues.push_back(in.clone());
-
-  torch::jit::IValue jit_results_ivalues = mod.forward(jit_inputs_ivalues);
-  std::vector<at::Tensor> jit_results;
-  if (jit_results_ivalues.isTensor()) {
-    jit_results.push_back(jit_results_ivalues.toTensor());
-  } else {
-    auto results = jit_results_ivalues.toTuple()->elements();
-    for (auto r : results) {
-      jit_results.push_back(r.toTensor());
-    }
-  }
+  // std::vector<torch::jit::IValue> jit_inputs_ivalues;
+  // std::vector<torch::jit::IValue> trt_inputs_ivalues;
+  // auto in = at::randint(5, dims[0], {at::kCUDA});
+  // jit_inputs_ivalues.push_back(in.clone());
+  // trt_inputs_ivalues.push_back(in.clone());
+  // //
+  // torch::jit::IValue jit_results_ivalues = mod.forward(jit_inputs_ivalues);
+  // std::vector<at::Tensor> jit_results;
+  // if (jit_results_ivalues.isTensor()) {
+  //   jit_results.push_back(jit_results_ivalues.toTensor());
+  // } else {
+  //   auto results = jit_results_ivalues.toTuple()->elements();
+  //   for (auto r : results) {
+  //     jit_results.push_back(r.toTensor());
+  //   }
+  // }
 
   std::cout << "Compiling graph as module" << std::endl;
   auto trt_mod = trtorch::CompileGraph(mod, compile_spec);
-  std::cout << "Running TRT module" << std::endl;
-  torch::jit::IValue trt_results_ivalues = trt_mod.forward(trt_inputs_ivalues);
-  std::vector<at::Tensor> trt_results;
-  if (trt_results_ivalues.isTensor()) {
-    trt_results.push_back(trt_results_ivalues.toTensor());
-  } else {
-    auto results = trt_results_ivalues.toTuple()->elements();
-    for (auto r : results) {
-      trt_results.push_back(r.toTensor());
-    }
-  }
-
-  for (size_t i = 0; i < trt_results.size(); i++) {
-    almostEqual(jit_results[i], trt_results[i].reshape_as(jit_results[i]));
-  }
+  // std::cout << "Running TRT module" << std::endl;
+  // torch::jit::IValue trt_results_ivalues = trt_mod.forward(trt_inputs_ivalues);
+  // std::vector<at::Tensor> trt_results;
+  // if (trt_results_ivalues.isTensor()) {
+  //   trt_results.push_back(trt_results_ivalues.toTensor());
+  // } else {
+  //   auto results = trt_results_ivalues.toTuple()->elements();
+  //   for (auto r : results) {
+  //     trt_results.push_back(r.toTensor());
+  //   }
+  // }
+  //
+  // for (size_t i = 0; i < trt_results.size(); i++) {
+  //   almostEqual(jit_results[i], trt_results[i].reshape_as(jit_results[i]));
+  // }
 
   std::cout << "Converted Engine saved to /tmp/engine_converted_from_jit.trt" << std::endl;
 

tests/core/conversion/converters/BUILD

@@ -59,6 +59,10 @@ converter_test(
     name = "test_pooling",
 )
 
+converter_test(
+    name = "test_quantization",
+)
+
 converter_test(
     name = "test_reduce",
 )

tests/core/conversion/converters/test_quantization.cpp

@@ -0,0 +1,63 @@
+#include <string>
+#include "NvInfer.h"
+#include "core/compiler.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+
+TEST(Converters, ATenFakeQuantizePerTensorConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%x.1 : Tensor):
+      %7 : int = prim::Constant[value=-128]()
+      %3 : float = prim::Constant[value=6.]()
+      %4 : int = prim::Constant[value=0]()
+      %8 : int = prim::Constant[value=127]()
+      %quant_input.1 : Tensor = aten::fake_quantize_per_tensor_affine(%x.1, %3, %4, %7, %8)
+      return (%quant_input.1))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::randint(1, 10, {1, 5, 5, 5}, {at::kCUDA}).to(at::kFloat);
+
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in}, nvinfer1::DataType::kINT8);
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}
+
+TEST(Converters, ATenFakeQuantizePerChannelConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%x.1 : Tensor):
+        %22 : int = prim::Constant[value=-128]()
+        %14 : int = prim::Constant[value=4]()
+        %9 : None = prim::Constant()
+        %35 : Device = prim::Constant[value="cuda:0"]()
+        %6 : int = prim::Constant[value=6]()
+        %3 : int = prim::Constant[value=5]()
+        %5 : float = prim::Constant[value=3.5]()
+        %13 : int = prim::Constant[value=1]()
+        %23 : int = prim::Constant[value=127]()
+        %4 : int[] = prim::ListConstruct(%3)
+        %11 : Tensor = aten::full(%4, %5, %6, %9, %35, %9)
+        %12 : int[] = prim::ListConstruct(%3)
+        %19 : Tensor = aten::full(%12, %13, %14, %9, %35, %9)
+        %quant_input.1 : Tensor = aten::fake_quantize_per_channel_affine(%x.1, %11, %19, %13, %22, %23)
+        return (%quant_input.1))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::randint(1, 10, {1, 5, 3, 3}, {at::kCUDA});
+
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in}, nvinfer1::DataType::kINT8);
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}

tests/util/run_graph_engine.cpp

@@ -63,11 +63,13 @@ std::vector<at::Tensor> RunEngine(std::string& eng, std::vector<at::Tensor> inpu
 std::vector<at::Tensor> RunGraphEngine(
     std::shared_ptr<torch::jit::Graph>& g,
     core::conversion::GraphParams& named_params,
-    std::vector<at::Tensor> inputs) {
+    std::vector<at::Tensor> inputs,
+    nvinfer1::DataType op_precision = nvinfer1::DataType::kFLOAT) {
   LOG_DEBUG("Running TRT version");
   auto in = toInputRanges(inputs);
   auto info = core::conversion::ConversionInfo(in);
   info.engine_settings.workspace_size = 1 << 20;
+  info.engine_settings.op_precision = op_precision;
   std::string eng = core::conversion::ConvertBlockToEngine(g->block(), info, named_params);
   return RunEngine(eng, inputs);
 }

tests/util/util.h

@@ -28,7 +28,8 @@ std::vector<at::Tensor> RunGraph(
 std::vector<at::Tensor> RunGraphEngine(
     std::shared_ptr<torch::jit::Graph>& g,
     core::conversion::GraphParams& named_params,
-    std::vector<at::Tensor> inputs);
+    std::vector<at::Tensor> inputs,
+    nvinfer1::DataType dtype = nvinfer1::DataType::kFLOAT);
 
 // Runs an arbitrary JIT graph with dynamic input sizes by converting it to
 // TensorRT and running inference and returns results