Rockchip — довольно крупная китайская компания, которая разрабатывает микросхемы для тв-приставок, смартфонов и планшетов. По данным на май 2025 года, годовой доход компании Rockchip находится в диапазоне от 200 млн до 500 млн долларов. Мне довелось поработать с одноплатным компьютером Orange Pi 5, оснащенным ARM процессором RK3588 и NPU (neural processing unit) с заявленной производительностью до 6 TOPS. У них есть подробная документация, примеры конвертации известных архитектур в формат rknn, например resnet and yolo. Есть отдельный репозиторий с моделями rknn_model_zoo и rknn-llm.
Однако примера для собственной произвольной архитектуры нет.

В этой статье я хочу поделиться своим опытом по конвертации нейросети в формат rknn с помощью библиотеки rknn-toolkit2. Анализ процесса преобразования поможет лучше разобраться с тем, как работает эта платформа. Это полезно как для учебных целей, так и для исследований. В этот пример легко добавить произвольный слой pytorch и посмотреть, как он преобразуется в инструкции для NPU сопроцессора rockchip.

git clone https://github.com/labintsev/rknn-toolkit2.git --depth 1
conda create -n toolkit2 python=3.11
conda activate toolkit2
pip install rknn-toolkit2 -i https://pypi.org/simple
cd rknn-toolkit2/examples/onnx/custom
python test.py

Описание пайплайна

1. Создание pytorch модели

Модель OneLayerModel — это очень простая нейросеть на PyTorch.
Модель подходит для задач классификации небольших изображений и демонстрирует базовый пайплайн преобразования PyTorch → ONNX → RKNN.

class OneLayerModel(nn.Module):
    def __init__(self):
        super(OneLayerModel, self).__init__()
        self.flatten = nn.Flatten()
        self.layer = nn.Linear(3072, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.layer(x)
        return x

Она состоит из двух основных слоёв:

1. Flatten — преобразует входной тензор размерности (1, 3, 32, 32) (1 изображение, 3 канала, 32x32 пикселя) в одномерный вектор длины 3072.

2. Linear — полносвязный слой, который принимает этот вектор и преобразует его в выход из 10 классов.

Вот как выглядят веса pytorch модели в Netron:

Важно!
Все модели, представленные в репозитории 'rknn-toolkit2', имеют формат 'NCHW' (channel first). За размерность входных данных отвечает параметр input_size_list в методе rknn.load_onnx. Однако установка input_size_list=[[1, 32, 32, 3]] не позволяет использовать этот формат входных данных. К сожалению, я так и не нашел параметр, который позволяет задать формат 'NHWC' (channel last) при построении модели. Если использовать модель в формате 'NHWC', то на этапе построения модели получаем ошибку:

ValueError: The channel of r_shape [1, 32, 32, 3] must be 3!

2. Конвертация pytorch модели в формат onnx

Функция build_pytorch_model создает экземпляр модели OneLayerModel и сохраняет её веса в файл с расширением .pt (формат PyTorch). Функция build_onnx_model создает экземпляр модели, загружает сохранённые веса из файла .pt, переводит модель в режим eval, экспортирует её в формат ONNX с помощью torch.onnx.export, и сохраняет результат в файл .onnx. Таким образом, эти функции реализуют последовательное сохранение модели сначала в формате PyTorch, затем в формате ONNX для дальнейшего использования на платформе Rockchip.
Вот так выглядит модель в формате onnx:

где
Reshape - операция преобразования размерности тензора.
Gemm - базовая операция матричного умножения.

Внимание!
Без параметра dynamo=True экспорт в onnx сломается с ошибкой:

ValueError: Export destination must be specified for torchscript-onnx export.

Вопросы к библиотеке onnx оставим для дальнейшего изучения, но пока работает - оставим как есть.

3. Конвертация onnx модели в rknn

Далее создается объект RKNN, который управляет процессом конвертации и инференса модели на платформе Rockchip. Выполняется конфигурация параметров для будущей модели, например, указывается целевая платформа (в данном случае 'rk3588'). Загружается ONNX-модель в объект RKNN для дальнейшей обработки и конвертации. На этом этапе происходит подготавка модели к конвертации в формат RKNN и последующему запуску на NPU Rockchip.

4. Создание и экспорт rknn модели

На этом этапе происходит конвертация ONNX-модели во внутренний формат RKNN, оптимизация графа и подготовка к запуску на NPU Rockchip. Если включена квантизация (do_quantization=True), веса и входные данные модели преобразуются в целочисленный формат для ускорения инференса. Файл dataset.txt содержит список путей к изображениям или numpy-массивам, которые используются для сбора статистики при квантизации. Эти данные позволяют корректно подобрать параметры преобразования float → int8, чтобы минимизировать потери точности при работе модели на NPU. Если квантизация не используется, этот файл можно не указывать.

Вот так выглядят промежуточные представления модели onnx:

Экспорт модели является завершающим этапом процесса конвертации нейронной сети в формат, оптимизированный для работы на устройствах Rockchip. Создается файл модели с расширением .rknn, в котором хранятся структура и веса модели. Модель оптимизирована под целевую платформу (в данном случае RK3588) и сохраняется в файл, указанный в переменной RKNN_MODEL.

Вот так выглядит модель в формате rknn:

5. Симуляция runtime

На этапе подготовки runtime настраиваются все необходимые компоненты для последующего запуска модели на целевом устройстве. Во время инициализации происходит загрузка оптимизированной RKNN-модели в память устройства и настройка всех системных ресурсов, требуемых для работы нейронной сети. Система проверяет доступность необходимых библиотек, выделяет память и настраивает параметры выполнения. Наш пример запускается на платформе x86, поэтому важно заметить, что в библиотеке rknn-toolkit2 реализован симулятор NPU RockChip. Он используется для запуска модели и анализа производительности инференса. Успешная инициализация подтверждает готовность системы.

После успешной инициализации среды выполнения начинается непосредственная работа с моделью. На этапе инференса происходит подача входных данных в подготовленную модель и получение результатов обработки. В нашем случае входные данные передаются в формате numpy-массива с указанием формата данных NCHW (batch, channels, height, width).

Процесс инференса включает в себя следующие этапы:

• Модель получает на вход подготовленные данные

• Выполняется последовательная обработка данных через все слои нейронной сети

• Система производит необходимые вычисления на NPU-ускорителе

• Формируется выходной результат, который сохраняется в переменной outputs

После завершения инференса модель готова к обработке следующей порции входных данных до тех пор, пока мы не вызовем метод rknn.release().

Анализ логов

В логах можно увидеть интересную таблицу с содержимым вычислительного графа rknn и вычислителем (CPU/NPU) на котором выполняется каждый узел:

D RKNN: [20:41:40.895] --------------------------------------------------------------------------------------------------------------------------------------------------------------------------
D RKNN: [20:41:40.895]                                                         Network Layer Information Table                                                      
D RKNN: [20:41:40.895] --------------------------------------------------------------------------------------------------------------------------------------------------------------------------
D RKNN: [20:41:40.895] ID   OpType             DataType Target InputShape                               OutputShape            Cycles(DDR/NPU/Total)    RW(KB)       FullName        
D RKNN: [20:41:40.895] --------------------------------------------------------------------------------------------------------------------------------------------------------------------------
D RKNN: [20:41:40.895] 0    InputOperator      INT8     CPU    \                                        (1,3,32,32)            0/0/0                    0            InputOperator:l_x_
D RKNN: [20:41:40.895] 1    Conv               INT8     NPU    (1,3,32,32),(3,3,1,1),(3)                (1,3,32,32)            831/4096/4096            3                            
D RKNN: [20:41:40.895] 2    Reshape            INT8     NPU    (1,3,32,32),(4)                          (1,3072,1,1)           0/0/0                    16           Reshape:flatten_1_tp_rs
D RKNN: [20:41:40.895] 3    Conv               INT8     NPU    (1,3072,1,1),(10,3072,1,1),(10)          (1,10,1,1)             1441/1536/1536           33           Conv:Gemm_1#2   
D RKNN: [20:41:40.895] 4    Reshape            INT8     CPU    (1,10,1,1),(2)                           (1,10)                 0/0/0                    0            Reshape:layer_1_mm_tp_rs
D RKNN: [20:41:40.895] 5    OutputOperator     INT8     CPU    (1,10)                                   \                      0/0/0                    0            OutputOperator:layer_1
D RKNN: [20:41:40.895] --------------------------------------------------------------------------------------------------------------------------------------------------------------------------

Таблица включает следующую ключевую информацию для каждого слоя:

Layer ID — уникальный идентификатор для каждой операции

Operation Type (OpType) — определяет тип выполняемой операции (например, Conv, Reshape)

Data Type — указывает используемый формат (в данном случае INT8)

Target Processor — показывает, выполняется ли операция на CPU или NPU

Input/Output Shapes — размеры тензоров данных

Cycles(DDR/NPU/Total) - показывает количество циклов выполнения для разных компонентов:

• DDR - циклы доступа к памяти

• NPU - циклы выполнения на нейропроцессоре

• Total - общее количество циклов

RW(KB) - объем данных, прочитанных/записанных в килобайтах во время выполнения операции

FullName - полное имя операции в вычислительном графе, включая:

• Тип операции (например, Conv, Reshape)

• Уникальный идентификатор

• Дополнительные параметры

Также в консоли можно прочитать подробности процесса преобразования модели. Большую часть логов занимают сообщения о ходе построения модели. Можно погрузиться в детали оптимизации, например, fuse ops:

D     convert_gemm_by_exmatmul: remove node = ['Gemm_1'], add node = ['flatten_1_tp', 'flatten_1_tp_rs', 'Gemm_1#1', 'layer_1_mm_tp', 'layer_1_mm_tp_rs']
D     unsqueeze_to_4d_transpose: remove node = [], add node = ['flatten_1_rs', 'flatten_1_tp-rs']
D     fuse_two_reshape: remove node = ['Reshape_3']
D     fuse_transpose_reshape: remove node = ['flatten_1_tp', 'layer_1_mm_tp']
D     fuse_two_reshape: remove node = ['flatten_1_rs']
D     fuse_two_reshape: remove node = ['flatten_1_tp-rs']
D     convert_exmatmul_to_conv: remove node = ['Gemm_1#1'], add node = ['Gemm_1#2']

Также выдается предупреждение о том, что тип данных для входа и выхода модели изменен с float32 на int8:

W build: The default input dtype of 'l_x_' is changed from 'float32' to 'int8' in rknn model for performance!
                       Please take care of this change when deploy rknn model with Runtime API!
W build: The default output dtype of 'layer_1' is changed from 'float32' to 'int8' in rknn model for performance!
                      Please take care of this change when deploy rknn model with Runtime API!

Заключение

Платформа Rockchip - это мощный инструмент для работы с нейронными сетями. Она позволяет запускать модели на специальном процессоре, который потребляет значительно меньше электроэнергии. С учетом накладных расходов и слабый ARM CPU, на 20 Вт одноплатном компьютере получается производительность уровня Ryzen 5 2600 (сравнивал с ним).

Особенно ценно то, что разработчики могут использовать не только готовые архитектуры, но и создавать собственные модели. Процесс преобразования модели из PyTorch в формат RKNN оказался довольно прозрачным. Сначала модель сохраняется в промежуточный формат ONNX, а затем преобразуется в формат, понятный процессору Rockchip. При этом происходит несколько оптимизаций. Главный прирост дает квантизация модели — веса и данные переводятся из формата с плавающей точкой в целочисленный, что делает вычисления быстрее и менее энергозатратными.

Таким образом, библиотека rknn-toolkit позволяет использовать как готовые решения, так и модели для уникальных задач. Это открывает интересные перспективы для разработчиков, которые хотят оптимизировать работу своих нейронных сетей.