megenginelite.utils#

class TensorBatchCollector(shape, dtype=LiteDataType.LITE_INT8, device_type=LiteDeviceType.LITE_CUDA, device_id=0, is_pinned_host=False, tensor=None)[source]#

A tensor utils is used to collect many single batch tensor to a multi batch size tensor, when the multi batch size tensor collect finish, the result tensor can be get and send to the model input for forwarding.

when collect single batch tensor, the single batch tensor is no need in the same device_type and device_id with the result tensor, however the dtype must match and the shape must match except the highest dimension.

Parameters:

shape – the multi batch size tensor shape, After collection, the result tensor shape.
dtype (LiteDataType) – the datatype of the single batch tensor and the result tensor, default value is LiteDataType.LITE_INT8.
device_type (LiteDeviceType) – the target device type the result tensor will allocate, default value is LiteDeviceType.LITE_CUDA.
device_id – the device id the result tensor will allocate, default 0.
is_pinned_host – Whether the memory is pinned memory, refer to CUDA pinned memory, default False.
tensor (LiteTensor) – the result tensor, user can also create the multi batch size tensor and then create the TensorBatchColletor, if tensor is not None, all the member, such as shape, dtype, device_type, device_id, is_pinned_host will get from the tensor, if the tensor is None and the result tensor will create by the TensorBatchCollector, default is None.

Note

when collect tensor, the single batch tensor or array shape must match the result tensor shape except the batch size dimension (the highest dimension)

Examples

import numpy as np
batch_tensor = TensorBatchCollector([4, 8, 8])
arr = np.ones([8, 8], "int8")
for i in range(4):
    batch_tensor.collect(arr)
    arr += 1
data = batch_tensor.to_numpy()
assert data.shape[0] == 4
assert data.shape[1] == 8
assert data.shape[2] == 8
for i in range(4):
    for j in range(64):
        assert data[i][j // 8][j % 8] == i + 1

collect(array)[source]#

Collect a single batch through an array and store the array data to an empty batch, the empty batch is the front batch id in free list.

Parameters:: array – an array maybe LiteTensor or numpy ndarray, the shape must match the result tensor shape except the highest dimension

collect_by_ctypes(data, length)[source]#

Collect a single batch through an ctypes memory buffer and store the ctypes memory data to an empty batch, the empty batch is the front batch id in free list.

Parameters:: array – an array maybe LiteTensor or numpy ndarray, the shape must match the result tensor shape except the highest dimension

collect_id(array, batch_id)[source]#

Collect a single batch through an array and store the array data to the specific batch_id.

Parameters:

array – an array maybe LiteTensor or numpy ndarray, the shape of array must match the result tensor shape except the highest dimension.
batch_id – the batch id to store the array data to the result tensor, if the batch_id has already collected, a warning will generate.

free(indexes)[source]#

free the batch ids in the indexes, after the batch id is freed, it can be collected again without warning.

Parameters:: indexes – a list of to be freed batch id

get()[source]#: After finish collection, get the result tensor

to_numpy()[source]#: Convert the result tensor to a numpy ndarray