TRTorch PTQ in Python

[narendasan]

TRTorch PTQ in Python

Goal: Implement support for using PyTorch DataLoaders inside TensorRT Int8Calibrators from Python

Explicit Requirements:

1. There should be no dependency on the TensorRT Python API
2. However the API for calibrators should remain compatible with the TensorRT Python API
3. Users should have an easy way to define the calibration algorithm

Proposed Supported Use-cases

1. Using PyTorch DataLoader in Calibrator

Primarily we want to support the use case where developers have existing data loaders perhaps from training and repurposing that infrastructure to quickly do calibration. All that should be required is providing a data loader to a constructor which will wrap it for use by the calibration APIs in TRT.

import torch
import trtorch 
from trtorch import ptq

...

class MyDataset(torch.utils.data.Dataset):
  'Characterizes a dataset for PyTorch'
  def __init__(self, list_IDs, labels):
        'Initialization'
        self.labels = labels
        self.list_IDs = list_IDs

  def __len__(self):
        'Denotes the total number of samples'
        return len(self.list_IDs)

  def __getitem__(self, index):
        'Generates one sample of data'
        # Select sample
        ID = self.list_IDs[index]

        # Load data and get label
        X = torch.load('data/' + ID + '.pt')
        y = self.labels[ID]

        return X, y
      
      
params = {'batch_size': 64,
          'shuffle': True,
          'num_workers': 6}

calibration_set = Dataset(partition['validation'], labels)
calib_dataloader = torch.utils.data.DataLoader(calibration_set, **params)


calibrator = trtorch.ptq.DataLoaderCalibrator(calib_dataloader, "path to cache", use_cache=True) #where true is use cache

trt_model = trtorch.compile(ts_model, {
  ...
  "calibrator": calibrator
})

Potential Implementation

We want to make sure there is no explicit dependency on the TensorRT Python API so we should reexport the bare minimum of the Int8Calibrator APIs and implement a class in python using the data loader

# File: trtorch/ptq.py

import torch
import trtorch._C._exposed_ptq_api as ptq

class DataLoaderCalibrator(ptq.IInt8Calibrator):
  def __init__(self, dl: DataLoader, path: str, use_cache: bool, algo: ptq.IInt8Calibrator=ptq.IInt8EntropyCalibrator2):
    super(self, algo)
    
  def get_batch_size(self):
     return self.batch_size

  def get_batch(self, names):
  	# Returns torch device handles 
    return [handles]

  def read_calibration_cache(self):
  	# If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None.

  def write_calibration_cache(self, cache):

2. Using standard TensorRT Calibrator

We want to make sure we maintain compatibility with existing TRT infrastructure, so if users have existing calibrator implementations in python they can reuse it with TRTorch without code changes.

import torch
import trtorch 
import tensorrt as trt

...

class MNISTEntropyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, training_data, cache_file, batch_size=64):
        # Whenever you specify a custom constructor for a TensorRT class,
        # you MUST call the constructor of the parent explicitly.
        trt.IInt8EntropyCalibrator2.__init__(self)

        self.cache_file = cache_file

        # Every time get_batch is called, the next batch of size batch_size will be copied to the device and returned.
        self.data = load_mnist_data(training_data)
        self.batch_size = batch_size
        self.current_index = 0

        # Allocate enough memory for a whole batch.
        self.device_input = cuda.mem_alloc(self.data[0].nbytes * self.batch_size)

    def get_batch_size(self):
        return self.batch_size

    # TensorRT passes along the names of the engine bindings to the get_batch function.
    # You don't necessarily have to use them, but they can be useful to understand the order of
    # the inputs. The bindings list is expected to have the same ordering as 'names'.
    def get_batch(self, names):
        if self.current_index + self.batch_size > self.data.shape[0]:
            return None

        current_batch = int(self.current_index / self.batch_size)
        if current_batch % 10 == 0:
            print("Calibrating batch {:}, containing {:} images".format(current_batch, self.batch_size))

        batch = self.data[self.current_index:self.current_index + self.batch_size].ravel()
        cuda.memcpy_htod(self.device_input, batch)
        self.current_index += self.batch_size
        return [self.device_input]


    def read_calibration_cache(self):
        # If there is a cache, use it instead of calibrating again. Otherwise, implicitly return None.
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()

    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)
            
calib = MNISTEntropyCalibrator(test_set, cache_file=calibration_cache)

trt_model = trtorch.compile(ts_model, {
  ...
  "calibrator": calib
})

3. Using cache file directly

Finally if users already have a cache file they can simply point to it and do calibration.

import torch
import trtorch 
from trtorch import ptq

...

calibrator = trtorch.ptq.CacheFileCalibrator("path to cache")

trt_model = trtorch.compile(ts_model, {
  ...
  "calibrator": calibrator
})

Potential Approach to Implementation

1. Re-export base PTQ APIs from TensorRT within the context of the TRTorch Python APIs
   1. This lets us avoid a direct dependency on the TRT Python API
2. Implement the DataLoader and cache calibrators in Python vs. C++ for flexibility and maintainability
3. We need ensure that there is some notion of symbol compatibly between a TRTorch Calibrator and a TensorRT Calibrator
   1. This may require digging into how PyBind works and if we can do something to ensure that it treats the two the same
4. We need some sort of contemplating or meta-programing system to allow the calibration algorithm to be changed without adding a different calibrator for each type of algorithm.

[narendasan]

cc: @peri044