跳到主要内容

模型编译

模型转换​

SyNAP 工具包允许您将模型从其原始格式编译为针对目标硬件优化的内部表示。转换工具和实用程序可以在 Docker 容器内的 Linux、MacOS 或 Windows 主机上运行。只需要 Docker 和 toolkit 镜像;无需安装其他依赖项。

支持的格式​

  • Tensorflow Lite (.tflite 扩展名)
  • ONNX (.onnx 扩展名)
  • TorchScript (.torchscript、.pt、.pth 扩展名)
  • Tensorflow (.pb 扩展名)
  • Caffe (.prototxt 扩展名)
备注

PyTorch 模型可以保存为不同的格式,仅支持 TorchScript 格式。更多信息请参见使用 PyTorch 模型。

备注

仅支持标准 Caffe 1.0 模型。自定义变体如 Caffe-SSD 或 Caffe-LSTM 模型或旧版(1.0 之前)模型需要特定的解析器,这些解析器目前在 SyNAP 工具包中不可用。也不支持 Caffe2 模型。

注意

对 .pb 和 .prototxt 格式的支持已弃用。

安装 Docker​

下面提供了一些安装提示。请注意,这些不能替代官方的 [Docker] 文档。更多详细信息,请参阅 Docker 网站。

Linux (Ubuntu)​

apt-get install docker.io

为了能够在不使用超级用户权限的情况下运行 Docker,在安装 Docker 后还需要运行以下两个命令(更多信息请参阅 linux postinstall:

# 如果 docker 组不存在则创建
sudo groupadd docker
# 将当前用户 "$USER" 添加到 docker 组
sudo usermod -aG docker $USER

macOS - Docker​

在 macOS 上安装 Docker 最简单的方法是通过 brew 包管理器。如果您还没有安装它,请按照官方 brew 网站的说明进行安装。安装 brew 后,您可以安装 Docker:

important

在 macOS 上,Docker GUI 在商业应用中不是免费使用的。一个有效的替代方案是 Colima。

brew install docker

请参阅上面 Linux 安装中的注释,了解如何在不使用超级用户权限的情况下运行 Docker。

macOS - Colima​

Colima 是 macOS 上的免费容器运行时,可以用作 Docker 的替代品。它没有 GUI,但易于安装和配置:

brew install colima
mkdir -p ~/.docker/cli-plugins
brew install docker-buildx
ln -sfn $(brew --prefix)/opt/docker-buildx/bin/docker-buildx ~/.docker/cli-plugins/docker-buildx
colima start --vm-type vz --mount-type virtiofs --cpu 4 --memory 8 --disk 80

执行上述命令后,您可以使用 [Colima] 来处理 Docker 容器。设置存储在配置文件 ~/.colima/default/colima.yaml 中,如果需要可以通过编辑文件来修改。每次重启 Mac 后都需要启动 Colima:

colima start

Windows​

在 Windows 上运行 Docker 的建议方式是在使用 WSL2 的 Linux 虚拟机中安装它,WSL2 可从 Windows 10 获得。

重要说明:

在 Windows 中直接运行 Docker 与虚拟机的存在不兼容。因此,在 WSL2 中使用 Linux 虚拟机通常是最佳选择。

WSL2 安装步骤​

  1. 以管理员身份运行 Windows PowerShell 应用程序并执行以下命令来安装 WSL2:

    wsl --install

    完成后,重启计算机。

  2. 像之前一样运行 Windows PowerShell 应用程序并安装 Ubuntu-22.04:

    wsl --install -d Ubuntu-22.04
  3. 运行 Windows Terminal 应用程序并选择 Ubuntu-22.04 发行版。从那里,按照上面 using-docker-ubuntu-label 中的说明安装 Docker 和 SyNAP 工具包。

有关 WSL2 安装和设置的更多信息,请参阅官方 Microsoft 文档: WSL 安装 和 WSL 环境。

安装 SyNAP 工具​

在安装 SyNAP 工具包之前,请确保您有一个正常工作的 Docker 安装。最简单的方法是运行 hello-world 镜像:

$ docker run hello-world

Unable to find image 'hello-world:latest' locally
latest: Pulling from libraryhello-world
...
...
Hello from Docker!
This message shows that your installation appears to be working correctly

如果上述命令没有产生预期的输出,请检查上一节中的说明或参阅适用于您平台的官方 Docker 文档。如果一切正常,您可以继续安装工具包。

SyNAP 工具包作为 Docker 镜像分发,只需从 SyNAP github 仓库下载镜像即可安装:

docker pull ghcriosynaptics-synaptoolkit:#SyNAP_Version#

此镜像不仅包含转换工具本身,还包含所有必需的依赖项和额外的支持实用程序。

您可以在这里找到工具包的最新版本。

运行 SyNAP 工具​

一旦安装了 Docker 和 SyNAP 工具包镜像,模型转换工具就可以直接在 Docker 容器内执行。源模型和转换后的模型可以通过在运行容器时挂载相应目录来访问主机文件系统。因此,使用拥有要转换文件所有权的相同用户/组来运行容器很重要。为了避免每次执行时手动指定这些选项,建议创建一个简单的别名并将其添加到用户的启动文件中(例如,.bashrc 或 .zshrc):

alias synap='docker run -i --rm -u $(id -u):$(id -g) -v $HOME:$HOME -w $(pwd) ghcr.io/synaptics-synap/toolkit:3.1.0'

这些选项的含义如下:

  • -i:
    以交互方式运行容器(对于从 stdin 读取数据的命令是必需的,例如 image_od)。

  • --rm:
    容器退出时删除它(停止的容器不再需要)。

  • -u $(id -u):$(id -g):
    以当前用户身份运行容器(这样文件将具有正确的访问权限)。

  • -v $HOME:$HOME:
    挂载用户的主目录,以便其全部内容在容器内可见。如果某些模型或数据位于主目录之外,可以通过重复 -v 选项来挂载额外的目录,例如添加:-v /mnt/data:/mnt/data。重要的是在容器内外指定相同的路径,以便绝对路径按预期工作。

  • -w $(pwd):
    将容器的工作目录设置为当前目录,以便正确解析命令行中指定的相对路径。

使用上述别名,只需将所需的 SyNAP 工具命令行作为参数传递,例如:

$ synap help

SyNAP 工具包

Docker 别名:
alias synap='docker run -i --rm -u $(id -u):$(id -g) -v $HOME:$HOME -w $(pwd) \
ghcr.io/synaptics-synap/toolkit:3.1.0'
如果需要挂载额外目录,请使用多个 -v 选项,例如:-v /mnt/dat:/mnt/dat

用法:
synap COMMAND ARGS
运行 'synap COMMAND --help' 获取有关命令的更多信息。

命令:
convert 转换和编译模型
help 显示帮助
image_from_raw 将图像文件转换为原始格式
image_to_raw 从原始格式生成图像文件
image_od 将目标检测框叠加到图像上
version 显示版本

如前所述,运行 docker 不需要是 root。如果您在执行上述命令时遇到权限拒绝错误,请参阅 Linux (Ubuntu)。

工具包提供了多个工具来转换和操作模型和图像。

可以使用 convert 命令执行模型转换。它需要以下输入:

  • 网络模型
  • 要转换模型的目标硬件(例如,VS680 或 VS640)
  • 生成转换后模型的目录名称
  • 可选的 YAML 元文件,可用于指定自定义转换选项(对于 .pb 模型是必需的)

输出生成三个文件:

  • model.synap:转换后的网络模型
  • model_info.txt:关于生成的模型的额外信息,供用户参考,包括:
    • 输入/输出张量属性
    • 子图分割
    • 层表
    • 操作表
    • 内存使用情况
  • quantization_info.txt:额外的量化信息(仅当模型使用工具包进行量化时)

还会生成一个额外的 cache 目录,以加快将来编译相同模型的速度。

示例:

$ synap convert --model mobilenet_v1_quant.tflite --target VS680 --out-dir mnv1
$ ls mnv1
model_info.txt model.synap cache

对于 Caffe 模型,权重不在 .prototxt 文件中,而是存储在单独的文件中,通常带有 .caffemodel 扩展名。此文件也必须作为输入提供给转换工具。示例:

$ synap convert --model mnist.prototxt --weights mnist.caffemodel --target VS680 --out-dir out
important

模型文件和指定的输出目录必须位于 Docker 容器内挂载的目录中或该目录下(请参阅上面 synap 别名中的 -v 选项)。

预处理​

网络输入张量中要提供的数据的大小、布局、格式和范围是在创建和训练网络模型时定义的。例如,典型的 mobilenet-v1 .tflite 模型将期望输入图像大小为 224x224,具有 NHWC 布局,通道按 RGB 顺序组织,每个像素值归一化(重新缩放)到 [-1, 1] 范围内。

不幸的是,在现实世界中使用时,要处理的图像很少以这种确切格式出现。例如,图像可能来自 1920x1080 YUV 格式的相机。此图像必须转换为 RGB,调整大小并归一化以匹配预期输入。许多库存在以执行这种类型的转换,但问题是这些计算非常计算密集,因此即使深度优化,在 CPU 上执行这些操作通常也需要比推理本身更多的时间。

另一种选择是重新训练网络以接受与运行时可用的数据格式相同的格式。虽然有时这是一个好主意,但它也带来了自己的问题。例如,有时可能无法或不切实际重新训练网络,尤其是对于多个输入大小和格式重复的任务。

为了简化并加快这一任务,SyNAP 工具包允许您在将模型转换时自动插入输入预处理代码。此代码直接在 NPU 中执行,在某些情况下可以比在 CPU 中执行等效操作快一个数量级。另一种将预处理添加到原始模型的方法是为仅将输入图像转换为所需的格式和大小而创建单独的“预处理模型”,然后执行两个模型而无需任何附加数据复制,请参见 Buffer Sharing。如果原始模型很大且输入可以以多种可能格式出现,这可能很方便。最常见的预处理选项已经预先安装。

可用的预处理选项专为图像设计,支持五种变换:

  • 格式转换(例如,YUV 到 RGB 或 RGB 到 BGR)
  • 裁剪
  • 调整大小和缩小(不保留比例)
  • 归一化到所需值范围(例如,将 [0, 255] 归一化到 [-1, 1])
  • 数据类型转换(从 uint8 转换为网络输入层的数据类型,例如 float16 或 int16)

预处理通过在输入规范的 .yaml 文件中指定 preprocess 部分来启用。此部分包含以下字段(标记 (*) 是必需的)。请注意,mean 和 scale 用于量化模型(请参阅输入规范中的 means 和 scale 字段),因此它们不会在这里出现。

type (*)​

此字段指定将提供给网络的输入数据格式。目前仅支持图像格式。SyNAP 工具包将添加所需的操作以将输入数据转换为网络输入张量预期的格式和布局。如果未指定网络输入张量的格式,则默认情况下假设为 rgb。如果此字段设置为空字符串或 "none",则不应用预处理。

并非所有转换都受支持:gray 输入只能用于输入张量具有一个通道的情况。所有其他输入格式(除 float32 外)只能用于输入张量具有三个通道的情况。

某些输入格式为单个网络张量生成多个数据输入。例如,如果指定 nv12,则转换后的网络将有两个输入:第一个用于 y 通道,第二个用于 uv 通道。预处理代码将组合这些两个输入的数据以提供单个 rgb 或 bgr 输入张量给网络。

以下表包含所有支持的输入格式摘要,以及每个生成的输入张量的属性及其含义。请注意,输入数据的布局始终为 NHWC,除了 rgb888-planar 和 float32 格式。在所有情况下,H 和 W 表示输入图像的高度和宽度。如果未明确指定输入图像的大小,则这些从网络输入张量中获取。在所有情况下,每个像素组件用 8 位表示。

float32 类型有点特殊,因为在这个例子中,输入不被认为是 8 位图像,而是原始 32 位浮点值,这些值被转换为张量的实际数据类型。因此,允许任何张量形状,并且通过 size 字段调整大小是不受支持的。

预处理类型输入#形状格式输入描述
yuv4440NHW1y8Y 分量
1NHW1u8U 分量
2NHW1v8V 分量
yuv4200NHW1y8Y 分量
1N(H/2)(W/2)1u8U 分量
2N(H/2)(W/2)1v8V 分量
nv120NHW1y8Y 分量
1N(H/2)(W/2)2uv8UV 组件交错
nv210NHW1y8Y 分量
1N(H/2)(W/2)2vu8VU 组件交错
gray0NHW1y8Y 分量
rgb0NHW3rgbRGB 组件交错
bgra0NHW4bgraBGRA 组件交错
rgb888p0N3HWrgbRGB 组件平面
rgb888p30NHW1r8Red 组件
1NHW1g8Green 组件
2NHW1b8Blue 组件
float320any浮点数据
备注

指定 dummy 预处理(例如从 rgb 输入到 rgb 张量)可以是使用 NPU 硬件而不是在软件中执行相同操作的规范化数据类型转换的方法。

size​

此可选字段允许您指定输入图像的大小作为包含 H 和 W 维度的列表。预处理将输入图像缩放到网络相应输入张量的尺寸。输入图像的比例不会被保留。如果未指定此字段,输入图像的 WxH 维度将与网络张量的 W 和 H 相同。

crop​

启用裁剪。如果指定,四个附加标量输入张量将添加到模型中(它们可以在生成的 model_info.txt 中看到)。这些输入包含单个 32 位整数,并在运行时用于指定裁剪矩形在输入图像中的位置和大小。如果启用了安全功能,这些附加输入将具有安全属性 "any",以便用户应用程序始终可以指定裁剪坐标,即使模型和其他输入/输出张量是安全的。裁剪输入在原始模型输入之后按以下顺序添加:

  • 裁剪矩形的宽度
  • 裁剪矩形的高度
  • 裁剪矩形的左坐标
  • 裁剪矩形的顶部坐标

这些输入应该使用 Tensor 标量 assign() 方法编写,该方法接受以像素为单位的值并将其转换为内部表示。预处理将指定的裁剪矩形缩放到网络相应输入张量的尺寸。输入图像的比例不会被保留。输入图像之外的区域将被忽略。裁剪坐标必须位于输入图像的维度内;否则,结果图像的内容是未定义的。

异构推理​

在某些情况下,将网络的不同部分在不同硬件上执行是有用的。例如,考虑一个对象检测网络,其中初始部分包含一堆卷积层,最后部分是一些后处理层,例如 TFLite_Detection_PostProcess。NPU 针对执行卷积进行了大量优化,但不支持后处理层,因此最佳方法是将网络的初始部分在 NPU 上执行,并将后处理在 CPU 上执行。

这可以通过在每层基础上指定要使用的委托来实现,使用与我们在 Mixed quantization 部分中看到的相同语法。例如,考虑再次在 Figure 4 中,我们可以指定所有层应该在 NPU 上执行,除了 conv5 和 conv6 以及我们希望在 GPU 上执行的层:

# Execute the entire model on the NPU, except conv5 and conv6
delegate:
'*': npu
conv5: gpu
conv5...: gpu

分布处理到不同硬件委托的另一个优点是,当模型组织在多个独立分支中时,可以并行执行分支,每个分支在不同的 HW 单元上执行,然后并行执行分支。

这样,整体推理时间可以减少到执行最慢分支的时间。分支并行化总是自动完成的,只要可能。

备注

分支并行化不应与内部层并行化混淆,内部层并行化总是活动的,只要可能。在上面的示例中,两个分支 (conv3, conv4) 和 (conv5, conv6) 并行执行,前者在 NPU 上执行,后者在 GPU 上执行。此外,每个卷积层在 NPU 和 GPU 硬件中并行化内部化,利用可用并行性。

模型转换教程​

让我们看看如何转换和运行典型的对象检测模型。

  1. 下载样本 ssd_mobilenet_v1_1_default_1.tflite 对象检测模型:

    https://tfhub.dev/tensorflow/lite-model/ssd_mobilenet_v1/1/default/1

  2. 创建转换元文件 ssd_mobilenet.yaml ,内容如下(重要提示:复制粘贴时必须保留换行符和格式,但复制粘贴到 PDF 时会丢失换行符和格式):

outputs:
- name: Squeeze
dequantize: true
format: tflite_detection_boxes y_scale=10 x_scale=10 h_scale=5 w_scale=5 anchors=${ANCHORS}
- name: convert_scores
dequantize: true
format: per_class_confidence class_index_base=-1

关于此文件内容的几点说明:

  • name: Squeeze 和 name: convert_scores 明确指定我们希望模型转换停止的输出张量。最后层 (TFLite_Detection_PostProcess) 是自定义层,不适用于 NPU 加速,因此它在 Detector 后处理器类中实现为软件。
  • dequantize: true 直接在 NPU 中执行从量化到浮点的转换。这比在软件中执行转换要快得多。
  • tflite_detection_boxes 和 convert_scores 表示这些张量中的内容和数据组织。
  • y_scale=10、x_scale=10、h_scale=5、w_scale=5 对应于网络中 TFLite_Detection_PostProcess 层中的参数。
  • ${ANCHORS} 在转换时替换为 TFLite_Detection_PostProcess 层中的 anchor 张量。这是为了能够计算在后处理期间进行后处理的边界框。
  • class_index_base=-1 表示此模型已使用额外的背景类作为索引 0 进行训练,因此我们在后处理期间从类索引中减去 1,以符合标准 coco 数据集标签。
  1. 转换模型(确保模型、元文件和输出目录在容器中可见,请参阅 -v 选项):
$ synap convert --model ssd_mobilenet_v1_1_default_1.tflite --meta ssd_mobilenet.yaml --target VS680 --out-dir compiled
  1. 将模型推送到板:
$ adb root
$ adb remount
$ adb shell mkdir /data/local/tmp/test
$ adb push compiled/model.synap /data/local/tmp/test
  1. 执行模型:
$ adb shell
# cd /data/local/tmp/test
# synap_cli_od -m model.synap $MODELS/object_detection/coco/sample/sample001_640x480.jpg

示例输出:

Input image: /vendor/firmware/.../sample/sample001_640x480.jpg (w = 640, h = 480, c = 3)
Detection time: 5.69 ms
# Score Class Position Size Description
0 0.70 2 395,103 69, 34 car
1 0.68 2 156, 96 71, 43 car
2 0.64 1 195, 26 287,445 bicycle
3 0.64 2 96,102 18, 16 car
4 0.61 2 76,100 16, 17 car
5 0.53 2 471, 22 167,145 car

模型分析​

在开发和优化模型时,了解网络各层执行时间分布情况很有用。这提供了网络哪些层执行效率高以及哪些表示瓶颈的指示。

要获得此信息,网络必须逐步执行,以便可以测量每个单独的时间。为此,网络必须通过调用 synap_convert.py 并使用 --profiling 选项来生成,例如:

$ synap convert --model mobilenet_v2_1.0_224_quant.tflite --target VS680 --profiling --out-dir mobilenet_profiling
备注

即使每层执行时间在 normal 和 profiling 模式下没有变化,编译具有启用分析功能的网络也将明显比没有分析功能的相同网络编译时间长,因为 NPU 执行必须多次启动和暂停以收集分析数据。因此,分析应该正常禁用,仅在需要调试时启用。

当使用 SyNAP 工具包转换模型时,层可以融合、替换为等效操作,并/或优化掉。因此,通常不可能找到分析信息中的项目与原始网络中的节点之间的一对一对应关系。例如,相邻卷积、ReLU 和池化层在可能时融合为单个 ConvolutionReluPoolingLayer 层。尽管进行了这些优化,对应关系通常不是太难找到。分析中显示的层对应于在模型转换时生成的 model_info.txt 文件中列出的层。

每次执行编译的模型时,分析信息将在 sysfs 中可用,请参见 Statistics and Usage 。由于此信息不是持久性的,当网络被销毁时它会消失,因此收集它的最简单方法是通过 synap_cli 程序。--profiling [filename](filename) 选项允许在网络被销毁之前将 sysfs network_profile 文件的副本保存到指定文件中:

$ adb push mobilenet_profiling $MODELS/image_classification/imagenet/model/
$ adb shell
# cd $MODELS/image_classification/imagenet/model/mobilenet_profiling
# synap_cli -m model.synap --profiling mobilenet_profiling.txt random

# cat mobilenet_profiling.txt
pid: 21756, nid: 1, inference_count: 78, inference_time: 272430, inference_last: 3108, iobuf_count: 2, iobuf_size: 151529, layers: 34
| lyr | cycle | time_us | byte_rd | byte_wr | type
| 0 | 152005 | 202 | 151344 | 0 | TensorTranspose
| 1 | 181703 | 460 | 6912 | 0 | ConvolutionReluPoolingLayer2
| 2 | 9319 | 51 | 1392 | 0 | ConvolutionReluPoolingLayer2
| 3 | 17426 | 51 | 1904 | 0 | ConvolutionReluPoolingLayer2
| 4 | 19701 | 51 | 1904 | 0 | ConvolutionReluPoolingLayer2
...
| 28 | 16157 | 52 | 7472 | 0 | ConvolutionReluPoolingLayer2
| 29 | 114557 | 410 | 110480 | 0 | FullyConnectedReluLayer
| 30 | 137091 | 201 | 2864 | 1024 | Softmax2Layer
| 31 | 0 | 0 | 0 | 0 | ConvolutionReluPoolingLayer2
| 32 | 0 | 0 | 0 | 0 | ConvolutionReluPoolingLayer2
| 33 | 670 | 52 | 1008 | 0 | ConvolutionReluPoolingLayer2

SyNAP 3.x 与 SyNAP 2.x 的兼容性​

SyNAP 3.x 完全向后兼容 SyNAP 2.x。

  • 可以执行使用 SyNAP 3.x 工具包编译的模型,使用 SyNAP 2.x 运行时。唯一限制是,在这种情况下,异构编译不可用,整个模型将在 NPU 上执行。这可以通过在转换模型时指定 --out-format nb 选项来完成。在这种情况下,工具包将生成旧版 model.nb 和 model.json 文件,而不是 model.synap 文件:

    $ synap convert --model mobilenet_v2_1.0_224_quant.tflite --target VS680 --out-format nb --out-dir mobilenet_legacy
  • 可以执行使用 SyNAP 2.x 工具包编译的模型,使用 SyNAP 3.x 运行时。

  • SyNAP 3.x API 是 SyNAP 2.x API 的扩展,因此所有现有应用程序可以无修改使用。

使用 PyTorch 模型​

PyTorch 框架支持非常灵活的模型,其中网络的架构和行为是使用 Python 类而不是固定图层(例如,在 TFLite 中)定义的。当保存模型时,通常只保存 state_dict,它包括可学习参数,而不是模型结构本身。(Saving and Loading Models in PyTorch)。

需要重新加载和执行它,因此无法直接从仅包含可学习参数的 .pt 文件导入 PyTorch 模型。

保存包括模型结构的 PyTorch 模型​

当保存 PyTorch 模型到 .pt 文件时,可以包括引用定义模型的 Python 类的文件。但是,重新创建模型仍然需要访问用于生成它的确切 Python 源代码。

保存 TorchScript 格式​

更便携的替代方法是保存模型到 TorchScript 格式。此格式包括可学习参数和模型结构,允许直接导入到其他工具中,例如 SyNAP 工具包。

有关如何保存到 TorchScript 格式的更多信息,请参阅 official PyTorch documentation。

跟踪与脚本​

保存到 TorchScript 格式可以执行脚本或跟踪:

  • 脚本:将模型转换为 torch.jit.ScriptModule 通过分析其结构。
  • 跟踪:记录运行模型时执行的操作,对于具有动态结构的模型很有用。

这两种技术都会产生相同的格式,并且可以使用跟踪保存的模型也可以直接导入。在线搜索“PyTorch 跟踪与脚本比较”以获取详细比较。

保存 TorchScript 模型的示例代码​

import torch
import torchvision

# An instance of your model
model = torchvision.models.mobilenet_v2(pretrained=True)

# Switch the model to eval mode
model.eval()

# Generate a torch.jit.ScriptModule via scripting
mobilenet_scripted = torch.jit.script(model)
mobilenet_scripted.save("mobilenet_scripted.torchscript")

# Generate a torch.jit.ScriptModule via tracing
example = torch.rand(1, 3, 224, 224)
mobilenet_traced = torch.jit.trace(model, example)
mobilenet_traced.save("mobilenet_traced.torchscript")
important

即使存在多个可能保存 PyTorch 模型到文件的方法,也没有 同意用于不同情况的扩展名,并且 .pt 或 .pth 扩展名通常用于任何文件格式,仅 TorchScript 模型可以导入到 SyNAP 工具包中, 如果模型是不同格式,导入将失败并出现错误消息。

备注

使用 TorchScript 模型进行混合量化或 异构推理时不太方便,因为模型层有时没有名称或名称在 导入过程中被修改和/或没有一对一对应关系 原始模型和导入的模型。在这种情况下,建议编译模型 使用 --preserve 选项,然后查看输出目录中的中间 build/model.onnx 文件。

一个更便携的替代方法是将模型导出到 ONNX 格式。 所需的代码非常相似,用于跟踪模型:

import torch
import torchvision

# An instance of your model
model = torchvision.models.mobilenet_v2(pretrained=True)

# Switch the model to eval model
model.eval()

# Export the model in ONNX format
torch.onnx.export(model, torch.rand(1, 3, 224, 224), "mobilenet.onnx")

导入 YOLO PyTorch 模型​

流行的 YOLO 库从 ultralytics 提供其网站上的预训练 .pt 模型。 所有这些模型都不是 TorchScript 格式,因此不能直接导入到 SyNAP 工具包中。 然而,非常容易将它们导出到 ONNX 或 TorchScript 格式,以便可以导入:

from ultralytics import YOLO

# Load an official YOLO model
model = YOLO("yolov8s.pt")

# Export the model in TorchScript format
model.export(format="torchscript", imgsz=(480, 640))

# Export the model in ONNX format
model.export(format="onnx", imgsz=(480, 640))

更多信息在 https://docs.ultralytics.com/modes/export/

大多数公共领域机器学习包提供类似的导出函数用于他们的 PyTorch 模型。