megenginelite.utils#
- class TensorBatchCollector(shape, dtype=LiteDataType.LITE_INT8, device_type=LiteDeviceType.LITE_CUDA, device_id=0, is_pinned_host=False, tensor=None)[源代码]#
张量小工具用于收集诸多单批量张量并合并成一个多批量张量，当多批量张量生成时，其结果会被用于模型输入进行前向推理
当收集单批量张量时， 这些张量必须与结果张量具有同样的设备类型和设备ID, 而且张量的数据类型和除过最高维之外的shape 也必须匹配
- 参数:
shape – 多批次张量的形状，合成后结果张量的形状
dtype (LiteDataType) – 单批次数据和结果张量的数据类型，默认为LiteDataType.LITE_INT8.
device_type (LiteDeviceType) – 结果张量分配内存的目标设备类型，默认为LiteDeviceType.LITE_CUDA
device_id – 结果张量待分配内存的设备ID，默认为0
is_pinned_host – 分配内存是否为页锁定内存，参考CUDA页锁定内存，默认为False
tensor (LiteTensor) – 结果张量，用户也可以先创建一个多批次张量然后再创建TensorBatchColletor， 当张量不为None 时， 诸如shape, dtype, device_type, device_id, is_pinned_host等参数都从tensor获取，当张量为None时， 结果张量用TensorBatchCollector创建，默认为None
备注
当收集张量时，单批次的张量或数组形状除过批次所在维度（最高维度）外必须和结果张量的形状匹配
实际案例
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)[源代码]#
收集数组中的单批次数据并将其存储为空批次，该空批次是释放列表前面的batch id
- 参数:
array – LiteTensor 或 numpy ndarray类型的数组，其形状除过最高维外必须和结果张量的形状相匹配
- collect_by_ctypes(data, length)[源代码]#
收集ctypes 类型的内存buffer中的单批次数据并将对应数据存储到某个空batch， 该空batch是释放列表前面的batch id
- 参数:
array – LiteTensor 或 numpy ndarray类型的数组，其形状除过最高维外必须和结果张量的形状相匹配