Performs inference in TensorRT engine. Args: engine_path (str): path to the TensorRT engine. val_batches (tf.data.Dataset): validation dataset (batches). batch_size (int): batch size used for inference and dataset batch splitting. top_k_value (int): value of
(
engine_path: str,
val_batches,
batch_size: int = 8,
top_k_value: int = 1,
)
| 86 | |
| 87 | |
| 88 | def infer( |
| 89 | engine_path: str, |
| 90 | val_batches, |
| 91 | batch_size: int = 8, |
| 92 | top_k_value: int = 1, |
| 93 | ) -> None: |
| 94 | """ |
| 95 | Performs inference in TensorRT engine. |
| 96 | |
| 97 | Args: |
| 98 | engine_path (str): path to the TensorRT engine. |
| 99 | val_batches (tf.data.Dataset): validation dataset (batches). |
| 100 | batch_size (int): batch size used for inference and dataset batch splitting. |
| 101 | top_k_value (int): value of `K` for the top K predictions used in the accuracy calculation. |
| 102 | |
| 103 | Raises: |
| 104 | RuntimeError: raised when loading images in the host fails. |
| 105 | """ |
| 106 | |
| 107 | def override_shape(shape: tuple) -> tuple: |
| 108 | """Overrides batch dimension if dynamic.""" |
| 109 | if TRT_DYNAMIC_DIM in shape: |
| 110 | shape = tuple( |
| 111 | [batch_size if dim == TRT_DYNAMIC_DIM else dim for dim in shape] |
| 112 | ) |
| 113 | return shape |
| 114 | |
| 115 | # Open engine as runtime |
| 116 | with open(engine_path, "rb") as f, trt.Runtime( |
| 117 | trt.Logger(trt.Logger.ERROR) |
| 118 | ) as runtime: |
| 119 | engine = runtime.deserialize_cuda_engine(f.read()) |
| 120 | |
| 121 | # Allocate buffers and create a CUDA stream. |
| 122 | inputs, outputs, dbindings = allocate_buffers(engine, batch_size) |
| 123 | |
| 124 | # Initiate test_accuracy |
| 125 | test_accuracy = tf.keras.metrics.SparseTopKCategoricalAccuracy( |
| 126 | k=top_k_value, name="top_k_accuracy", dtype=tf.float32 |
| 127 | ) |
| 128 | test_accuracy.reset_states() |
| 129 | |
| 130 | # Contexts are used to perform inference. |
| 131 | with engine.create_execution_context() as context: |
| 132 | |
| 133 | # Resolves dynamic shapes in the context |
| 134 | for binding in engine: |
| 135 | binding_idx = engine.get_binding_index(binding) |
| 136 | binding_shape = engine.get_binding_shape(binding_idx) |
| 137 | if engine.binding_is_input(binding_idx): |
| 138 | binding_shape = override_shape(binding_shape) |
| 139 | context.set_binding_shape(binding_idx, binding_shape) |
| 140 | |
| 141 | if isinstance(val_batches, tf.Tensor): |
| 142 | # Load images in Host (flatten and copy to page-locked buffer in Host) |
| 143 | data = val_batches.numpy().astype(np.float32).ravel() |
| 144 | pagelocked_buffer = inputs[0].host |
| 145 | np.copyto(pagelocked_buffer, data) |
no test coverage detected