1. Экспорт модели в формат ONNX

Общее

opset - важный параметр. Официально китайский SDK для компиляции onnx моделей поддерживает все операции из opset=12. Но нужно пробовать с версиями выше, поскольку в них есть более производительные операции, что сказывается на итоговой скорости работы модели

https://onnxruntime.ai/docs/reference/compatibility.html

https://github.com/onnx/onnx/blob/main/docs/Operators.md

Ultralytics

У моделей ultralytics есть свой интерфейс для конвертации в onnx: https://docs.ultralytics.com/modes/export/

Можно указывать параметр batch: int, хотя о нем не сказано в документации

Пример:

from ultralytics import YOLO
model = YOLO("yolov8n")
model.export(
format="onnx",
imgsz=640,
opset=12,
batch=4)

PyTorch

Если есть pytorch модель, ее можно экспортировать с помощью torch.onnx.export

Пример:

import torch
import torchvision.models as models
BATCH_SIZE = 1
ONNX_FILE = "output.onnx"
model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.DEFAULT)
input_shape = (BATCH_SIZE, 3, 224, 224) # Define the input shape of your model
input_names = ["input"] # Specify the input names
output_names = ["output"] # Specify the output names
torch.onnx.export(
model,
torch.randn(input_shape),
ONNX_FILE,
input_names=input_names,
output_names=output_names,
opset_version=12)

2. Компиляция модели в формат TPU

То, что описано в этой секции будет во многом дублировать официальную документацию, поэтому не помешает изучить оригинал тоже.

Окружение

Для начала нам понадобится:

• Docker образ sophgo/tpuc_dev:v2.1
  docker pull sophgo/tpuc_dev:v2.1
• SDK - https://sophon-file.sophon.cn/sophon-prod-s3/drive/23/08/17/11/Release_230701-public.zip
  Все версии можно найти тут https://developer.sophon.ai/site/index/material/all/all.html
• Из SDK нам понадобится распакованный tpu_nntc

workspace # рабочая папка, название может быть любым
└── tpu_nntc
├── bin
├── doc
├── include
├── kernel
├── lib
├── scripts
├── userlayer
└── wheel

Теперь запустим докер контейнер в котором и будет происходить компиляция. Для этого:

1. Находясь в рабочей папке (PWD=workspace) нужно запустить:
docker run -v $PWD:/workspace -p 8001:8001 -it sophgo/tpuc_dev:v2.1
Теперь мы находимся внутри контейнера.

2. Далее нужно проинициализировать окружение:

cd tpu_nntc
source scripts/envsetup.sh

3. Все готово, можно переходить к компиляции.

Компиляция в FP32

В зависимости от модели и размера батча, компиляция может потребовать значительное время или количество оперативной памяти.

Для компиляции нужно вызвать bmneto.compile, c правильными параметрами. Вызов в python скрипте. Пример такого вызова:

import torch
BATCH_SIZE = 1
ONNX_FILE = "input-model.onnx"
WORKDIR = "<рабочая папка компиляции>"

bmneto.compile(
model=ONNX_FILE,
outdir=WORKDIR,
target="BM1684",
input_names=["input"], # здесь должно быть название, которое было задано на шаге экспорта
shapes=[[BATCH_SIZE ,3,224,224]],
net_name="model", # имя модели которое будет сохранено в результирующем файле
dyn=False, # динамическая модель имеет меньшую производительность, поэтому тут False
opt=2 # уровень оптимизации, целое значение от 0 до 2)

💡 Результирующая модель будет лежать в {WORKDIR}/compilation.bmodel