在NVIDIA Drive PX上利用TensorRT 3 进行快速INT8推理

import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from PIL import Imageimport ctypes
import tensorrt as trt

CHANNEL = 3
HEIGHT = 512
WIDTH = 1024class PythonEntropyCalibrator(trt.infer.EntropyCalibrator):
  def __init__(self, input_layers, stream):
    trt.infer.EntropyCalibrator.__init__(self)       
    self.input_layers = input_layers
    self.stream = stream
  self.d_input = cuda.mem_alloc(self.stream.calibration_data.nbytes)
    stream.reset()

  def get_batch_size(self):
    return self.stream.batch_size

  def get_batch(self, bindings, names):
    batch = self.stream.next_batch()
    if not batch.size:   
      return None
      
    cuda.memcpy_htod(self.d_input, batch)
    for i in self.input_layers[0]:
      assert names[0] != i

    bindings[0] = int(self.d_input)
    return bindings

  def read_calibration_cache(self, length):
    return None

  def write_calibration_cache(self, ptr, size):
    cache = ctypes.c_char_p(int(ptr))
    with open('calibration_cache.bin', 'wb') as f:
      f.write(cache.value)
    return None

class ImageBatchStream():
  def __init__(self, batch_size, calibration_files, preprocessor):
    self.batch_size = batch_size
    self.max_batches = (len(calibration_files) // batch_size) + 
                       (1 if (len(calibration_files) % batch_size) 
                        else 0)
    self.files = calibration_files
    self.calibration_data = np.zeros((batch_size, CHANNEL, HEIGHT, WIDTH), 
                                     dtype=np.float32)
    self.batch = 0
    self.preprocessor = preprocessor

  @staticmethod
  def read_image_chw(path):
    img = Image.open(path).resize((WIDTH,HEIGHT), Image.NEAREST)
    im = np.array(img, dtype=np.float32, order='C')
    im = im[:,:,::-1]
    im = im.transpose((2,0,1))
    return im
         
  def reset(self):
    self.batch = 0
     
  def next_batch(self):
    if self.batch < self.max_batches:
      imgs = []
      files_for_batch = self.files[self.batch_size * self.batch : 
                        self.batch_size * (self.batch + 1)]
      for f in files_for_batch:
        print("[ImageBatchStream] Processing ", f)
        img = ImageBatchStream.read_image_chw(f)
        img = self.preprocessor(img)
        imgs.append(img)
      for i in range(len(imgs)):
        self.calibration_data[i] = imgs[i]
      self.batch += 1
      return np.ascontiguousarray(self.calibration_data, dtype=np.float32)
    else:
      return np.array([])

import glob
from random import shuffle
import numpy as np
from PIL import Imageimport tensorrt as trt

import labels  #from cityscapes evaluation scriptimport calibrator    #calibrator.py

MEAN = (71.60167789, 82.09696889, 72.30508881)
MODEL_DIR = '/data/fcn8s/'
CITYSCAPES_DIR = '/data/Cityscapes/'
TEST_IMAGE = CITYSCAPES_DIR + 'leftImg8bit/val/lindau/lindau_000042_000019_leftImg8bit.png'
CALIBRATION_DATASET_LOC = CITYSCAPES_DIR + 'leftImg8bit/train/*/*.png'

CLASSES = 19
CHANNEL = 3
HEIGHT = 512
WIDTH = 1024def sub_mean_chw(data):
  data = data.transpose((1,2,0)) # CHW -> HWC
  data -= np.array(MEAN) # Broadcast subtract
  data = data.transpose((2,0,1)) # HWC -> CHW
  return data
             
def color_map(output):
  output = output.reshape(CLASSES, HEIGHT, WIDTH)
  out_col = np.zeros(shape=(HEIGHT, WIDTH), dtype=(np.uint8, 3))
  for x in range (WIDTH):
    for y in range (HEIGHT):
    out_col[y,x] = labels.id2label[labels.trainId2label[np.argmax(output[:,y,x])].id].color
  return out_col

def create_calibration_dataset():
  # Create list of calibration images (filename)
  # This sample code picks 100 images at random from training set
  calibration_files = glob.glob(CALIBRATION_DATASET_LOC)
  shuffle(calibration_files)
  return calibration_files[:100]def main():
  calibration_files = create_calibration_dataset()
  
  # Process 5 images at a time for calibration
  # This batch size can be different from MaxBatchSize (1 in this example)
  batchstream = calibrator.ImageBatchStream(5, calibration_files, sub_mean_chw)
  int8_calibrator = calibrator.PythonEntropyCalibrator(["data"], batchstream)
  
  # Easy to use TensorRT lite package
  engine = trt.lite.Engine(framework="c1",
                           deployfile=MODEL_DIR + "fcn8s.prototxt",
                           modelfile=MODEL_DIR + "fcn8s.caffemodel",
                           max_batch_size=1,
                           max_workspace_size=(256 << 20),
                           input_nodes={"data":(CHANNEL,HEIGHT,WIDTH)},
                           output_nodes=["score"],
                           preprocessors={"data":sub_mean_chw},
                           postprocessors={"score":color_map},
                           data_type=trt.infer.DataType.INT8,
                           calibrator=int8_calibrator,
                           logger_severity=trt.infer.LogSeverity.INFO)
                           
  test_data = calibrator.ImageBatchStream.read_image_chw(TEST_IMAGE)
  out = engine.infer(test_data)[0]
  test_img = Image.fromarray(out, 'RGB')
  test_img.show()

bool getBatch(void* bindings[], const char* names[], int nbBindings) override {
    return false;
  }
  
  const void* readCalibrationCache(size_t& length) override
  {
    mCalibrationCache.clear();
    std::ifstream input(mCacheFile, std::ios::binary);
    input >> std::noskipws;
    if (input.good()) {
      std::copy(std::istream_iterator(input),
      std::istream_iterator<char>(),
      std::back_inserter<char>(mCalibrationCache));
    }
    length = mCalibrationCache.size();
    return length ? &mCalibrationCache[0] : nullptr;
  }private:
  std::string mCacheFile;
  std::vector<char> mCalibrationCache;};