Implements the INT8 MinMax Calibrator.
| 34 | import common |
| 35 | |
| 36 | class EngineCalibrator(trt.IInt8MinMaxCalibrator): |
| 37 | """ |
| 38 | Implements the INT8 MinMax Calibrator. |
| 39 | """ |
| 40 | |
| 41 | def __init__(self, cache_file): |
| 42 | """ |
| 43 | :param cache_file: The location of the cache file. |
| 44 | """ |
| 45 | super().__init__() |
| 46 | self.cache_file = cache_file |
| 47 | self.image_batcher = None |
| 48 | self.batch_allocation = None |
| 49 | self.batch_generator = None |
| 50 | |
| 51 | def set_image_batcher(self, image_batcher: ImageBatcher): |
| 52 | """ |
| 53 | Define the image batcher to use, if any. If using only the cache file, an image batcher doesn't need |
| 54 | to be defined. |
| 55 | :param image_batcher: The ImageBatcher object |
| 56 | """ |
| 57 | self.image_batcher = image_batcher |
| 58 | self.size = int(np.dtype(self.image_batcher.dtype).itemsize * np.prod(self.image_batcher.shape)) |
| 59 | self.batch_allocation = common.cuda_call(cudart.cudaMalloc(self.size)) |
| 60 | self.batch_generator = self.image_batcher.get_batch() |
| 61 | |
| 62 | def get_batch_size(self): |
| 63 | """ |
| 64 | Overrides from trt.IInt8MinMaxCalibrator. |
| 65 | Get the batch size to use for calibration. |
| 66 | :return: Batch size. |
| 67 | """ |
| 68 | if self.image_batcher: |
| 69 | return self.image_batcher.batch_size |
| 70 | return 1 |
| 71 | |
| 72 | def get_batch(self, names): |
| 73 | """ |
| 74 | Overrides from trt.IInt8MinMaxCalibrator. |
| 75 | Get the next batch to use for calibration, as a list of device memory pointers. |
| 76 | :param names: The names of the inputs, if useful to define the order of inputs. |
| 77 | :return: A list of int-casted memory pointers. |
| 78 | """ |
| 79 | if not self.image_batcher: |
| 80 | return None |
| 81 | try: |
| 82 | batch, _, _ = next(self.batch_generator) |
| 83 | log.info("Calibrating image {} / {}".format(self.image_batcher.image_index, self.image_batcher.num_images)) |
| 84 | common.memcpy_host_to_device(self.batch_allocation, np.ascontiguousarray(batch)) |
| 85 | |
| 86 | return [int(self.batch_allocation)] |
| 87 | except StopIteration: |
| 88 | log.info("Finished calibration batches") |
| 89 | return None |
| 90 | |
| 91 | def read_calibration_cache(self): |
| 92 | """ |
| 93 | Overrides from trt.IInt8MinMaxCalibrator. |