Package: optimize-tensorrt
Всем привет )Этот код нужен что-бы оптимизировать модель в TensorRT и инференсить в лайве.Важный момент: engine файлы не совместимы между системами, нужно оптимизировать под каждую систему отдельно.Recommended pipelineЗаранее перегнать модель из фреймворка в ONNX. Для этого есть несколько функций в этом модуле.python -c "import onnx_converter; help(onnx_converter)"Скачать эту либу и onnx файл модели на необходимое железо.Поставить все необходимые зависимостиУказать все пути для энвов из пунктаInstallationЗапустить оптимизатор onnx2trt:python onnx2trt.py --helpДля INT8 необходимо заранее создать calibration файл через onnx2trt_handmade.python -c "from onnx2trt_handmade import onnx2trt_handmade; help(onnx2trt_handmade)"Запомнить куда сохранился итоговый engine файл.Для инференса использовать класс InferenceTRT.python -c "from inference import InferenceTRT; help(InferenceTRT)"InstallationEnvsCUDA_HOME: Путь до папки с кудой.LD_LIBRARY_PATH:Путь доlib64в папке с кудой.Путь доtargets/x86_64-linux/lib/в папке с кудой.Путь доlibв папке с TensorRT.trtexecДля запуска onnx2trt необходимо указывать путь до trtexec. Как правило, он хранится по пути/targets/x86_64-linux-gnu/bin/trtexecв папке с TensorRT.TensorRTTested on version8.2.2CUDAPyCUDAConvert to ONNXfastai2onnxКонвертируем модельку из FastAI в ONNX.Args:fastai_path(str): Path to fastai model file.input_shape(tuple or list): Input shape (with batch dimension).onnx_path(str): Path to save converted model.input_names(tuple or list): Model's input names.output_names(tuple or list): Model's output names.torch2onnxКонвертируем модельку из PyTorch в ONNX.Args:model(torch.nn.Module): PyTorch model.input_shape(tuple or list): Input shape (with batch dimension).onnx_path(str): Path to save converted model.input_names(tuple or list): Model's input names.output_names(tuple or list): Model's output names.keras2onnxКонвертируем модельку из Keras в ONNX.Args:model(): Keras model.input_shape(tuple or list): Input shape (with batch dimension).onnx_path(str): Path to save converted model.Optimizeonnx2trtСамый оптимальный метод для оптимизации модельки под TensorRT. В конце конвертации выводит бенчмарки. Не поддерживает калибровку для INT8.NB! Для каждой машины нужно конвертировать отдельно, просто перетащить файлик engine не сработает.Args:onnx_path(str): Path to onnx model (opset9).engine_path(str): Path to save engine.dtype(str or np.dtype): Model data type.use_sparsity(bool): Use sparsity boolean.trtexec_path(str): Path to trtexec binary.workspace_size(int): Size of workspace in MB.onnx2trt_handmadeОптимизация модельки с использованием TensorRT Python API.
Лучше всего использовать один раз для создания calibration файлика, его можно переносить между системами. И под каждую оптимизировать через trtexec.Args:onnx_path(str): Path to ONNX model.engine_path(str): Path to save TensorRT engine.dtype(str): Model data type. (fp32, fp16 or int8).calib(trt.IInt8EntropyCalibrator2): Calibration class. Only for INT8.CalibrationКласс для агрегации данных во время калибровки INT8.__init__Args:calibration_files(list or tuple): List of paths to images for calibration.preprocessor(function): Image preprocessing function.shape(list or tuple): Input image shape. With batch_size.cache_file(str): Path to save calibration.InferenceInferenceTRTКласс для инференса модельки оптимизированной в engine.__init__Initialize class variables.Args:engine_path(str): Path to TensorRT engine file.device(cuda.Device): Cuda device to run inference on.load_engineLoad TensorRT engine from file.Args:engine_path(str): Path to TensorRT engine file.predictGet model prediction.Args:inputs(np.ndarray): Model inputs array.Returns:Model prediction as np.ndarray.