Оптимизация модели FastAI2 на NVIDIA Jetson Orin с помощью TensorRT и Torch2TRT: неправильный размер батча.

Вопрос или проблема

У меня есть Jetson Orin с последней версией Jetpack 6.0 и CUDA 12, работающей на Ubuntu 22.04.

Я установил PyTorch, и он имеет поддержку CUDA:

Python 3.10.12 (main, 11 сентября 2024, 15:47:36) [GCC 11.4.0] на linux
Введите "help", "copyright", "credits" или "license" для получения дополнительной информации.
>>> import torch
>>> 
>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
>>> print(f"Используемое устройство: {device}")
Используемое устройство: cuda

Я использую TensorRT для оптимизации моей модели Fastai 2.0:

>>> import torch2trt
>>> import tensorrt
>>> print(tensorrt.__version__)
8.6.2

Я оптимизирую модель следующим образом:

# Параметры
detect_batch_size = 8  # Установите размер партии на 8
global_win_h = 200     # Пример высоты
global_win_w = 200     # Пример ширины

# Создайте входной тензор с размером партии 8
x = torch.ones((detect_batch_size, 3, global_win_h, global_win_w)).cuda()

# Создайте пользовательский логгер для подавления предупреждений
class CustomLogger(trt.Logger):
    def __init__(self):
        super().__init__(trt.Logger.ERROR)  # Установите в ERROR для подавления предупреждений

# Создайте экземпляр пользовательского логгера
custom_logger = CustomLogger()

def load_cnn_model(logger, should_use_tensor_rt=True):
    '''
    Эта функция загружает CNN модель и возвращает её

    :param: logger - объект для логирования
    :param: should_use_tensor_rt - булевое значение, указывающее, нужно ли использовать NVIDIA TensorRT
    :return: кортеж - модель обучающегося и классы для распознавания
    '''

    # Загрузка модуля CNN
    model_load_time_start = time.time()

    # Используйте CUDA
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    logger.info(f"Нужно использовать CUDA: {use_cuda}")
    default_device(use=use_cuda)

    # Загрузка обучающегося
    logger.info(f"Загрузка модели на платформе: {platform.system()} ...")
    platform_designator = platform.system()
    if platform_designator == 'Linux':
        logger.info(f"Платформа: {platform_designator}, изменяем путь к Windows ...")
        pathlib.WindowsPath = pathlib.PosixPath
    learner = load_learner("JetsonNanoFaceTrackingTrainAndTest/fast_ai_trained_models/face.pkl", cpu=False)
    classes = learner.dls.vocab
    logger.info(f"Классы: {classes} ...")
    learner.model.to(device)
    learner.model.eval()

    # Оптимизация для TensorRT
    logger.info(f"Проверка использования TensorRT: {should_use_tensor_rt} ...")
    if should_use_tensor_rt:

        # Проверка, существует ли уже оптимизированная модель
        opt_model = Path(tensor_rt_optimized_model)
        logger.info(f"Проверка оптимизированной модели: {tensor_rt_optimized_model} ...")
        if not opt_model.exists():

            logger.info("Оптимизированная модель не существует ...")

            # Преобразование модели в половинную точность
            #learner.model = learner.model.half()

            learner.dls.to(device="cuda")
            learner.model.to(device="cuda")

            try:

                logger.info("Оптимизация модели для TensorRT ...")
                learner = torch2trt(learner.model, [x], logger=custom_logger) # Используйте пользовательский логгер, чтобы отключить предупреждения
                logger.info("Оптимизация модели завершена ...")

                logger.info(f"Сохранение оптимизированной модели: {tensor_rt_optimized_model} ...")
                torch.save(learner.state_dict(), tensor_rt_optimized_model)

            except Exception as err:
                logger.error(f"Ошибка: Не удалось создать оптимизированную модель TensorRT ...")
                sys.exit(1)

        else:

            logger.info(f"Загрузка существующей оптимизированной модели TensorRT: {tensor_rt_optimized_model} ...")
            learner = TRTModule()
            learner.load_state_dict(torch.load(tensor_rt_optimized_model))

    # Измерение времени загрузки
    model_load_time_end = time.time()
    logger.info(f"Время загрузки модели: {model_load_time_end - model_load_time_start}")

    # Возвращаем модель и ожидаемые для распознавания классы

Я могу успешно загрузить оптимизированную модель, а затем пытаюсь сделать вывод, используя следующий код на моем наборе данных с пакетной инференцией:

# Создание пользовательского набора данных
infer_dataset = InferenceDataset(image_list_to_infer, augmentations=preprocess)

logger.info(f"Создание загрузчика данных: {len(image_list_to_infer)} | {infer_dataset} | {len(infer_dataset)} ...")

# Создание DataLoader для пакетной обработки
infer_loader = torch.utils.data.DataLoader(infer_dataset,
                                            detect_batch_size,
                                            shuffle=True,
                                            num_workers=1,
                                            pin_memory=True)

logger.info(f"Размер загрузчика для вывода: {infer_loader}")

# Выполнение пакетного вывода
torch.cuda.synchronize()
start_inference = time.time()
with torch.no_grad():

    # Переведите модель в режим вывода
    accumulated_output = []
    logger.debug("Установка режима вывода ...")
    infer_counter = 0
    with torch.inference_mode():
        for data in infer_loader:
            data = data.cuda()
            logger.debug(f"Вызов обучающегося: {data} ...")
            logger.debug(f"Обучающийся: {learner}")

            # Определение класса
            #output = learner.model(data)

            infer_counter += 1
            logger.info(f"Счетчик вывода: {infer_counter}")

            # TensorRT требует этот синтаксис
            output = learner(data)

            logger.debug(f"Завершен вызов обучающегося: {output} ...")
            for tensor_output in output.data.cpu():
                accumulated_output.append(tensor_output)

# Время вывода
torch.cuda.synchronize()
end_inference = time.time()
logger.info(f"Время вывода: {end_inference - start_inference} | [{len(accumulated_output)}]")

Вывод завершается неудачей, и я получаю следующее сообщение в консоли:

[19/10/2024-15:22:12] [TRT] [E] 3: [executionContext.cpp::validateInputBindings::2103] Код ошибки 3: Ошибка использования API (Проверка параметров не прошла на: runtime/api/executionContext.cpp::validateInputBindings::2103, условие: profileMaxDims.d[i] >= dimensions.d[i]. Поставленный размер привязки [8,3,200,200] для bindings[0] превышает минимальный ~ максимальный диапазон по индексу 0, максимальный размер в профиле равен 1, минимальный размер в профиле равен 1, но указанный размер равен 8

Мне не удается понять, как решить эту проблему. Я считаю, что при первоначальном создании оптимизированной модели я установил размер партии на ‘8’. Тем не менее, я все еще вижу эту проблему.

Как я могу это исправить? Спасибо.

Ответ или решение

Оптимизация модели FastAI2 на NVIDIA Jetson Orin с использованием TensorRT и Torch2TRT: Решение проблемы с размерам пакета

Ваша проблема связана с несовпадением размерности входных данных, передаваемых в модель TensorRT, и размерности, для которой модель была оптимизирована. Эта ошибка возникает при нарушении ограничений, установленных для входных данных во время оптимизации модели с помощью TensorRT. Давайте разберем, почему это происходит и как вы можете это исправить.

Проблема с размером пакета

Ваша ошибка указывает на то, что оптимизированная модель ожидает входные данные с размерностью (1, 3, 200, 200), в то время как вы передаете данные с размерностью (8, 3, 200, 200). Это происходит потому, что TensorRT создает профили для входных данных, определяя минимальные и максимальные размеры для каждой размерности. В вашем случае минимальная и максимальная размерности для пакета были установлены на 1, и когда вы пытаетесь передать данные с размером пакета 8, возникает ошибка.

Решение проблемы

Чтобы исправить данную проблему, вам нужно убедиться, что TensorRT оптимизируется с учетом необходимого размера пакета. Вот шаги, которые помогут вам это сделать:

  1. Настройка профиля во время оптимизации: При оптимизации модели с помощью torch2trt вам нужно задать параметры входного тензора, включая минимальный и максимальный размеры. Вы можете сделать это следующим образом:

    from torch2trt import torch2trt, TrtModule
    
    # Определите минимальный и максимальный размер пакета
    min_batch_size = 1
    max_batch_size = 8
    
    # Создайте входной тензор для оптимизации с размером `max_batch_size`
    x = torch.ones((max_batch_size, 3, global_win_h, global_win_w)).cuda()
    
    # Создайте диапазон для размеров пакета
    input_shapes = [(min_batch_size, 3, global_win_h, global_win_w), (max_batch_size, 3, global_win_h, global_win_w)]
    
    # Оптимизируйте модель с указанными размерами входного тензора
    learner = torch2trt(learner.model, [x], logger=custom_logger, inputs=input_shapes)
  2. Оптимизация для нескольких размеров пакета: Вы можете также оптимизировать модель для нескольких размеров пакета, добавив дополнительные профили, например, для размеров пакетов 2, 4 и 8.

  3. Изменение размера пакета на этапе инференса: Если вы оптимизируете модель для фиксированного размера пакета, будьте уверены, что при инференсе вы используете тот же размер. Если вы хотите использовать более крупный размер пакета для инференса, оптимизируйте модель на этот размер заранее.

  4. Пересоздание оптимизированной модели: Если ваши данные не соответствуют минимальным и максимальным профилям, вы можете пересоздать оптимизированную модель, чтобы она соответствовала данным, которые вы собираетесь использовать. Убедитесь, что вы обновляете все соответствующие параметры.

Пример кода для исправления

Убедитесь, что ваш код для оптимизации модели выглядит следующим образом с учетом минимального и максимального размеров:

# Оптимизация модели
if should_use_tensor_rt:
    x = torch.ones((detect_batch_size, 3, global_win_h, global_win_w)).cuda()
    try:
        # Оптимизируем модель с учетом профиля
        logger.info("Optimizing model for TensorRT ...")
        learner = torch2trt(learner.model, [x], logger=custom_logger, max_batch_size=max_batch_size)
        logger.info("Done optimizing model ...")

        # Сохраняем оптимизированную модель
        logger.info(f"Saving optimized model: {tensor_rt_optimized_model} ...")
        torch.save(learner.state_dict(), tensor_rt_optimized_model)

    except Exception as err:
        logger.error(f"Error: Failed to create optimized TensorRT model ...")
        sys.exit(1)

Итоги и рекомендации

Следуя этим шагам, вы сможете избежать проблемы неправильного размера пакета при инференсе вашей модели, оптимизированной с использованием TensorRT. Убедитесь, что вы правильно определяете минимальные и максимальные размеры для входных тензоров во время настройки TensorRT. Это гарантирует, что ваша модель сможет обрабатывать входящие данные с любым подходящим размером пакета в рамках установленных вами ограничений.

Если у вас будут дополнительные вопросы или проблемы, пожалуйста, не стесняйтесь обращаться за поддержкой.

Оцените материал
Добавить комментарий

Капча загружается...