Merge pull request #51 from NVIDIA/docs

Adds a documenation site for TRTorch

Author: narendasan

.gitignore

@@ -24,4 +24,7 @@ cpp/ptq/datasets/data/
 tests/accuracy/datasets/data/*
 ._.DS_Store
 *.tar.gz
-*.tgz
+*.tgz
+docsrc/_build
+docsrc/_api
+docsrc/_tmp

core/lowering/lowering.cpp

@@ -25,6 +25,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g) {
     torch::jit::FuseLinear(g);
     passes::RemoveDropout(g);
     passes::FuseFlattenLinear(g);
+    passes::Conv2DToConvolution(g);
     passes::UnpackAddMM(g);
     passes::UnpackLogSoftmax(g);
     //passes::RemoveDimExeception(g);

core/lowering/passes/BUILD

@@ -6,6 +6,7 @@ cc_library(
         "passes.h",
     ],
     srcs = [
+        "conv2d_to_convolution.cpp",
         "exception_elimination.cpp",
         "fuse_flatten_linear.cpp",
         "remove_dropout.cpp",

core/lowering/passes/conv2d_to_convolution.cpp

@@ -0,0 +1,34 @@
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace lowering {
+namespace passes {
+
+void Conv2DToConvolution(std::shared_ptr<torch::jit::Graph>& graph) {
+    std::string conv2d_pattern = R"IR(
+        graph(%x, %w, %b, %s, %p, %d, %g):
+            %4 : Tensor = aten::conv2d(%x, %w, %b, %s, %p, %d, %g)
+            return (%4))IR";
+    std::string convolution_pattern = R"IR(
+        graph(%x, %w, %b, %s, %p, %d, %g):
+            %1 : bool = prim::Constant[value=1]()
+            %2 : int[] = prim::Constant[value=[0, 0]]()
+            %3 : bool = prim::Constant[value=0]()
+            %4 : Tensor = aten::_convolution(%x, %w, %b, %s, %p, %d, %1, %2, %g, %1, %1, %3)
+            return (%4))IR";;
+
+    // replace matmul + add pattern to linear
+    torch::jit::SubgraphRewriter map_conv2d_to_convolution;
+    map_conv2d_to_convolution.RegisterRewritePattern(
+        conv2d_pattern, convolution_pattern);
+    map_conv2d_to_convolution.runOnGraph(graph);
+    LOG_GRAPH("Post map conv2d -> _convolution: " << *graph);
+}
+
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace trtorch

core/lowering/passes/passes.h

@@ -7,6 +7,7 @@ namespace core {
 namespace lowering {
 namespace passes {
 
+void Conv2DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void FuseFlattenLinear(std::shared_ptr<torch::jit::Graph>& graph);
 void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph);
 void UnpackAddMM(std::shared_ptr<torch::jit::Graph>& graph);

cpp/api/include/trtorch/logging.h

@@ -9,12 +9,18 @@ namespace logging {
  * Emum for setting message severity
  */
 enum Level {
-    kINTERNAL_ERROR, // Only print messages for internal errors
-    kERROR,          // Print all internal errors and errors (default)
-    kWARNING,        // Print warnings and errors
-    kINFO,           // Print all info, warnings and errors
-    kDEBUG,          // Print all debug info, info, warnings and errors
-    kGRAPH,          // Print everything including the intermediate graphs of the lowering phase
+    /// Only print messages for internal errors
+    kINTERNAL_ERROR,
+    /// Print all internal errors and errors (default)
+    kERROR,
+    /// Print warnings and errors
+    kWARNING,
+    /// Print all info, warnings and errors
+    kINFO,
+    /// Print all debug info, info, warnings and errors
+    kDEBUG,
+    /// Print everything including the intermediate graphs of the lowering phase
+    kGRAPH,
 };
 
 // Are these ones necessary for the user?
@@ -37,11 +43,15 @@ TRTORCH_API void set_is_colored_output_on(bool colored_output_on);
 
 /**
  * @brief Get the current reportable log level
+ *
+ * @return TRTORCH_API get_reportable_log_level
  */
 TRTORCH_API Level get_reportable_log_level();
 
 /**
  * @brief Is colored output enabled?
+ *
+ * @return TRTORCH_API get_is_colored_output_on
  */
 TRTORCH_API bool get_is_colored_output_on();
 

cpp/api/include/trtorch/ptq.h

@@ -6,6 +6,7 @@
 #include <iostream>
 #include <sstream>
 
+#ifndef DOXYGEN_SHOULD_SKIP_THIS
 namespace nvinfer1 {
 class IInt8Calibrator;
 class IInt8EntropyCalibrator2;
@@ -17,18 +18,40 @@ template<typename Example>
 class Iterator;
 }
 }
+#endif //DOXYGEN_SHOULD_SKIP_THIS
 
 namespace trtorch {
 namespace ptq {
 
+/**
+ * @brief Generic Int8Calibrator implementation based on a specified
+ * TensorRT calibration algorithm and a LibTorch DataLoader
+ *
+ * @tparam Algorithm: class nvinfer1::IInt8Calibrator (Default: nvinfer1::IInt8EntropyCalibrator2) - Algorithm to use
+ * @tparam DataLoaderUniquePtr: std::unique_ptr<torch::data::DataLoader> - DataLoader type
+ */
 template<typename Algorithm, typename DataLoaderUniquePtr>
 class Int8Calibrator : Algorithm {
     using DataLoader = typename DataLoaderUniquePtr::element_type;
     using Batch = typename DataLoader::super::BatchType;
 public:
+    /**
+     * @brief Construct a new Int8Calibrator object
+     *
+     * Using the provided DataLoader, construct a calibrator that can be used for PTQ with TRTorch
+     *
+     * @param dataloader: std::unqiue_ptr<torch::data::DataLoader> - A unique pointer to the DataLoader, should be what is returned from the make_data_loader factory
+     * @param cache_file_path: const std::string& - A path to store / find the calibration cache
+     * @param use_cache : bool - Whether to use the cache (if it exists)
+     */
     Int8Calibrator(DataLoaderUniquePtr dataloader, const std::string& cache_file_path, bool use_cache)
       : dataloader_(dataloader.get()), it_(dataloader_->end()), cache_file_path_(cache_file_path), use_cache_(use_cache) {}
 
+    /**
+     * @brief Get the Batch Size for the next batch (always 1 due to issues with TRT and explicit batch)
+     *
+     * @return int
+     */
     int getBatchSize() const override {
         // HACK: TRTorch only uses explict batch sizing, INT8 Calibrator does not
         // work when reporting the batch size here and having explicity batching.
@@ -37,6 +60,15 @@ class Int8Calibrator : Algorithm {
         //return static_cast<int>(dataloader_->options().batch_size);
     }
 
+    /**
+     * @brief Get the next Batch
+     *
+     * @param bindings: void*[] - An array of binding pointers (fed in from TensorRT calibrator), these buffers should be filed with batch data for each input
+     * @param names: const char*[] - Names of bindings
+     * @param nbBindings: int - Number of bindings
+     * @return true - There is a new batch for the calibrator to consume
+     * @return false - There is not a new batch for the calibrator to consume
+     */
     bool getBatch(void* bindings[], const char* names[], int nbBindings) override {
         // HACK: doesnt seem like the first try in the initializer list works
         if (! it_created_) {
@@ -60,6 +92,14 @@ class Int8Calibrator : Algorithm {
         return true;
     }
 
+    /**
+     * @brief Read calibration cache
+     *
+     * How to read from the calibration cache, only enabled if use_cache is set
+     *
+     * @param length
+     * @return const void* - Pointer to cache data
+     */
     const void* readCalibrationCache(size_t& length) override {
         if (use_cache_) {
             std::stringstream ss;
@@ -81,6 +121,14 @@ class Int8Calibrator : Algorithm {
         return nullptr;
     }
 
+    /**
+     * @brief Write calibration cache
+     *
+     * Write a the calibration cache provided by TensorRT to a specified file
+     *
+     * @param cache: const void* - cache data
+     * @param length: size_t - length of cache
+     */
     void writeCalibrationCache(const void* cache, size_t length) override {
         std::ofstream cache_file(cache_file_path_, std::ios::binary);
         cache_file.write(reinterpret_cast<const char*>(cache), length);
@@ -89,37 +137,87 @@ class Int8Calibrator : Algorithm {
         logging::log(logging::Level::kINFO, ss.str());
     }
 
+    /**
+     * @brief operator to cast to nvinfer1::IInt8Calibrator*
+     *
+     * Convience function to convert to a IInt8Calibrator* to easily be assigned to the ptq_calibrator field in ExtraInfo
+     *
+     * @return nvinfer1::IInt8Calibrator*
+     */
     operator nvinfer1::IInt8Calibrator* () {
         return reinterpret_cast<nvinfer1::IInt8Calibrator*>(this);
     }
 
 private:
+    /// Pointer to the dataloader
     DataLoader* dataloader_;
+    /// Iterator used to traverse the dataloader
     torch::data::Iterator<Batch> it_;
+    /// Path to cache file
     const std::string& cache_file_path_;
+    /// Size of cache
     size_t cache_size_ = 0;
+    /// Whether to use the cache or not
     bool use_cache_;
+    /// Cache data
     std::vector<char> cache_;
+    /// If the iterator has been created, DataLoaders can only have 1 live iterator,
+    /// due to some issues this cannot be created at construction, so it is set in the first
+    /// batch, controlled by this flag
     bool it_created_ = false;
 };
 
+/**
+ * @brief Generic Int8Calibrator implementation based on a specified
+ * TensorRT calibration algorithm that only reads from a calibration file
+ *
+ * @tparam Algorithm: class nvinfer1::IInt8Calibrator (Default: nvinfer1::IInt8EntropyCalibrator2) - Algorithm to use
+ */
 template<typename Algorithm>
 class Int8CacheCalibrator : Algorithm {
 public:
+    /**
+     * @brief Construct a new Int 8 Cache Calibrator object
+     *
+     * @param cache_file_path
+     */
     Int8CacheCalibrator(const std::string& cache_file_path)
       : cache_file_path_(cache_file_path) {}
 
+    /**
+     * @brief Get the Batch Size for the next batch (always 1 due to issues with TRT and explicit batch)
+     *
+     * @return int
+     */
     int getBatchSize() const override {
         // HACK: TRTorch only uses explict batch sizing, INT8 Calibrator does not
         // work when reporting the batch size here and having explicity batching.
         // So we just report batch size 1 (warnings will still be printed out).
         return 1;
     }
 
+    /**
+     * @brief Get the next Batch
+     *
+     * Not used always returns false
+     *
+     * @param bindings: void*[] - An array of binding pointers (fed in from TensorRT calibrator), these buffers should be filed with batch data for each input
+     * @param names: const char*[] - Names of bindings
+     * @param nbBindings: int - Number of bindings
+     * @return false
+     */
     bool getBatch(void* bindings[], const char* names[], int nbBindings) override {
         return false;
     }
 
+    /**
+     * @brief Read calibration cache
+     *
+     * How to read from the calibration cache, only enabled if use_cache is set
+     *
+     * @param length
+     * @return const void* - Pointer to cache data
+     */
     const void* readCalibrationCache(size_t& length) override {
         std::stringstream ss;
         ss << "Reading Calibration Cache from " << cache_file_path_;
@@ -143,6 +241,15 @@ class Int8CacheCalibrator : Algorithm {
         return cache_size_ ? cache_.data() : nullptr;
     }
 
+
+    /**
+     * @brief Write calibration cache
+     *
+     * Write a the calibration cache provided by TensorRT to a specified file
+     *
+     * @param cache: const void* - cache data
+     * @param length: size_t - length of cache
+     */
     void writeCalibrationCache(const void* cache, size_t length) override {
         std::ofstream cache_file(cache_file_path_, std::ios::binary);
         cache_file.write(reinterpret_cast<const char*>(cache), length);
@@ -151,13 +258,23 @@ class Int8CacheCalibrator : Algorithm {
         logging::log(logging::Level::kINFO, ss.str());
     }
 
+    /**
+     * @brief operator to cast to nvinfer1::IInt8Calibrator*
+     *
+     * Convience function to convert to a IInt8Calibrator* to easily be assigned to the ptq_calibrator field in ExtraInfo
+     *
+     * @return nvinfer1::IInt8Calibrator*
+     */
     operator nvinfer1::IInt8Calibrator* () {
         return reinterpret_cast<nvinfer1::IInt8Calibrator*>(this);
     }
 
 private:
+    /// Path to cache file
     const std::string& cache_file_path_;
+    /// Size of cache
     size_t cache_size_ = 0;
+    /// Cache data
     std::vector<char> cache_;
 };