feat(//cpp/ptq): Add a feature to the dataset to use less than the full

test set for calibration

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

cpp/ptq/datasets/cifar10.cpp

@@ -1,5 +1,6 @@
 #include "cpp/ptq/datasets/cifar10.h"
 
+#include "torch/torch.h"
 #include "torch/data/example.h"
 #include "torch/types.h"
 
@@ -63,15 +64,19 @@ std::pair<torch::Tensor, torch::Tensor> read_batch(const std::string& path) {
 }
 
 std::pair<torch::Tensor, torch::Tensor> read_train_data(const std::string& root) {
-    torch::Tensor images, targets;
+    std::vector<torch::Tensor> images, targets;
     for(uint32_t i = 1; i <= 5; i++) {
         std::stringstream ss;
         ss << root << '/' << kTrainFilenamePrefix << i << ".bin";
         auto batch = read_batch(ss.str());
-        images = torch::stack({images, batch.first});
-        targets = torch::stack({targets, batch.second});
+        images.push_back(batch.first);
+        targets.push_back(batch.second);
     }
-    return std::make_pair(images, targets);
+
+    torch::Tensor image_tensor = std::accumulate(++images.begin(), images.end(), *images.begin(), [&](torch::Tensor a, torch::Tensor b) {return torch::cat({a, b}, 0);});
+    torch::Tensor target_tensor = std::accumulate(++targets.begin(), targets.end(), *targets.begin(), [&](torch::Tensor a, torch::Tensor b) {return torch::cat({a, b}, 0);});
+
+    return std::make_pair(image_tensor, target_tensor);
 }
 
 std::pair<torch::Tensor, torch::Tensor> read_test_data(const std::string& root) {
@@ -93,6 +98,7 @@ CIFAR10::CIFAR10(const std::string& root, Mode mode)
 
     images_ = std::move(data.first);
     targets_ = std::move(data.second);
+    assert(images_.sizes()[0] == images_.sizes()[0]);
 }
 
 torch::data::Example<> CIFAR10::get(size_t index) {
@@ -115,5 +121,12 @@ const torch::Tensor& CIFAR10::targets() const {
     return targets_;
 }
 
+CIFAR10&& CIFAR10::use_subset(int64_t new_size) {
+    assert(new_size <= images_.sizes()[0]);
+    images_ = images_.slice(0, 0, new_size);
+    targets_ = targets_.slice(0, 0, new_size);
+    return std::move(*this);
+}
+
 } // namespace datasets
 

cpp/ptq/datasets/cifar10.h

@@ -34,6 +34,10 @@ class CIFAR10 : public torch::data::datasets::Dataset<CIFAR10> {
     // Returns all targets stacked into a single tensor
     const torch::Tensor& targets() const;
 
+    // Trims the dataset to the first n pairs
+    CIFAR10&& use_subset(int64_t new_size);
+
+
 private:
     Mode mode_;
     torch::Tensor images_, targets_;

cpp/ptq/main.cpp

@@ -32,6 +32,7 @@ int main(int argc, const char* argv[]) {
     // Create the calibration dataset
     const std::string data_dir = std::string(argv[2]);
     auto calibration_dataset = datasets::CIFAR10(data_dir, datasets::CIFAR10::Mode::kTest)
+                                    .use_subset(320)
                                     .map(torch::data::transforms::Normalize<>({0.4914, 0.4822, 0.4465},
                                                                               {0.2023, 0.1994, 0.2010}))
                                     .map(torch::data::transforms::Stack<>());
@@ -41,19 +42,19 @@ int main(int argc, const char* argv[]) {
 
     std::string calibration_cache_file = "/tmp/vgg16_TRT_ptq_calibration.cache";
 
-    auto calibrator = trtorch::ptq::make_int8_calibrator(std::move(calibration_dataloader), calibration_cache_file, true);
+    auto calibrator = trtorch::ptq::make_int8_calibrator(std::move(calibration_dataloader), calibration_cache_file, false);
     //auto calibrator = trtorch::ptq::make_int8_cache_calibrator(calibration_cache_file);
 
 
     std::vector<std::vector<int64_t>> input_shape = {{32, 3, 32, 32}};
     // Configure settings for compilation
     auto extra_info = trtorch::ExtraInfo({input_shape});
     // Set operating precision to INT8
-    extra_info.op_precision = torch::kChar;
+    extra_info.op_precision = torch::kFI8;
     // Use the TensorRT Entropy Calibrator
     extra_info.ptq_calibrator = calibrator;
-    // Increase the default workspace size;
-    extra_info.workspace_size = 1 << 30;
+    // Set max batch size for the engine
+    extra_info.max_batch_size = 32;
 
     mod.eval();
 
@@ -92,6 +93,14 @@ int main(int argc, const char* argv[]) {
 
         auto outputs = trt_mod.forward({images});
         auto predictions = std::get<1>(torch::max(outputs.toTensor(), 1, false));
+        predictions = predictions.reshape(predictions.sizes()[0]);
+
+        if (predictions.sizes()[0] != targets.sizes()[0]) {
+            // To handle smaller batches util Optimization profiles work
+            predictions = predictions.slice(0, 0, targets.sizes()[0]);
+        }
+
+        std:: cout << predictions << targets << std::endl;
 
         total += targets.sizes()[0];
         correct += torch::sum(torch::eq(predictions, targets)).item().toFloat();