Skip to content

# Program Listing for File ptq.h

↰ Return to documentation for file ( `cpp/api/include/trtorch/ptq.h` )

    #pragma once

    #include <string>
    #include <vector>
    #include <memory>
    #include <iostream>
    #include <sstream>

    #ifndef DOXYGEN_SHOULD_SKIP_THIS
    namespace nvinfer1 {
    class IInt8Calibrator;
    class IInt8EntropyCalibrator2;
    }

    namespace torch {
    namespace data {
    template<typename Example>
    class Iterator;
    }
    }
    #endif //DOXYGEN_SHOULD_SKIP_THIS

    namespace trtorch {
    namespace ptq {

    template<typename Algorithm, typename DataLoaderUniquePtr>
    class Int8Calibrator : Algorithm {
        using DataLoader = typename DataLoaderUniquePtr::element_type;
        using Batch = typename DataLoader::super::BatchType;
    public:
        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) {}

        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;
            //return static_cast<int>(dataloader_->options().batch_size);
        }

        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_) {
                it_ = dataloader_->begin();
                it_created_ = true;
            }

            if (it_ == dataloader_->end()) {
                return false;
            }

            auto batch = *it_;

            for (int i = 0; i < nbBindings; i++) {
                auto data = batch.data;
                data = data.to(at::kCUDA).contiguous();
                bindings[i] = data.data_ptr();
            }

            it_ = ++it_;
            return true;
        }

        const void* readCalibrationCache(size_t& length) override {
            if (use_cache_) {
                std::stringstream ss;
                ss << "Reading Calibration Cache from " << cache_file_path_;
                logging::log(logging::Level::kINFO, ss.str());
                cache_.clear();
                std::ifstream cache_file(cache_file_path_, std::ios::binary);
                cache_file >> std::noskipws;
                if (cache_file.good()) {
                    std::copy(std::istream_iterator<char>(cache_file),
                                std::istream_iterator<char>(),
                                std::back_inserter(cache_));
                    ss << "Cache read";
                    logging::log(logging::Level::kDEBUG, ss.str());
                }
                cache_size_ = cache_.size();
                return cache_size_ ? cache_.data() : nullptr;
            }
            return nullptr;
        }

        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);
            std::stringstream ss;
            ss << "Saved Calibration Cache to " << cache_file_path_;
            logging::log(logging::Level::kINFO, ss.str());
        }

        operator nvinfer1::IInt8Calibrator* () {
            return reinterpret_cast<nvinfer1::IInt8Calibrator*>(this);
        }

    private:
        DataLoader* dataloader_;
        torch::data::Iterator<Batch> it_;
        const std::string& cache_file_path_;
        size_t cache_size_ = 0;
        bool use_cache_;
        std::vector<char> cache_;
        bool it_created_ = false;
    };

    template<typename Algorithm>
    class Int8CacheCalibrator : Algorithm {
    public:
        Int8CacheCalibrator(const std::string& cache_file_path)
          : cache_file_path_(cache_file_path) {}

        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;
        }

        bool getBatch(void* bindings[], const char* names[], int nbBindings) override {
            return false;
        }

        const void* readCalibrationCache(size_t& length) override {
            std::stringstream ss;
            ss << "Reading Calibration Cache from " << cache_file_path_;
            logging::log(logging::Level::kINFO, ss.str());
            cache_.clear();
            std::ifstream cache_file;
            cache_file.open(cache_file_path_, std::ios::in | std::ios::binary);
            cache_file.unsetf(std::ios::skipws);
            cache_file.seekg(0, std::ios::beg);
            cache_.reserve(cache_file.tellg());
            cache_file.seekg(0, std::ios::beg);
            if (cache_file.good()) {
                std::cout << "Trying to read cache" << std::endl;
                std::copy(std::istreambuf_iterator<char>(cache_file),
                            std::istreambuf_iterator<char>(),
                            std::back_inserter(cache_));
                ss << "Cache read";
                logging::log(logging::Level::kDEBUG, ss.str());
            }
            cache_size_ = cache_.size();
            return cache_size_ ? cache_.data() : nullptr;
        }

        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);
            std::stringstream ss;
            ss << "Saved Calibration Cache to " << cache_file_path_;
            logging::log(logging::Level::kINFO, ss.str());
        }

        operator nvinfer1::IInt8Calibrator* () {
            return reinterpret_cast<nvinfer1::IInt8Calibrator*>(this);
        }

    private:
        const std::string& cache_file_path_;
        size_t cache_size_ = 0;
        std::vector<char> cache_;
    };

    } // namespace ptq
    } // namespace trtorch