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