Implementación del modelo AI: implementación en Python de la cuantificación INT8 del modelo TensorRT

Implementación del modelo de IA: implementación en Python de la cuantificación INT8 del modelo TensorRT

Este artículo se publicó por primera vez en la cuenta pública [DeepDriving], bienvenido a prestar atención.

Descripción general

32En la actualidad, los parámetros de los modelos de aprendizaje profundo se expresan básicamente en bits de punto flotante ( ) durante la fase de entrenamiento FP32, de modo que se puede utilizar un rango dinámico mayor para actualizar los parámetros durante el proceso de entrenamiento. Sin embargo, en la etapa de inferencia, FP32la precisión utilizada consumirá más recursos informáticos y espacio de memoria. Por esta razón, al implementar el modelo, a menudo se utiliza un método para reducir la precisión del modelo, utilizando punto flotante de bits () o bit 16con FP16signo 8. entero ( INT8) Para representar. Generalmente no hay pérdida de precisión al FP32convertir a , pero la conversión a 2 puede causar una pérdida de precisión mayor, especialmente cuando los pesos del modelo se distribuyen en un rango dinámico grande.FP16FP32INT8

Aunque hay una cierta pérdida de precisión, la conversión INT8también traerá muchos beneficios, como reducir CPUla ocupación del espacio de almacenamiento y la memoria, mejorar el rendimiento informático, etc. Esto es muy significativo en plataformas integradas con recursos informáticos limitados.

Para convertir el tensor del parámetro del modelo de FP32a INT8, es decir, el rango al que se asigna el rango dinámico del tensor de coma flotante [-128,127], puede utilizar la siguiente fórmula:

$$X_{}=Clip\left(Round\right(x\_\_{}_{}/escala\left)\right)\_\_$$

Entre ellos, $Clip$和$Round$ representa operaciones de truncamiento y redondeo respectivamente. Como se puede ver en la fórmula anterior,FP32convertirINT8es establecer una$Laescalase$ utiliza para el mapeo. Este proceso de mapeo se llama cuantificación. La fórmula anterior es una fórmula de cuantificación simétrica$.$

La clave de la cuantificación es encontrar un factor de escala apropiado para que la precisión del modelo cuantificado sea lo más cercana posible al modelo original. Hay dos formas de cuantificar un modelo:

• La cuantificación posterior al entrenamiento ( ) consiste en calcular el factor de escala mediante un proceso de calibración ( ) Post-training quantization，PTQdespués de que el modelo se entrena para lograr el proceso de cuantificación.Calibration

• El entrenamiento consciente de la cuantificación ( Quantization-aware training，QAT) calcula el factor de escala durante el proceso de entrenamiento del modelo, lo que permite compensar el error de precisión causado por las operaciones de cuantificación y de cuantificación inversa durante el proceso de entrenamiento.

Este artículo solo presenta cómo llamar a la TensorRTinterfaz Pythonpara lograr INT8la cuantificación. En cuanto INT8al conocimiento teórico de la cuantificación, dado que implica mucho contenido, escribiré un artículo especial para presentarlo cuando tenga tiempo.

Implementación específica de la cuantificación TensorRT INT8

Calibrador en TensorRT

Durante el proceso de cuantificación posterior al entrenamiento, TensorRTes necesario calcular el factor de escala de cada tensor en el modelo, proceso que se denomina calibración. El proceso de calibración requiere proporcionar datos representativos para TensorRTejecutar el modelo en este conjunto de datos y luego recopilar estadísticas para cada tensor para encontrar un factor de escala óptimo. Encontrar el factor de escala óptimo requiere equilibrar las dos fuentes de error de discretización (que se hace mayor a medida que el rango representado por cada valor cuantificado se hace mayor) y error de truncamiento (cuyos valores están restringidos a los límites del rango representable), proporcionando TensorRTvarios diferentes calibradores:

• IInt8EntropyCalibrator2 : el calibrador de entropía recomendado actualmente; de ​​forma predeterminada, la calibración se produce antes de la fusión de capas, recomendado para su uso CNNen modelos.

• IInt8MinMaxCalibrator : Este calibrador utiliza todo el rango de la distribución de activación para determinar el factor de escala. Por defecto, la calibración ocurre antes de la fusión de capas, recomendado NLPen los modelos para la tarea.

• IInt8EntropyCalibrator : este calibrador es TensorRTel calibrador de entropía original. De forma predeterminada, la calibración se produce después de la fusión de capas y su uso está actualmente en desuso.

• IInt8LegacyCalibrator : este calibrador requiere parametrización por parte del usuario; de forma predeterminada, la calibración se produce después de la fusión de capas y no se recomienda.

TensorRTAl construir INT8un motor modelo, se realizan los siguientes pasos:

1. Construya un 32motor de modelo de bits, luego ejecute este motor en el conjunto de datos de calibración y registre un histograma de la distribución de los valores de activación para cada tensor;
2. Construya una tabla de calibración a partir del histograma y calcule un factor de escala para cada tensor;
3. Construir un motor basado en la tabla de calibración y la definición del modelo INT8.

El proceso de calibración puede ser lento, pero la tabla de calibración generada en el segundo paso se puede generar en un archivo y se puede reutilizar. Si el archivo de la tabla de calibración ya existe, el calibrador leerá la tabla de calibración directamente desde el archivo sin ejecutar el proceso anterior. dos pasos paso. Además, a diferencia de los archivos del motor, las tablas de calibración se pueden utilizar en todas las plataformas. Por lo tanto, durante la implementación real del modelo, primero podemos GPUgenerar la tabla de calibración en una computadora de uso general y luego Jetson Nanousarla en una plataforma integrada. Para facilitar la codificación, podemos utilizar Pythonla programación para implementar INT8el proceso de cuantificación para generar una tabla de calibración.

Implementación

1. Cargar datos de calibración

Primero defina una clase de carga de datos para cargar datos de calibración. Los datos de calibración aquí son una JPGimagen en el formato. Después de leer la imagen, es necesario escalarla, normalizarla, intercambiar canales y otras operaciones de preprocesamiento de acuerdo con los requisitos de datos de entrada del modelo:

class CalibDataLoader:
    def __init__(self, batch_size, width, height, calib_count, calib_images_dir):
        self.index = 0
        self.batch_size = batch_size
        self.width = width
        self.height = height
        self.calib_count = calib_count
        self.image_list = glob.glob(os.path.join(calib_images_dir, "*.jpg"))
        assert (
            len(self.image_list) > self.batch_size * self.calib_count
        ), "{} must contains more than {} images for calibration.".format(
            calib_images_dir, self.batch_size * self.calib_count
        )
        self.calibration_data = np.zeros((self.batch_size, 3, height, width), dtype=np.float32)

    def reset(self):
        self.index = 0

    def next_batch(self):
        if self.index < self.calib_count:
            for i in range(self.batch_size):
                image_path = self.image_list[i + self.index * self.batch_size]
                assert os.path.exists(image_path), "image {} not found!".format(image_path)
                image = cv2.imread(image_path)
                image = Preprocess(image, self.width, self.height)
                self.calibration_data[i] = image
            self.index += 1
            return np.ascontiguousarray(self.calibration_data, dtype=np.float32)
        else:
            return np.array([])

    def __len__(self):
        return self.calib_count

El código de operación de preprocesamiento es el siguiente:

def Preprocess(input_img, width, height):
    img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (width, height)).astype(np.float32)
    img = img / 255.0
    img = np.transpose(img, (2, 0, 1))
    return img

2. Implementar el calibrador

Para implementar la función de un calibrador, debe heredar TensorRTuna de las cuatro clases de calibrador proporcionadas y luego anular varios métodos del calibrador principal:

• get_batch_size: batchtamaño usado para obtener
• get_batch: utilizado para obtener batchdatos de un
• read_calibration_cache: Se utiliza para leer la tabla de calibración desde un archivo.
• write_calibration_cache: Se utiliza para escribir la tabla de calibración desde la memoria a un archivo.

Como es el modelo que necesito cuantificar CNN, elijo heredar IInt8EntropyCalibrator2el calibrador:

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

class Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file=""):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.data_loader = data_loader
        self.d_input = cuda.mem_alloc(self.data_loader.calibration_data.nbytes)
        self.cache_file = cache_file
        data_loader.reset()

    def get_batch_size(self):
        return self.data_loader.batch_size

    def get_batch(self, names):
        batch = self.data_loader.next_batch()
        if not batch.size:
            return None
        # 把校准数据从CPU搬运到GPU中
        cuda.memcpy_htod(self.d_input, batch)

        return [self.d_input]

    def read_calibration_cache(self):
        # 如果校准表文件存在则直接从其中读取校准表
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()

    def write_calibration_cache(self, cache):
        # 如果进行了校准，则把校准表写入文件中以便下次使用
        with open(self.cache_file, "wb") as f:
            f.write(cache)
            f.flush()

3. Generar motor INT8

FP32Presenté el proceso de generación de motores modelo en un artículo anterior, pero en ese artículo se C++implementó. Llamar Pythona la implementación de la interfaz es en realidad más simple: el código específico es el siguiente:

def build_engine():
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    config = builder.create_builder_config()
    parser = trt.OnnxParser(network, TRT_LOGGER)
    assert os.path.exists(onnx_file_path), "The onnx file {} is not found".format(onnx_file_path)
    with open(onnx_file_path, "rb") as model:
        if not parser.parse(model.read()):
            print("Failed to parse the ONNX file.")
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None

    print("Building an engine from file {}, this may take a while...".format(onnx_file_path))

    # build tensorrt engine
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 * (1 << 30))  
    if mode == "INT8":
        config.set_flag(trt.BuilderFlag.INT8)
        calibrator = Calibrator(data_loader, calibration_table_path)
        config.int8_calibrator = calibrator
    else mode == "FP16":
        config.set_flag(trt.BuilderFlag.FP16)

    engine = builder.build_engine(network, config)
    if engine is None:
        print("Failed to create the engine")
        return None
    with open(engine_file_path, "wb") as f:
        f.write(engine.serialize())

    return engine

El código anterior se OnnxParserutiliza primero para analizar el modelo y luego configestablecer la precisión del motor. Si está construyendo INT8un motor, debe establecer la configuración correspondiente Flagy pasarle el objeto calibrador previamente implementado, de modo que los TensorRTdatos de calibración se lean automáticamente para generar una tabla de calibración al construir el motor.

Resultados de la prueba

Para verificar INT8el efecto de la cuantificación, hice una prueba comparativa en YOLOv5varios modelos que utilicé en la tarjeta gráfica. GeForce GTX 1650 TiLos resultados de la prueba de inferencia que requiere mucho tiempo con diferentes precisiones son los siguientes:

Modelo	Introducir dimensiones	Precisión del modelo	Tiempo de razonamiento (ms)
yolov5s.onnx	640x640	INT8	7
yolov5m.onnx	640x640	INT8	10
yolov5l.onnx	640x640	INT8	15
yolov5s.onnx	640x640	FP32	12
yolov5m.onnx	640x640	FP32	23
yolov5l.onnx	640x640	FP32	45

yolov5lLos resultados de detección de objetivos del modelo FP32y INT8la precisión se muestran en las dos imágenes siguientes:

Se puede ver que los resultados de la prueba aún son relativamente cercanos.

Referencias

1. https://developer.nvidia.com/zh-cn/blog/tensorrt-int8-cn/
2. https://developer.nvidia.com/blog/chieving-fp32-accuracy-for-int8-inference-using-quantization-aware-training-with-tensorrt/
3. https://docs.nvidia.com/deeplearning/tensorrt/archives/tensorrt-861/developer-guide/index.html#working-with-int8
4. https://github.com/xuanandsix/Tensorrt-int8-quantization-pipline.git