Merge pull request #565 from NVIDIA/tensorrt_8_update

TensorRT 8.0 update (rebased for 1.9 master)

Author: narendasan

README.md

@@ -81,7 +81,7 @@ These are the following dependencies used to verify the testcases. TRTorch can w
 - Libtorch 1.9.0 (built with CUDA 11.1)
 - CUDA 11.1 (10.2 on Jetson)
 - cuDNN 8.1
-- TensorRT 7.2.3
+- TensorRT 8.0.1.6
 
 ## Prebuilt Binaries and Wheel files
 

WORKSPACE

@@ -70,20 +70,20 @@ http_archive(
 http_archive(
     name = "cudnn",
     build_file = "@//third_party/cudnn/archive:BUILD",
-    sha256 = "98a8784e92862f20018d20c281b30d4a0cd951f93694f6433ccf4ae9c502ba6a",
+    sha256 = "39412acd9ef5dd27954b6b9f5df75bd381c5d7ceb7979af6c743a7f4521f9c77",
     strip_prefix = "cuda",
     urls = [
-        "https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.1.1.33/11.2_20210301/cudnn-11.2-linux-x64-v8.1.1.33.tgz",
+        "https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.2.1.32/11.3_06072021/cudnn-11.3-linux-x64-v8.2.1.32.tgz",
     ],
 )
 
 http_archive(
     name = "tensorrt",
     build_file = "@//third_party/tensorrt/archive:BUILD",
-    sha256 = "d3a1f478e304b48878604fac70ce7920fece71f9cac62f925c9c59c197f5d087",
-    strip_prefix = "TensorRT-7.2.3.4",
+    sha256 = "def6a5ee50bed25a68a9c9e22ec671a8f29ee5414bde47c5767bd279e5596f88",
+    strip_prefix = "TensorRT-8.0.1.6",
     urls = [
-        "https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/7.2.3/tars/TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz",
+        "https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/8.0.1/tars/tensorrt-8.0.1.6.linux.x86_64-gnu.cuda-11.3.cudnn8.2.tar.gz",
     ],
 )
 

core/conversion/conversionctx/ConversionCtx.cpp

@@ -69,10 +69,10 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
       case nvinfer1::DataType::kINT8:
         TRTORCH_CHECK(builder->platformHasFastInt8(), "Requested inference in INT8 but platform does not support INT8");
         cfg->setFlag(nvinfer1::BuilderFlag::kINT8);
-        TRTORCH_CHECK(
-            settings.calibrator != nullptr,
-            "Requested inference in INT8 but no calibrator provided, set the ptq_calibrator field in the CompileSpec struct with your calibrator");
-        cfg->setInt8Calibrator(settings.calibrator);
+        if (settings.calibrator == nullptr) {
+          LOG_INFO(
+              "INT8 kernels are enabled but not calibrator was provided, assuming source model was trained quantization aware");
+        }
         break;
       case nvinfer1::DataType::kFLOAT:
         break;
@@ -90,6 +90,10 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
     cfg->clearFlag(nvinfer1::BuilderFlag::kTF32);
   }
 
+  if (settings.sparse_weights) {
+    cfg->setFlag(nvinfer1::BuilderFlag::kSPARSE_WEIGHTS);
+  }
+
   if (settings.refit) {
     cfg->setFlag(nvinfer1::BuilderFlag::kREFIT);
   }
@@ -130,9 +134,9 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
 }
 
 ConversionCtx::~ConversionCtx() {
-  builder->destroy();
-  net->destroy();
-  cfg->destroy();
+  delete builder;
+  delete net;
+  delete cfg;
   for (auto ptr : builder_resources) {
     free(ptr);
   }
@@ -150,14 +154,11 @@ torch::jit::IValue* ConversionCtx::AssociateValueAndIValue(const torch::jit::Val
 }
 
 std::string ConversionCtx::SerializeEngine() {
-  auto engine = builder->buildEngineWithConfig(*net, *cfg);
-  if (!engine) {
-    TRTORCH_THROW_ERROR("Building TensorRT engine failed");
+  auto serialized_network = builder->buildSerializedNetwork(*net, *cfg);
+  if (!serialized_network) {
+    TRTORCH_THROW_ERROR("Building serialized network failed in TensorRT");
   }
-  auto serialized_engine = engine->serialize();
-  engine->destroy();
-  auto engine_str = std::string((const char*)serialized_engine->data(), serialized_engine->size());
-  serialized_engine->destroy();
+  auto engine_str = std::string((const char*)serialized_network->data(), serialized_network->size());
   return engine_str;
 }
 

core/conversion/conversionctx/ConversionCtx.h

@@ -25,14 +25,14 @@ struct Device {
 
 struct BuilderSettings {
   std::set<nvinfer1::DataType> enabled_precisions = {nvinfer1::DataType::kFLOAT};
-  std::vector<nvinfer1::DataType> input_dtypes;
+  bool sparse_weights = false;
   bool disable_tf32 = false;
   bool refit = false;
   bool debug = false;
   bool strict_types = false;
   bool truncate_long_and_double = false;
   Device device;
-  nvinfer1::EngineCapability capability = nvinfer1::EngineCapability::kDEFAULT;
+  nvinfer1::EngineCapability capability = nvinfer1::EngineCapability::kSTANDARD;
   nvinfer1::IInt8Calibrator* calibrator = nullptr;
   uint64_t num_min_timing_iters = 2;
   uint64_t num_avg_timing_iters = 1;

core/conversion/converters/impl/batch_norm.cpp

@@ -40,6 +40,7 @@ auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
       LOG_DEBUG("momentum disregarded");
       LOG_DEBUG("training disregarded");
       LOG_DEBUG("cudnn disregarded");
+      TRTORCH_CHECK(orig_shape.nbDims > 2, "Unable to create batch normalization layer from node: " << *n);
 
       // Expand spatial dims from 1D to 2D if needed
       bool expandDims = (orig_shape.nbDims < 4);

core/conversion/converters/impl/conv_deconv.cpp

@@ -30,7 +30,8 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
   LOG_DEBUG("out_padding: " << out_padding);
   LOG_DEBUG("groups: " << groups);
 
-  // Expand spatial dims from 1D to 2D if needed
+  TRTORCH_CHECK(orig_dims.nbDims > 2, "Unable to create convolution layer from node: " << *n);
+
   bool expandDims = (orig_dims.nbDims < 4);
   if (expandDims) {
     in = addPadding(ctx, n, in, 4);

core/conversion/converters/impl/interpolate.cpp

@@ -109,17 +109,9 @@ void resize_layer_size(
   resize_layer->setResizeMode(mode);
   resize_layer->setName(util::node_info(n).c_str());
 
-  // if interpolation mode is linear, align corners must have been set to true.
-  // else, don't use align corners.
-  if (mode == nvinfer1::ResizeMode::kLINEAR) {
-#if NV_TENSORRT_MAJOR < 7 || (NV_TENSORRT_MAJOR == 7 && NV_TENSORRT_MINOR < 1) // IF TRT VERSION <= 7.0
-    TRTORCH_CHECK(align_corners, "resize layer (linear) only supports align_corners=True in TensorRT <= 7.0");
-    resize_layer->setAlignCorners(true);
-#else
-    resize_layer->setAlignCorners(align_corners);
-#endif
+  if (align_corners) {
+    resize_layer->setCoordinateTransformation(nvinfer1::ResizeCoordinateTransformation::kALIGN_CORNERS);
   }
-
   auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
 
   LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());

core/conversion/converters/impl/pooling.cpp

@@ -48,7 +48,7 @@ bool AdaptivePoolingConverter(
 
   auto orig_dims = in->getDimensions();
   bool expandDims = (orig_dims.nbDims < 4);
-
+  TRTORCH_CHECK(orig_dims.nbDims > 2, "Unable to create pooling layer from node: " << *n);
   if (expandDims) {
     in = addPadding(ctx, n, in, 4, false, false);
   }
@@ -122,6 +122,7 @@ bool PoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& args,
 
   // Max Pool needs at least 4D input
   auto orig_dims = in->getDimensions();
+  TRTORCH_CHECK(orig_dims.nbDims > 2, "Unable to create pooling layer from node: " << *n);
   bool expandDims = (orig_dims.nbDims < 4);
 
   if (expandDims) {

core/plugins/impl/interpolate_plugin.cpp

@@ -105,35 +105,35 @@ std::vector<int64_t> InterpolatePlugin::getOutputSize() {
   return size_;
 }
 
-int InterpolatePlugin::getNbOutputs() const {
+int InterpolatePlugin::getNbOutputs() const noexcept {
   if (mode_ == "adaptive_max_pool2d") {
     return 2;
   } else {
     return 1;
   }
 }
 
-const char* InterpolatePlugin::getPluginType() const {
+const char* InterpolatePlugin::getPluginType() const noexcept {
   return "Interpolate";
 }
 
-const char* InterpolatePlugin::getPluginVersion() const {
+const char* InterpolatePlugin::getPluginVersion() const noexcept {
   return "1";
 }
 
-const char* InterpolatePlugin::getPluginNamespace() const {
+const char* InterpolatePlugin::getPluginNamespace() const noexcept {
   return "trtorch";
 }
 
-nvinfer1::IPluginV2DynamicExt* InterpolatePlugin::clone() const {
+nvinfer1::IPluginV2DynamicExt* InterpolatePlugin::clone() const noexcept {
   return new InterpolatePlugin(in_shape_, out_shape_, size_, scales_, mode_, align_corners_, use_scales_);
 }
 
 nvinfer1::DimsExprs InterpolatePlugin::getOutputDimensions(
     int outputIndex,
     const nvinfer1::DimsExprs* inputs,
     int nbInputs,
-    nvinfer1::IExprBuilder& exprBuilder) {
+    nvinfer1::IExprBuilder& exprBuilder) noexcept {
   nvinfer1::DimsExprs output(inputs[0]);
 
   // TODO: This should enable the case of using this plugin with dynamic shape, scale factor and align corners == true
@@ -165,15 +165,15 @@ nvinfer1::DimsExprs InterpolatePlugin::getOutputDimensions(
 }
 
 nvinfer1::DataType InterpolatePlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs)
-    const {
+    const noexcept {
   return nvinfer1::DataType::kFLOAT;
 }
 
-int InterpolatePlugin::initialize() {
+int InterpolatePlugin::initialize() noexcept {
   return 0;
 }
 
-void InterpolatePlugin::serialize(void* buffer) const {
+void InterpolatePlugin::serialize(void* buffer) const noexcept {
   std::string data = serializeToString();
   size_t size = getSerializationSize();
 
@@ -197,23 +197,32 @@ std::string InterpolatePlugin::serializeToString() const {
   return data_str.str();
 }
 
-size_t InterpolatePlugin::getSerializationSize() const {
+size_t InterpolatePlugin::getSerializationSize() const noexcept {
   return serializeToString().size();
 }
 
 bool InterpolatePlugin::supportsFormatCombination(
     int pos,
     const nvinfer1::PluginTensorDesc* inOut,
     int nbInputs,
-    int nbOutputs) {
-  TRTORCH_ASSERT(nbInputs == 1, "Expected a single tensor as input to interpolate plugin");
-
+    int nbOutputs) noexcept {
+  if (nbInputs != 1) {
+    LOG_ERROR("Expected a single tensor as input to interpolate plugin");
+  }
   if (mode_ == "adaptive_max_pool2d") {
-    TRTORCH_ASSERT(nbOutputs == 2, "Expected 2 tensors as output to interpolate plugin");
-    TRTORCH_ASSERT(0 <= pos && pos <= 2, "There should be exactly 3 connections to the plugin - 1 input, 2 output");
+    if (nbOutputs != 2) {
+      LOG_ERROR("Expected 2 tensors as output to interpolate plugin");
+    }
+    if (pos < 0 || pos > 2) {
+      LOG_ERROR("There should be exactly 3 connections to the plugin - 1 input, 2 output");
+    }
   } else {
-    TRTORCH_ASSERT(nbOutputs == 1, "Expected a single tensor as output to interpolate plugin");
-    TRTORCH_ASSERT(0 <= pos && pos <= 1, "There should be exactly 2 connections to the plugin - 1 input, 1 output");
+    if (nbOutputs != 1) {
+      LOG_ERROR("Expected a single tensor as output to interpolate plugin");
+    }
+    if (pos < 0 || pos > 1) {
+      LOG_ERROR("There should be exactly 2 connections to the plugin - 1 input, 1 output");
+    }
   }
 
   const nvinfer1::PluginTensorDesc& in = inOut[0];
@@ -232,15 +241,15 @@ void InterpolatePlugin::configurePlugin(
     const nvinfer1::DynamicPluginTensorDesc* in,
     int nbInputs,
     const nvinfer1::DynamicPluginTensorDesc* out,
-    int nbOutputs) {
+    int nbOutputs) noexcept {
   dtype_ = nvinfer1::DataType::kFLOAT;
 }
 
 size_t InterpolatePlugin::getWorkspaceSize(
     const nvinfer1::PluginTensorDesc* inputs,
     int nbInputs,
     const nvinfer1::PluginTensorDesc* outputs,
-    int nbOutputs) const {
+    int nbOutputs) const noexcept {
   return 0;
 }
 
@@ -250,7 +259,7 @@ int InterpolatePlugin::enqueue(
     const void* const* inputs,
     void* const* outputs,
     void* workspace,
-    cudaStream_t stream) {
+    cudaStream_t stream) noexcept {
   at::Tensor input =
       at::from_blob((void*)inputs[0], util::toVec(inputDesc->dims), [](void*) {}, {at::kCUDA}).to(torch::kFloat);
   at::Tensor output =
@@ -317,21 +326,21 @@ InterpolatePluginCreator::InterpolatePluginCreator() {
   mFC.fields = mPluginAttributes.data();
 }
 
-const char* InterpolatePluginCreator::getPluginNamespace() const {
+const char* InterpolatePluginCreator::getPluginNamespace() const noexcept {
   return "trtorch";
 }
 
-const char* InterpolatePluginCreator::getPluginName() const {
+const char* InterpolatePluginCreator::getPluginName() const noexcept {
   return "Interpolate";
 }
 
-const char* InterpolatePluginCreator::getPluginVersion() const {
+const char* InterpolatePluginCreator::getPluginVersion() const noexcept {
   return "1";
 }
 
 nvinfer1::IPluginV2* InterpolatePluginCreator::createPlugin(
     const char* name,
-    const nvinfer1::PluginFieldCollection* fc) {
+    const nvinfer1::PluginFieldCollection* fc) noexcept {
   std::vector<int64_t> in_shape;
   std::vector<int64_t> out_shape;
   std::vector<int64_t> out_size;
@@ -370,12 +379,12 @@ nvinfer1::IPluginV2* InterpolatePluginCreator::createPlugin(
 nvinfer1::IPluginV2* InterpolatePluginCreator::deserializePlugin(
     const char* name,
     const void* serialData,
-    size_t serialLength) {
+    size_t serialLength) noexcept {
   name_ = name;
   return new InterpolatePlugin((const char*)serialData, serialLength);
 }
 
-const nvinfer1::PluginFieldCollection* InterpolatePluginCreator::getFieldNames() {
+const nvinfer1::PluginFieldCollection* InterpolatePluginCreator::getFieldNames() noexcept {
   return nullptr;
 }
 
@@ -384,4 +393,4 @@ REGISTER_TRTORCH_PLUGIN(InterpolatePluginCreator);
 } // namespace impl
 } // namespace plugins
 } // namespace core
-} // namespace trtorch
+} // namespace trtorch