Function to allocate buffers and bindings for TensorRT inference. Args: engine (trt.ICudaEngine): batch_size (int): batch size to be used during inference. Returns: inputs (List): list of input buffers. outputs (List): list of output buffers. db
(engine: trt.ICudaEngine, batch_size: int)
| 47 | |
| 48 | |
| 49 | def allocate_buffers(engine: trt.ICudaEngine, batch_size: int) -> [list, list, list]: |
| 50 | """ |
| 51 | Function to allocate buffers and bindings for TensorRT inference. |
| 52 | |
| 53 | Args: |
| 54 | engine (trt.ICudaEngine): |
| 55 | batch_size (int): batch size to be used during inference. |
| 56 | |
| 57 | Returns: |
| 58 | inputs (List): list of input buffers. |
| 59 | outputs (List): list of output buffers. |
| 60 | dbindings (List): list of device bindings. |
| 61 | """ |
| 62 | inputs = [] |
| 63 | outputs = [] |
| 64 | dbindings = [] |
| 65 | |
| 66 | for binding in engine: |
| 67 | binding_shape = engine.get_binding_shape(binding) |
| 68 | if binding_shape[0] == TRT_DYNAMIC_DIM: # dynamic shape |
| 69 | size = batch_size * abs(trt.volume(binding_shape)) |
| 70 | else: |
| 71 | size = abs(trt.volume(binding_shape)) |
| 72 | dtype = trt.nptype(engine.get_binding_dtype(binding)) |
| 73 | # Allocate host and device buffers |
| 74 | host_mem = cuda.pagelocked_empty(size, dtype) |
| 75 | device_mem = cuda.mem_alloc(host_mem.nbytes) |
| 76 | # Append the device buffer to device bindings |
| 77 | dbindings.append(int(device_mem)) |
| 78 | |
| 79 | # Append to the appropriate list (input/output) |
| 80 | if engine.binding_is_input(binding): |
| 81 | inputs.append(HostDeviceMem(host_mem, device_mem)) |
| 82 | else: |
| 83 | outputs.append(HostDeviceMem(host_mem, device_mem)) |
| 84 | |
| 85 | return inputs, outputs, dbindings |
| 86 | |
| 87 | |
| 88 | def infer( |
no test coverage detected