TensorRT CUDA Buffers

class dolphin.CudaTrtBuffers(stream: Optional[pycuda.driver.Stream] = None)[source]

Bases: object

To be used with the darray class. This class actually manages the darray used by the engine, both for inputs and outputs.

To ease the use of the darray, this class can be understood as a dict in order to name inputs and outputs.

Note that the names of inputs and outputs have to match the names of the inputs and outputs of the engine.

The constructor of this class takes an optional cuda.Stream as an argument.

allocate_input(name: str, shape: tuple, buffer_size: int, dtype: object, buffer_full_hook: Optional[callable] = None, flush_hook: Optional[callable] = None, allocate_hook: Optional[callable] = None, append_one_hook: Optional[callable] = None, append_multiple_hook: Optional[callable] = None) → None[source]

Method to allocate an input buffer. This methods creates a dolphin.Bufferizer and adds it to the inputs dict with the given name.

Parameters:

  • name (str) – Name of the input.

  • shape (tuple) – Shape of a single element in the buffer.

  • buffer_size (int) – Size of the buffer.

  • dtype (object) – Dtype of the buffer.

  • buffer_full_hook (callable) – callable function called each time the buffer is full.

  • flush_hook (callable) – callable function called each time the buffer is flushed.

  • allocate_hook (callable) – callable function called each time the allocate method is called.

  • append_one_hook (callable) – callable function called each time the append_one method is called.

  • append_multiple_hook (callable) – callable function called each time the append_multiple method is called.

allocate_output(name: str, shape: tuple, dtype: dtype) → None[source]

Method to allocate an output buffer. Oppositely to the inputs, the outputs are not dolphin.Bufferizer. They are darray.

Parameters:

  • name (str) – Name of the output.

  • shape (tuple) – Shape of a single element in the buffer.

  • dtype (dolphin.dtype) – Dtype of the buffer.

flush(value: Any = 0) → None[source]

Method to flush all the input buffers. Note that this method will trigger the flush_hook of each dolphin.Bufferizer.

Parameters:

value (int, optional) – value to initialize the inputs with, defaults to 0

append_one_input(name: str, data: darray)[source]

Method to append a single element to the input buffer. Note that this method will trigger the append_one_hook of the dolphin.Bufferizer. It will also trigger the buffer_full_hook if the buffer is full.

Parameters:

  • name (str) – Name of the input.

  • data (dolphin.darray) – Data to append.

append_multiple_input(name: str, data: darray)[source]

Method to append multiple elements to the input buffer. Note that this method will trigger the append_multiple_hook of the dolphin.Bufferizer. It will also trigger the buffer_full_hook if the buffer is full.

Parameters:

  • name (str) – Name of the input.

  • data (dolphin.darray) – Data to append.

property input_shape: Dict[str, tuple]

Property to get the shape of the inputs. Returns a dict with the name of the input as key and the shape as value.

Returns:

Shape of the inputs.

Return type:

Dict[str, tuple]

property input_dtype: Dict[str, dtype]

Property to get the dtype of the inputs. Returns a dict with the name of the input as key and the dtype as value.

Returns:

Dtype of the inputs.

Return type:

Dict[str, dolphin.dtype]

property output_shape: Dict[str, tuple]

Property to get the shape of the outputs. Returns a dict with the name of the output as key and the shape as value.

Returns:

Shape of the outputs.

Return type:

Dict[str, tuple]

property output_dtype

Property to get the dtype of the outputs. Returns a dict with the name of the output as key and the dtype as value.

Returns:

Dtype of the outputs.

Return type:

Dict[str, dolphin.dtype]

property full: bool

Property to check if the buffer is full. Returns True if at least one of the input buffer is full.

Returns:

True if at least one of the input buffer is full.

Return type:

bool

property output: Dict[str, darray]

Property to get the output of the buffer. Returns a dict with the name of the output as key and the output as value.

Returns:

Output of the buffer.

Return type:

Dict[str, dolphin.darray]

property input_bindings: List[darray]

Property to get the input bindings. ‘input bindings’ refers here to the list of input allocations

Returns:

Input bindings.

Return type:

List[dolphin.darray]

property output_bindings: List[int]

Property to get the output bindings. ‘output bindings’ refers here to the list of output allocations

Returns:

Output bindings.

Return type:

List[int]

property bindings

Property to get the bindings.

Returns:

Bindings.

Return type:

List[int]