megenginelite.utils

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

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.

参数
  • 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.

注解

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

实际案例

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)[源代码]

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.

参数

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)[源代码]

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.

参数

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)[源代码]

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

参数
  • 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)[源代码]

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

参数

indexes – a list of to be freed batch id

get()[源代码]

After finish collection, get the result tensor

get_tensor_at(idx)[源代码]

get the tensor from the internal big tensor by the idx, make sure the idx is not freed, return the tensor

参数

idx – the tensor index in the internal big tensor

to_numpy()[源代码]

Convert the result tensor to a numpy ndarray