Triton Inference Server推理引擎

https://github.com/triton-inference-server/server

Triton 从入门到精通

https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/index.html

深度学习部署神器——triton-inference-server入门教程指北

triton-inference-server的backend(一)——关于推理框架的一些讨论

Triton Inference Server是Nvida开源的机器学习推理引擎(可以理解为同TF Serving对等的产品),其提供了多种开箱即用的功能帮助我们快速落地AI模型到生产环境以提供业务使用。当我们团队人手资源受限或开发时间不足的情况下,没有能力自研一套机器学习推理引擎,Triton是一个非常好的选择。

从Triton的不断完善的进程来,其已经逐渐的具有了一个深度学习推理引擎应该所具有的的全部功能(当然还是有一些不够完善的地方)。

Triton Inference Server核心的几个功能:

  • 多框架支持:Triton支持了几乎所有主流的机器学习框架,例如Tensorflow、TensorRT、Pytorch、Python、ONNX等;同时也可以custom backend的方式来扩展解码引擎。
  • 高性能:Triton提供了dynamic batching、concurrent execution、optimal model configuration、model ensemble、dali model 等策略来提升在线推理的性能;同时也提供了perf analyze和model analyze工具来辅助我们进行性能调优。
  • MLOps:Triton提供了Prometheus exporter、模型在线更新、http server 、grpc server等多种在线服务策略以满足用户生产场景多样化的部署和运维需求

triton可以充当服务框架去部署你的深度学习模型,其他用户可以通过http或者grpc去请求你搭建的服务,相当于你用flask搭了个服务供别人请求,当然相比flask的性能高很多了。triton也可以摘出C-API充当多线程推理服务框架,去除http和grpc部分,适合本地部署多模型,比如你有很多模型要部署,然后分时段调用,或者有pipeline,有了triton就省去你处理显存、内存和线程的麻烦。

借用官方的图,triton的使用场景结构如下:

涉及到运维部分,我也不是很懂,抛去K8S后,结构清爽了些:

triton的一些优点

通过上述的两个结构图,可以大概知道triton的一些功能和特点

  • 支持HTTP/GRPC
  • 支持多backend,TensorRT、libtorch、onnx、paddle、tvm啥的都支持,也可以自己custom,所以理论上所有backend都可以支持
  • 单GPU、多GPU都可以支持,CPU也支持
  • 模型可以在CPU层面并行执行
  • 很多基本的服务框架的功能都有,模型管理比如热加载、模型版本切换、动态batch,类似于之前的tensorflow server
  • 开源,可以自定义修改,很多问题可以直接issue,官方回复及时
  • NVIDIA官方出品,对NVIDIA系列GPU比较友好,也是大厂购买NVIDIA云服务器推荐使用的框架
  • 很多公司都在用triton,真的很多,不管是互联网大厂还是NVIDIA的竞品都在用,用户多代表啥不用我多说了吧
大厂都在用triton

如何学习triton

两年前开始学习的时候,官方资料比较匮乏, 只能通过看源码来熟悉triton的使用方式,所幸知乎上有个关于TensorRT serving不错的教程[2],跟着看了几篇大致了解了triton的框架结构。那会triton叫做TensorRT serving,专门针对TensorRT设计的服务器框架,后来才变为triton,支持其他推理后端的。

现在triton的教程比较多了,官方的docs写着比较详细,还有issue中各种用例可以参考,B站上也有视频教程[3],比两年前的生态要好了不少。

当然,最重要的,还是上手使用,然后看源码, 然后客制化。

源码学习

从triton的源码中可以学到:

  • C++各种高级语法
  • 设计模式
  • 不同backend(libtorch、TensorRT、onnxruntime等)如何正确创建推理端,如何多线程推理
  • C++多线程编程/互斥/队列
  • API接口暴露/SDK设计
  • CMAKE高级用法

等等等等,不列举了,对于程序员来说,好的源码就是好的学习资料。当然,也可以看老潘的文章哈。

triton系列教程计划

triton相关系列也会写一些文章,目前大概规划是这些:

  • 什么是triton以及triton入门、triton编译、triton运行
  • triton管理模型、调度模型的方式
  • triton的backend介绍、自定义backend
  • 自定义客户端,python和c++
  • 高级特性、优先级、rate limiter等等

编译和安装

一般来说,如果想快速使用triton,直接使用官方的镜像[4]最快。但是官方镜像有个尴尬点,那就是编译好的镜像需要的环境一般都是最新的,和你的不一定一致

比如22.09版本的镜像需要的显卡驱动为520及以上,如果想满足自己的显卡驱动,就需要自行编译了。官方也提供了使用镜像的快速使用方法

# 第一步,创建 model repository 
git clone -b r22.09 https://github.com/triton-inference-server/server.git
cd server/docs/examples
./fetch_models.sh

# 第二步,从 NGC Triton container 中拉取最新的镜像并启动
docker run --gpus=1 --rm --net=host -v ${PWD}/model_repository:/models nvcr.io/nvidia/tritonserver:22.09-py3 tritonserver --model-repository=/models

# 第三步,发送
# In a separate console, launch the image_client example from the NGC Triton SDK container
docker run -it --rm --net=host nvcr.io/nvidia/tritonserver:22.09-py3-sdk
/workspace/install/bin/image_client -m densenet_onnx -c 3 -s INCEPTION /workspace/images/mug.jpg

# Inference should return the following
Image '/workspace/images/mug.jpg':
    15.346230 (504) = COFFEE MUG
    13.224326 (968) = CUP
    10.422965 (505) = COFFEEPOT

triton官方仓库

两年前的triton只有一个大仓库,tensorrt_backend也默认在triton主仓库中,但是现在tensorrt_backend被拆分出来了,很显然triton除了支持tensorrt还支持很多其他的后端,我们可以自定义使用很多后端。

现在是目前的triton包含的一些仓库:

  • [server[5]] triton服务外层框架,包含了http收发请求,服务内存分配等一些功能代码
  • [core[6]] triton主框架,如果处理请求、后端管理、模型调度啥的全在这里
  • [common[7]] 通用工具,没啥好说的,打日志的代码在这里
  • [backend[8]] backend后端框架代码,存放了一些后端通用父类,自定义后端可以集成这些类仿写新的后端
  • [third_party[9]] triton使用的第三方库的汇总,主要是cmake里头会包含
  • [tensorrt_backend[10]] tensorrt后端代码
  • [pytorch_backend[11]] libtorch后端代码

最开始的时候,server、core、common、backend这些代码仓库都是合在一起的,后来都拆分出来了,增加了triton的灵活性。

比如,上述的core仓库可以单独暴露出cAPI作为动态链接库供其他程序调用,去掉http、grpc的外层请求接口,直接一步到位调用。

一般来说,我们都是从最主要的server开始编,编译的时候会链接core、common、backend中的代码,其他自定义backend(比如tensorrt_backend)在编译的时候也需要带上common、core、backend这三个仓库,这些关系我们可以从相应的CMakeList中找到。

自行编译

如果想要研究源码,修改源码实现客制化,那么自行编译是必须的。

triton的编译和安装其实很简单,唯一的难点就是需要加速,因为triton在编译的时候会clone很多第三方库,第三方库也会克隆它们需要的第三方库,这些库当然都是国外的,所以有个好的网络环境很重要。

比如在编译triton的时候需要下载grpc这个库,grpc又依赖很多第三方其他库,网络不好的话,会经常遇到下面的问题:

Failed to recurse into submodule path 'third_party/bloaty'
CMake Error at /tmp/tritonbuild/tritonserver/build/_deps/repo-third-party-build/grpc-repo/tmp/grpc-repo-gitclone.cmake:52 (message):
  Failed to update submodules in:
  '/tmp/tritonbuild/tritonserver/build/_deps/repo-third-party-build/grpc-repo/src/grpc'


make[3]: *** [_deps/repo-third-party-build/CMakeFiles/grpc-repo.dir/build.make:99: _deps/repo-third-party-build/grpc-repo/src/grpc-repo-stamp/grpc-repo-download] Error 1
make[3]: Leaving directory '/tmp/tritonbuild/tritonserver/build'
make[2]: *** [CMakeFiles/Makefile2:590: _deps/repo-third-party-build/CMakeFiles/grpc-repo.dir/all] Error 2
make[2]: Leaving directory '/tmp/tritonbuild/tritonserver/build'
make[1]: *** [CMakeFiles/Makefile2:145: CMakeFiles/server.dir/rule] Error 2
make[1]: Leaving directory '/tmp/tritonbuild/tritonserver/build'

开加速是最好的办法,不管是UI还是命令行,都有相应的软件可以用,比如clash。

如果你的服务器实在是开不了加速,也有其他办法,那就是将triton库中大部分重量级库的git地址全换为国内的

怎么替换,我是在gitee中,同步github上的仓库,比如triton的core仓库,同步过来,就可以使用国内的地址了。

当然也需要将这些库的submodule中的库也修改为国内源,比如grpc这个库依赖很多第三方库,克隆的时候,这是要一个一个下载的:

改起来稍微麻烦,还需要注意,要改特定commit分支的git地址:

ver/build/_deps/repo-third-party-build/grpc-repo/src/grpc/third_part目录然后手动git clone xxx,然后执行一下git submodule init / git submodule update下就可以带进去。

示例:

root@64da25af2629:/tmp/tritonbuild/tritonserver/build/_deps/repo-third-party-build/grpc-repo/src/grpc# git submodule init
root@64da25af2629:/tmp/tritonbuild/tritonserver/build/_deps/repo-third-party-build/grpc-repo/src/grpc# git submodule update
Submodule path 'third_party/googletest': checked out 'c9ccac7cb7345901884aabf5d1a786cfa6e2f397'

太麻烦了,不过确实为一种办法呃呃。

还有一点,triton每次build都会clone,是因为其用了cmake中的ExternalProject_Add指令,假如我们已经有下载好的grpc,那么直接替换到server/build/_deps/repo-third-party-build/grpc-repo/src中然后将/data/oldpan/software/server/build/_deps/repo-third-party-src/CMakeLists.txt

注释掉git下载部分,修改自己本地的就行,就不需要每次再clone一遍了。

#
# Get the protobuf and grpc source used for the GRPC endpoint. We must
# use v1.25.0 because later GRPC has significant performance
# regressions (e.g. resnet50 bs128).
#
ExternalProject_Add(grpc-repo
  PREFIX grpc-repo
  # GIT_REPOSITORY "https://gitee.com/Oldpann/grpc.git"
  # GIT_TAG "v1.25.x"
  SOURCE_DIR "${CMAKE_CURRENT_BINARY_DIR}/grpc-repo/src/grpc"
  CONFIGURE_COMMAND ""
  BUILD_COMMAND ""
  INSTALL_COMMAND ""
  TEST_COMMAND ""
  PATCH_COMMAND python3 ${CMAKE_CURRENT_SOURCE_DIR}/tools/install_src.py --src <SOURCE_DIR> ${INSTALL_SRC_DEST_ARG} --dest-basename=grpc_1.25.0
)

说了这么多,总之,最好的办法当然还是开科学,全局一下就OK,省去那么多麻烦事儿。

搞定好网络问题,编译triton就很简单了!

git clone --recursive https://github.com/triton-inference-server/server.git
cd server
python build.py  --enable-logging --enable-stats --enable-tracing --enable-gpu  --endpoint=http --repo-tag=common:r22.06 --repo-tag=core:r22.06 --repo-tag=backend:r22.06 --repo-tag=thirdparty:r22.06 --backend=ensemble --backend=tensorrt

在克隆好的server的目录下执行以上命令(下面是我的设置,我们可以个根据自己的需求进行修改)就可以了。

执行这个命令后triton就会构建docker在docker中编译,最终会创建3个镜像:

  • tritonserver:latest
  • tritonserver_buildbase:latest
  • tritonserver_cibase:latest

最终编译好的tritonserver_buildbase:latest镜像,我们可以在其中开发,因为环境都帮忙配好了,只需要再执行编译命令,就可以编译了,我们也可以自定义源码进行个性功能的开发。

在镜像中开发

需要注意,在编译的时候需要pull官方默认的镜像,而这个镜像是有显卡驱动限制的,比如r22.06需要显卡驱动版本为470。

同志们看看自己的显卡驱动,别下了不能用hhh

可以通过triton镜像历史[12]查看镜像版本要求:

接上,我们不是编译好了triton镜像,直接进去就可以开发了:

docker run -v/home/oldpan/code:/code -v/home/oldpan/software:/software  -d tritonserver_buildbase:latest /usr/bin/sh -c "while true; do echo hello world; sleep 20;done"

在docker中修改triton的源码,继续执行以下命令就可以编译,和之前的区别就是加了 --no-container-build参数。

python build.py  --enable-logging --enable-stats --enable-tracing --enable-gpu  --endpoint=http --repo-tag=common:r22.06 --repo-tag=core:r22.06 --repo-tag=backend:r22.06 --repo-tag=thirdparty:r22.06 --backend=ensemble --no-container-build --build-dir=./build

我们如果想编译debug版本的triton,可以在命令中添加:--build-type=Debug

另外,原始triton镜像中已经有tensorrt,如果想换版本,可以删除原始docker中的旧的tensorrt,自行安装新的tensorrt即可:

  • https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html[13]

说下运行流程吧!

讲了这么多铺垫,接下来简单说下运行流程。

这里通过代码简单梳理下triton运行的整体流程,之后的具体细节,放到接下来的篇章讲解

首先一开始,main函数在servers/main.cc下,triton在启动的时候会执行以下函数:

// src/servers/main.cc 经过简化
int
main(int argc, char** argv)
{
  // 解析参数
  TRITONSERVER_ServerOptions* server_options = nullptr;
  if (!Parse(&server_options, argc, argv)) {
    exit(1);
  }
  ...
  // 这里创建server
  TRITONSERVER_Server* server_ptr = nullptr;
  FAIL_IF_ERR(
      TRITONSERVER_ServerNew(&server_ptr, server_options), "creating server"); // 这里创建server
  FAIL_IF_ERR(
      TRITONSERVER_ServerOptionsDelete(server_options),
      "deleting server options");
  std::shared_ptr<TRITONSERVER_Server> server(
      server_ptr, TRITONSERVER_ServerDelete);
  ...
  // 启动HTTP, GRPC, 以及性能统计的端口
  if (!StartEndpoints(server, trace_manager, shm_manager)) {
    exit(1);
  }
  // Trap SIGINT and SIGTERM to allow server to exit gracefully
  signal(SIGINT, SignalHandler);
  signal(SIGTERM, SignalHandler);
  // 等待kill信号区关闭triton 
  while (!exiting_) {
   ...
      // 做一些监控模型仓库是否变动的操作
  }
  // 优雅地关闭triton
  TRITONSERVER_Error* stop_err = TRITONSERVER_ServerStop(server_ptr);
  // 如果无法优雅地关掉,旧直接exit即可
  if (stop_err != nullptr) {
    LOG_TRITONSERVER_ERROR(stop_err, "failed to stop server");
    exit(1);
  }
  // 停止监控http、grpc
  StopEndpoints();
  ...
  return 0;
}

TRITONSERVER_ServerNew这个函数中,会:

  • new一个triton类InferenceServer对象
  • 根据参数设置配置一下,执行一堆Set函数
  • 配置好参数后,Init服务,这里初始化服务的状态,校验参数
  • 创建各种模块,经常使用的有后端管理TritonBackendManager以及模型仓库管理ModelRepositoryManager
  • 再进行一些检查、配置一些状态

在启动过程中最重要的是模型仓库,运行triton当然你要有模型,要不然你开它干嘛?

这里我使用的模型仓库目录结构如下(是一个识别姿态的hrnet,hrnet官方有很多预训练模型,转tensorrt也很简单):

debug目录下有一个模型文件夹叫做hrnet-pose-estimate-debug的模型文件夹,这个文件夹地址(/path/to/hrnet-pose-estimate-debug)需要传给triton启动命令行,文件夹内的四个子模型文件夹,会被triton检测到并且一一加载。

需要注意的是,除了hrnet_pose_estimate这个其余三个在目录的1子目录下有个so或者model.plan,这代表hrnet-trt-staticimage_preprocess还有pose_postprocess都属于model,使用了backend,backend会在各自的config中指明:

name: "hrnet-trt-static"
backend: "tensorrt"

因为hrnet-trt-static是tensorrt的模型,所以backend设置为tensorrt,model.plan就是tensorrt的engine。其backend的so文件我放到了其他位置(放到和model.plan同目录也是可以的),而另外两个预处理和后处理的backend就放到了模型仓库中,也就是libtorch_image_preprocess.solibtriton_pose_postprocess,包含了你的backend代码,封装成so供triton调用

关于backend、model以及modelinstanc的关系,说实话稍微复杂点,各自有完整的生命周期,这个嘛,之后文章说,感兴趣的也可以提前看官方文档的介绍:

  • https://github.com/triton-inference-server/backend[14]

然后我们就启动triton吧!

# 执行以下函数,模型目录通过 --model-repository 指定     tensorrt的backend通过  --backend-directory 指定
./tritonserver --model-repository=/path/to/hrnet-pose-estimate-debug --backend-directory=/workspace/backends/tensorrt_backend/ 

模型加载成功之后会输出:

...
I1016 08:25:37.952055 51771 server.cc:587] 
+------------------+----------------------------------------------------------------+----------------------------------------------------------------+
| Backend          | Path                                                           | Config                                                         |
+------------------+----------------------------------------------------------------+----------------------------------------------------------------+
| image_preprocess | /workspace/triton-models/debug/hrnet-pose-estimate-debug/image | {"cmdline":{"auto-complete-config":"false","min-compute-capabi |
|                  | _preprocess/1/libtriton_image_preprocess.so                    | lity":"6.000000","backend-directory":"/workspace/backends/tens |
|                  |                                                                | orrt_backend/","default-max-batch-size":"4"}}     |
|                  |                                                                |                                                                |
| pose_postprocess | /workspace/triton-models/debug/hrnet-pose-estimate-debug/pose_ | {"cmdline":{"auto-complete-config":"false","min-compute-capabi |
|                  | postprocess/1/libtriton_pose_postprocess.so                    | lity":"6.000000","backend-directory":"/workspace/backends/tens |
|                  |                                                                | orrt_backend/","default-max-batch-size":"4"}}     |
|                  |                                                                |                                                                |
| tensorrt         | /workspace/backends/tensorrt_backend/li              | {"cmdline":{"auto-complete-config":"false","min-compute-capabi |
|                  | btriton_tensorrt.so                                            | lity":"6.000000","backend-directory":"/workspace/backends/tens |
|                  |                                                                | orrt_backend/","default-max-batch-size":"4"}}     |
|                  |                                                                |                                                                |
+------------------+----------------------------------------------------------------+----------------------------------------------------------------+

I1016 08:25:37.952252 51771 server.cc:630] 
+---------------------+---------+--------+
| Model               | Version | Status |
+---------------------+---------+--------+
| hrnet-trt-static    | 1       | READY  |
| hrnet_pose_estimate | 1       | READY  |
| image_preprocess    | 1       | READY  |
| pose_postprocess    | 1       | READY  |
+---------------------+---------+--------+

I1016 08:25:38.051742 51771 metrics.cc:650] Collecting metrics for GPU 0: NVIDIA GeForce RTX 3080
I1016 08:25:38.055197 51771 tritonserver.cc:2159] 
+----------------------------------+------------------------------------------------------------------------------------------------------------------+
| Option                           | Value                                                                                                            |
+----------------------------------+------------------------------------------------------------------------------------------------------------------+
| server_id                        | triton                                                                                                           |
| server_version                   | 2.23.0                                                                                                           |
| server_extensions                | classification sequence model_repository model_repository(unload_dependents) schedule_policy model_configuration |
|                                  |  system_shared_memory cuda_shared_memory binary_tensor_data statistics trace                                     |
| model_repository_path[0]         | /workspace/triton-models/debug/hrnet-pose-estimate-debug                                                         |
| model_control_mode               | MODE_NONE                                                                                                        |
| strict_model_config              | 1                                                                                                                |
| rate_limit                       | OFF                                                                                                              |
| pinned_memory_pool_byte_size     | 268435456                                                                                                        |
| cuda_memory_pool_byte_size{0}    | 300021772                                                                                                        |
| response_cache_byte_size         | 0                                                                                                                |
| min_supported_compute_capability | 6.0                                                                                                              |
| strict_readiness                 | 1                                                                                                                |
| exit_timeout                     | 30                                                                                                               |
+----------------------------------+------------------------------------------------------------------------------------------------------------------+

I1016 08:25:38.055627 51771 http_server.cc:3303] Started HTTPService at 0.0.0.0:8000
I1016 08:25:38.097213 51771 http_server.cc:178] Started Metrics Service at 0.0.0.0:8001

加载好之后,我们开启了http端口,端口号为8000,另一个是metric接口,端口号8001

此时可以使用http请求试一下。

简单请求

请求的话有http和grpc协议,我对http协议熟悉些,所以就搞http吧。

官方也提供了客户端,C++和python的都可以有,可以直接使用官方的,也可以根据官方提供的http协议构造自己的客户端,只要会构造body,一切都很简单。

请求协议可以参考官方:

  • https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md[15]

这里我们用python简单构造一个body

# 构造triton的输入body
json_buf = b'{\"inputs\":[{\"name\":\"INPUT\",\"datatype\":\"BYTES\",\"shape\":[1],\"parameters\":{\"binary_data_size\":' + \
        bytes(str(len(data)), encoding = "utf8") + b'}}],\"outputs\":[{\"name\":\"RESULT\",\"parameters\":{\"binary_data\":true}}]}'
push_data = json_buf + data

print("Inference-Header-Content-Length ",str(len(json_buf)), " Content-Length ",str(len(data) + len(json_buf)))
# 构造triton-header
header = {"Content-Type": "application/octet-stream", "Accept": "*/*",
          "Inference-Header-Content-Length":str(len(json_buf)),
          "Content-Length":str(len(data) + len(json_buf))}

server_url = "127.0.0.1:8000"
model_name = "hrnet_pose_estimate"

# 请求
response = post('http://' + server_url + '/v2/models/' + model_name + '/infer', data=push_data, headers=header)

就可以发送请求,结果也会传回response里。我们也可以使用curl命令,直接传递构造好的body(这个body将上述的push_data写到本地即可):

[oldpan@have-fun client]$ curl -v --max-time 1 --request POST 'http://192.168.1.102:9006/v2/models/hrnet_pose_estimate/infer' --header 'Inference-Header-Content-Length: 230' --header 'Content-Type: application/octet-stream' --data-binary '@data.txt' --output temp_res
Note: Unnecessary use of -X or --request, POST is already inferred.
*   Trying 192.168.1.102:9006...
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0* Connected to 172.29.210.105 (172.29.210.105) port 9006 (#0)
> POST /v2/models/aocr_cnprint_trt8p/infer HTTP/1.1
> Host: 192.168.1.102:9006
> User-Agent: curl/7.71.1
> Accept: */*
> Inference-Header-Content-Length: 230
> Content-Type: application/octet-stream
> Content-Length: 1573102
> Expect: 100-continue

* Mark bundle as not supporting multiuse
< HTTP/1.1 100 Continue
} [56480 bytes data]
* We are completely uploaded and fine
* Mark bundle as not supporting multiuse
< HTTP/1.1 200 OK
< Content-Type: application/json
< Inference-Header-Content-Length: 394
< Content-Length: 6794

{ [6794 bytes data]
100 1542k  100  6794  100 1536k   127k  28.8M --:--:-- --:--:-- --:--:-- 28.9M

结果就不发了,验证没啥问题。

关于如何使用curl直接请求triton,有一些相关链接可以参考:

  • https://github.com/triton-inference-server/server/issues/2563[16]
  • https://github.com/triton-inference-server/server/issues/1822[17]
参考资料
https://www.bilibili.com/video/BV1KS4y1v7zd/?spm_id_from=333.788&vd_source=eec038509607175d58cdfe2e824e8ba2[18]
https://github.com/triton-inference-server/server/releases[19]
https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/index.html[20]
https://zhuanlan.zhihu.com/p/462817679[21]
参考资料
[1]
triton: https://github.com/openai/triton

[2]
教程: https://www.zhihu.com/people/pwlazy/posts

[3]
视频教程: https://www.bilibili.com/video/BV1KS4y1v7zd/?spm_id_from=333.337.search-card.all.click&vd_source=eec038509607175d58cdfe2e824e8ba2

[4]
官方的镜像: https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/index.html

[5]
[server: https://github.com/triton-inference-server/server

[6]
[core: https://github.com/triton-inference-server/core

[7]
[common: https://github.com/triton-inference-server/common

[8]
[backend: https://github.com/triton-inference-server/backend

[9]
[third_party: https://github.com/triton-inference-server/third_party.git

[10]
[tensorrt_backend: https://github.com/triton-inference-server/tensorrt_backend

[11]
[pytorch_backend: https://github.com/triton-inference-server/pytorch_backend

[12]
triton镜像历史: https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/rel_22-06.html

[13]
https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html: https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html

[14]
https://github.com/triton-inference-server/backend: https://github.com/triton-inference-server/backend

[15]
https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md: https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md

[16]
https://github.com/triton-inference-server/server/issues/2563: https://github.com/triton-inference-server/server/issues/2563

[17]
https://github.com/triton-inference-server/server/issues/1822: https://github.com/triton-inference-server/server/issues/1822

[18]
https://www.bilibili.com/video/BV1KS4y1v7zd/?spm_id_from=333.788&vd_source=eec038509607175d58cdfe2e824e8ba2: https://www.bilibili.com/video/BV1KS4y1v7zd/?spm_id_from=333.788&vd_source=eec038509607175d58cdfe2e824e8ba2

[19]
https://github.com/triton-inference-server/server/releases: https://github.com/triton-inference-server/server/releases

[20]
https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/index.html: https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/index.html

[21]
https://zhuanlan.zhihu.com/p/462817679: https://zhuanlan.zhihu.com/p/462817679

讨论:如何用1024张显卡训练一个模型?

文章来源于一个知乎问题:如何判断候选人有没有千卡GPU集群的训练经验?确实对于普通开发者来说,大概率从未尝试过使用数千张GPU训练一个模型,这方面确实是一个很好的研究方向,也是成为顶尖算法工程师所必需的必经之路,因此记录下知乎的一些回答,用于学习和记录。虽然目前还没有机会能够调用数千张GPU用于模型训练,但对于目前几十张GPU进行并行训练也有帮助。

高赞回答1:如何用千卡进行训练

最近看到知乎一个回答,把千卡训练的难度吹上天了。但其实真正用过千卡就会发现也就那么几个点。于是想写一篇文章简单讲讲。

本文将包括三个部分:

  • 首先我们将讨论千卡训练的难题,以及应该在什么时候使用千卡训练;
  • 接着,我们将讨论如何在一千张卡上开始训练,如何让他达到近乎线性的性能提升;
  • 最后我们将展开讨论一些千卡训练当中仍然悬而未决(至少对于开源社区来说)的问题。

为什么千卡训练是困难的?

其实那篇回答在这部分说的没错。千卡训练和八卡训练的区别是—显卡多了一百多倍。

这意味着什么呢?

  1. 通信时间增加
  2. 故障概率增加

这俩问题都很好理解。

时间上,PyTorch 内部支持 NCCL / Gloo / MPI 三个通信后端(请务必使用 NCCL。)其中训网络最常用的 AllReduce 操作【从多个sender那里接收数据,最终combine到一个节点上】会根据具体硬件配置走 Ring AllReduce 和 Tree AllReducering allreduce和tree allreduce的具体区别是什么?。Ring 的时间复杂度是 O(pn),Tree 的时间复杂度是 O(log⁡pn)。就算是理论上 128 节点也比单节点慢至少七倍,实践当中跨节点通信要远比单节点慢得多。

故障上,一个节点出问题的概率是 p,128 个节点就是 1−(1−p128)。也就是说如果一个操作在一个训练当中的出错概率是 1%,那么在 128 节点当中的出错概率就是 72.37%。

此外,随着规模的增大,许多问题都会变得难以忍受。比如数据增强要花 0.1s,一亿条数据就是 278 个小时(当然这只是胡拆的一个数字,实际有各种机制所以不会有这么大影响。

因此,钱多烧手并不是使用千卡训练的理由。闲得蛋疼可能是,但你得多蛋疼才能想出这么折磨自己的 idea?

因此,千卡训练解决的问题是大模型&大数据问题如果你的训练时间没有超过 8192 GPU 日,那么你绝对不需要一千张显卡。

看到这里,绝大多数人已经可以关掉这篇文章了。除非你的模型和数据都以 B(十亿)来作为计量单位。当然如果你正在厕所里手机没电想看点儿东西解闷儿的话(虽然我很怀疑是否会有人把他打出来……那么可以继续往下看

如何使用一千张卡训练?

如何提高计算效率?

这件事情其实是一个 case by case 的事情。因为通信、计算速度啥的受硬件影响更多。同样是 A100 集群,我全 DGX 节点,每一张 A100 都是 SXM 接口并配一块儿专属的 IB 网卡。你一个小破普惠服务器插 8 张 PCI-E A100,IB 卡一个节点只给一张。那咱俩遇到的问题就完全不是一个问题。

因此,要讨论如何提高训练效率、减少训练耗时,我们首先要了解训练耗时在哪里。那么,一个训练步的耗时在哪里呢?需要谨记,没有 profile 的优化是没有意义的。

你可能会说,forward backward sync。很好,这说明你了解 PyTorch 的基本流程。不过现实当中要复杂得多。

  1. dataset 读取数据,构建输出
  2. dataloader collate 数据,进行数据预处理
  3. 模型 forward 计算输出
  4. loss compute
  5. 模型 backward 计算梯度
  6. 模型 sync 梯度
  7. 优化器 step 更新权重
  8. 打印 log

当然这是可以无限细分下去的,但一般这些就够了。需要注意的是,除了 4-7 的耗时是真耗时,其他都需要通过异步操作来盖掉。这也是我们的优化目标。

异步执行在 PyTorch 的 dataloader、CUDA 和分布式当中都存在。前者可以通过设置 num_workers 和 prefetch_count 为 0 来关闭,后两者可以通过 cuda.synchornize 和 dist.barrier 来执行手动同步。在 profile 时,我们需要首先需要测整个 step 的时长。然后再在每次测量前执行手动同步来计算每个部分的时长。如果前者的总耗时等于后者 4-7 的耗时之和,那么通常不需要执行任何操作。但这种情况在千卡操作中几乎不可能发生。

第 6 步通信往往需要耗费大量时间。因此,我们还需要进一步优化通信。

以下内容是对PyTorch Distributed的概括,有感兴趣的同学建议通读并背诵全文。

计算-通信重叠

在 PyTorch 当中,梯度的通信和反向传播是交叠进行的。也就是说,每完成一层的梯度计算,都会立即触发当前层的同步。实现起来也很简单,每个进程在完成自己第 k 层的梯度计算后都会触发一个钩子来给计数器+1s。当计数器达到进程数时开火进行梯度通信。有很多同学在计算梯度过程中遇到过 RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. 错误,这就是因为有的模块没有参与计算 loss,导致梯度同步卡住了。需要注意,当 find_unused_parameters=True 时,PyTorch 分布式使用 nn.Module.__init__ 当中定义子模块的反向顺序来作为梯度桶的构建顺序。因此,确保模块定义和调用的顺序一致是一个良好的实践。

梯度合桶

尽管理论上来说,同步发生的越及时,重合度越高,性能越好。但实际上每次发起通信都是有上头的。因此,现实当中梯度同步并不是越多越好越快越好。为此,PyTorch 引入了梯度合桶机制,通过把多个 Tensor 装在一个桶里再通信桶来减少通信次数从而减少总耗时。合桶的 bucket_cap_mb 默认是 25MiB,这对于绝大多数模型来说都是太小的。目前已经有提升这个默认值的特性需求,但是这个还是调一下更好。

梯度累加

当你做完所有操作之后,惊喜的发现 TMD 怎么同步时间还是单节点的好几倍。这其实是正常情况……实际上超过 256 卡的训练想要把通信盖掉就是一件不可能的事情。你说老师我看 FB 论文说他们 256 卡就是线性提升啊…那这里不得不提的一个策略就是梯度累加了。梯度累加会执行 k 次 forward+backward 之后再执行优化器步进。这有很多好处,首先对于大模型 batch size 通常不能开多大,梯度累加可以提升等效 batch size。其次累加期间的 backward 不需要通信梯度,加快了训练速度。

少即是快

Python 是一种很慢的语言。当然你说 JIT trace+torch.compile 有提升我也不反对,但对于最高效率来说,只有必须要存在的代码和不存在的代码两种。

抱抱脸的 Transformers 就是一个反例。两个子模块就能写完的 TransformerLayer 他们硬是能写出来一堆…偏偏他们还信奉 Single Model File Policy……我寻思你这完全不考虑继承的封这么多层是要搞鸡毛啊?正例反而是 PyTorch……(笑死,我竟然会夸脸书代码写得好。具体来说就是 nn.functional 当中的各种实现。你会发现他们第一行往往是 handle_torch_func。熟悉 Python 装饰器的小伙汁通常要问了,为啥这里不用个装饰器统一一下?因为装饰器会引入额外的函数调用,额外的函数调用就是额外的上头。

因此,如果你想确保最高的效率,写一个简单的训练代码和模型代码非常重要。毕竟,1%的效率提升,节省的可能是数百个 GPU 日。

如何平稳训练

这一段当中中咱们只讨论你能控制的问题。

捕捉不致命的异常

故障率高的问题其实很好解决。在训练当中,大部分异常都是非致命异常,接住他们就好了。我之前写过一个装饰器,catch,它的作用就是接住异常,然后调回调函数(默认当然就是把错误打印到 log 里)。所有你需要做的只是使用它来装饰所有非 fatal 的操作。

在实际应用当中,我们遇到的最常见的问题是存 ckpt 写满了磁盘(不准笑,从商汤到上海 AI Lab,这个问题在哪儿都日常出现。咱也不知道为啥肯买那么多显卡但不肯多插点儿硬盘,咱也不敢问)。接住所有保存操作,如果你有闲心可以在回调里删一下之前的 ckpt。没闲心的话…大不了重训一次嘛(逃。)第二常见的问题,你猜对了……存 log 写满了硬盘……所以所有 logging 操作也都是要 catch 的。这就是为啥我都用 tmux 然后开很长的缓存窗口,总是能抢救一些 log 出来的。

咳咳,说点儿正经的。任何联网操作都是需要 catch 的,常见的联网操作主要包括从 ceph 读取数据和…写 log 到远程(逃。其他就没啥了吧,我见过有大哥尝试恢复 OOM 的,但效果似乎不是很好,至少我自己没用过。简单来说,唯一不应捕捉的错误是集群炸了。

那有的大兄弟就说了,集群没爆炸,但是有两张卡突然掉了咋办。这个咱第三部分再讨论。

管好模型的输出

模型训着训着发散了几乎是每个训大模型的人都会遇到的问题。输出和 loss 只要有 nan 果断丢掉。梯度先 clip by value 再 clip by norm 都是常规操作。哦对了,还有初始化……关于大模型收敛性的论文有一堆,此处不再赘述。

比更大,还更大,再更大

弹性训练

实际上当你的训练超过 2048 个 GPU 日时,在整个训练过程当中发生单个 GPU 甚至单个节点下线是再正常不过的事情了。

PyTorch 在 1.10 就引入了 torchelastic 弹性训练机制。这个东西,用过的都骂娘。等下,让我先骂一遍。呸。ok 咱们继续吧。

我印象当中在微软的最后一轮面试当中被问到了这个问题:如何设计一个弹性分布式系统。

我的回答很教科书。每 k 分钟,系统会做一次 AllGather 来统计存活进程数,然后选举出一个主进程。主进程会计算好每个进程的 rank 和 local rank 然后广播给各个进程。所有进程每次前向传播开始时向主进程发送一个心跳包来汇报状态。主进程会根据心跳包来确定这一个 step 参与同步的机器有多少。

但很可惜,2024 年了。还是没人去写。他妈的。

层次化梯度同步

我一直认为梯度同步不应该以 GPU/进程为单位。而应该分为大同步(节点间同步)和小同步(节点内同步)。小同步可以更高频的进行,大同步则可以更慢的执行。这样不仅能提高实际的梯度同步频率,降低同步总耗时,并且还能天然的去结合小 batch 和大 batch 训练的优点—节点内小 batch 关注个体,节点间大 batch 关注整体。

有没有发现所有东西都很简单?
是这样的,千卡训练是任何一个普通CS本科生花三个月就能学会的东西。没有任何复杂的地方。

延伸阅读

PyTorch DDP 设计笔记

PyTorch 微调菜谱

分析与优化使用 PyTorch 训练机器学习模型

使用 Nsight Systems 来分析 GPU 负载

DLProf

NVProf

NCCL AllReduce 设计

Spike No More

Understanding the difficulty of training deep feedforward neural networks

高赞回答2:关于千卡训练的难度分析

千卡,其实是个相对模糊的概念,对于多数人而言,这就跟你告诉你,我有1兆资产和10兆资产一样,你只知道很多很多,但是完全傻傻分不清楚,这到底是多少,能买多少东西。千卡,也是一样。按照常见的一机8卡GPU的类型来看,用过125台机器训练的,就算得上是千卡了。从这个角度上说,其实门槛也没有那么那么的高。可事实呢?真正的大模型要用多少机器训练?答案是远超千卡 —— 看看下面的GPT-4信息,这是千卡吗?比万卡还要多!!!

How much compute was used to train GPT-4?
2.15e25 floating

Some key facts about how this enormous model was trained: Used 25,000 Nvidia A100 GPUs simultaneously. Trained continuously for 90–100 days. Total compute required was 2.15e25 floating point operations.Sep 27, 2023

另外,从根本上说,千卡训练是一个复合体。很多时候,大家就知道一件事,“牛逼”,除了知道喊666之外,就少有了解到底牛逼在哪里。以至于说,觉得非常的高大上。我可以这么说一句,千卡及其以上的训练对于绝大多数人和企业而言,这就是个屠龙术 —— 除了某些个,用手指头数的出来的地方,别的地方完全没有这样的需求,也没有这样的资源来进行千卡训练。这是什么意思?意思就是。如果真有这样经验的人,流了出来,大概率很难对口的找工作,因为他在千卡集训中的训练经验和工程实践,大概率根本别的地方用不上。另一层意思是,如果你只是个一般人,那你想想就得了,就跟你可以意淫下某个自己喜欢的明星,但别真去追求,真要去了,你大概是连榜一大哥的待遇都不会有,注定了人财两空。我当然知道有人会说,那难道流出来的人,不能去那几个指头数的出来的地方吗?可以的,但是别着急,你们往后面看就知道了 —— 去还是能去的,但是如果他不是跟着大佬一起跑,到了新地方他们真不见的能继续干。

再来说说,千卡训练是个什么复合体 —— 至少是,科学,工程和人情世故。先看大模型训练中的任务怎么从“量变到质变”的 —— 任何一个小规模训练上的问题,放大几百几千倍之后,都有可能成为不可忽视的问题。比如,数据预处理,小的时候,你也许完全不在意,这到底是多少个毫秒搞定的。但是,如果你现在有上百T的数据要处理,手一抖,写个不那么高效的算法,多处理个几天,甚至几周都有可能。当然,更可怕的就不是慢,而是坏了 —— 一个小bug可以坏了整条pipeline。有的时候,你甚至都不能说是bug,但是反正不爽就是了。比如,两个人写预处理,一个人把图片弄成了BGR,一个弄成了RGB —— 又不是不能用,但是就是膈应人;又比如,数据原图太大了,要统一缩放,然后有人做的时候直接就缩成了方形,然后呢?我们之后需要的模型要正常长宽比的数据又该怎么办呢?再来一次嘛?再搞两个礼拜?你说这是个什么问题?可以看成是科学,当然也可以看成是工程的一部分,这两个就是紧密结合在一起的,单单你会调模型,在这个千卡训练的事情上,你是玩不转的。因为很多问题卡脖子的问题,根本就是,或者大概率是工程问题 —— 什么网络通信,什么磁盘空间,什么读取速度,什么数值稳定,小规模的时候,你都可以不用管,想怎么搞怎么搞,怎么搞可以怎么有。可是上了规模之后,很多东西都被限定死了,根本不是你想怎么干就能怎么干的。我说的这个话,大概很多用pytorch+cuda的朋友也不见的认同,毕竟这套组合下,没有太多技术支持的团队也干成了这样的事情。但是,这背后是因为,使用能支持这样训练的云服务本身就意味着,付了更多的钱,在已经白嫖了nvidia一波的前提下,外加meta(pytorch),外加微软(deepspeed),外加……,又变相雇佣了一个专门的支持团队。但是,这些都不改变一个事实 —— 那就是,这都是你跟着前人的脚步前进,有人替你已经把这条路上的坑,踩的差不多了。可是,如果你要做些原创性的工作呢?必然是会遇到很多前人都不会有的问题。

多说一句,也许有人会说,“我不关心别的,我就只关心pytorch+cuda下,做训练的经验”。那我告诉你,这本质上这跟你单机单卡训练就不应该有什么不一样,跟是不是pytorch,用不用cuda都没什么大关系 —— 你想想最理想情况下,这是不是就应该跟单机单卡训练一样么,无非就是现在的这个“单机”的GPU内存是所有的机器GPU内存的总和,能让你用一个更大的batch size和学习率。至于,GPU内部怎么通信,数据怎么通信,各个机器怎么通信,gradient传播怎么实现,需要你这个训模师知道吗?你在单机单卡的时候都不用知道,在单机多卡的时候不用知道,在小规模分布式训练的时候不用知道,那为什么到了千卡的时候,你就应该知道了?理想情况下,就算到了百万卡,也不用做建模的你去知道这里的各种工程实践。

那千卡训练到底难在哪里了?首先,就是难在之前提及的工程上面了 —— 简单假设一个卡在一天内不挂掉的概率是p,那么现在现在千卡同时一天内不挂掉的概率是多少?算算你就知道,对于p^1000,其实有机器挂掉才是正常的。如果是万卡呢?你可见的是,机器N越多,p^N就越小,这个事情就是越难。有人要说“我单机训练的时候,几年都遇不到问题,老黄的GPU稳定的一塌糊涂。”对此,我也只有呵呵,没人告诉你训练不下去都是GPU的问题。你大概是选择性忘记了各种自己训练中遇到的事情 —— 比如,上次实验中断,GPU进程没杀干净,还占着内存;和人共享的服务器上,有个卧龙觉得你训练的时候CPU占用率低了点,给你加了点任务;好巧不巧,默认的缓存地址没地方了,包装不上了,预训练模型下不来了…… 说这么多看似和训练无关的事情是因为,所有这些都要能自动化,因为里面有一个地方翻车了,你训练就进行不下去了。通信连不上,磁盘满了,遇上了奇葩的GPU,IO巨慢 …… 不管什么原因挂掉了,关键的是之后应该怎么办?有没有可能对这个出问题的机器进行热替换?怎么办才能最大程度不影响训练进程?怎么样才能在下次避免同样的问题。当然,实际情况可以更加复杂,GPU不见的是同批次的,模型也可以大到,哪怕在A100上也只能这个机器放一部分,那个机器放一部分……

但是也别误解,以为千卡训练,就对训模师而言,其实没什么挑战。这样的理解显然是错的。这对训模师的实操来说,肯定是一个巨大的挑战。完全是拿着卖白菜钱(想想你年薪才多少,算你年薪百万好了),操着卖白粉的心(这千卡训练要花多少钱?你年薪都不够它的一个零头)。因为这机器一开,实在是太烧金了。而且可见的是,你必然是要去debug的 —— 为什么小模型的时候,训练的挺好的,一变大就翻车了?或者说,虽然没翻车,但是为什么性能就涨了一丢丢?或者为什么前面训练挺稳定的,到了后面的loss curve就会有很大的spike?有经验的训模师能更早,更快的发现问题。也能更快和更好的解决问题。很多时候,也真不见的看log就能看出来点啥的,看数据,看gradient的大小分布,和其他模型的训练进行记录做比对,甚至做可视化,都是很有必要的。而这所有的一切,都需要你很有经验 —— 同样的log,有人就能一眼看出来问题在哪里,有人就只能对着发呆,或者机械性的说“换一组参数再试一下”。同样觉得可能哪里有问题,有人就能知道应该来验证这个猜想是对是错,有人就只能天马行空的给出一堆,谁也不知道对不对的原因。所以,一旦这条路线被摸索出来之后,其实也就没什么难度了 —— 数据,脚本,机器都在那里了,我就问你,我在服务器上run那条千卡训练命令,跟你run的能有什么不同?所以,真正的关键不是在于有没有用过千卡GPU训练过模型,而是有没有从头至尾,一路披荆斩棘的自己淌出来一条可重复的模型技术路线!!

当然,如果你要以为,这事情就只是技术,那也是太年轻了点。机器一开,要多少钱,这账真要算准从训模师的角度说是很难的,毕竟具体价格都是大公司之间协议的,属于商业机密,但是估算个大概的数目不难。按照aws的p4d算(8卡A100,见下图),便宜的算法,千卡训练一个月,需要花费 $11.57/每台小时*24小时/天*125台*30天 = $1,041,340;按照阿里云的算法,单卡年费¥170533.8,也就是¥19.47每小时,但是算上多卡的费用,这实际上比上面aws的价格更贵。当然,你也许能用更便宜的价格拿到机器,比如别找这么大的云服务平台,找个小的,但是再少还能少多少,算打5折,这都是50多万美刀,350多万人民币一个月。要知道,这可是训练一次的价格哦。一个能用的模型背后,可是5x,甚至10x更多的不能用模型哦,所以烧个几千万,真跟玩一样。

正是因为这么贵,所以也同样表明了,为什么一定要找有经验的训模师 —— 你要知道,你自己的每个实验决定,都是变相的花出去十几,几十甚至上百万的美金。早发现问题,早停下来;早解决问题,早开始;知道怎么偷懒,什么样的ablation study可以跳什么必须做,什么时候可以用小模型替代,什么时候可以用一个老的checkpoint来个jump start,什么时候直接白嫖论文上的结论就行……,所有这些都和花多少钱才能把这个事情办了,输出一个达标的模型,直接相关 。相反的,万一你要找个拉垮的训模师,前面不知道怎么计划,代码不知道怎么验证,训练起来了不懂怎么有效监控,有了异常不知道如何排除,……,最后都要靠着模型训练全完了之后做evaluation才知道行不行的那种。那么就算预算全花完了,什么都没有训练出来,我也没有什么好奇怪的。

铺垫这么多,终于可以来谈谈最后一个层面 —— 人情世故了。你看,千卡训练这个事情,有这么大的风险翻车,要花这么多的预算,那么现在问题来了,你要是部门领导,你让谁来干这个事情?哦,你想放权给下面的经理,让他来找人?又或者找个刚来的博士?找个顶校+顶会的博士?不管你怎么找,可问题是,你就这么信得过他吗?你怎么保证他,能干这个,能干好这个,不会中途跑路,不会磨洋工…… 要知道,这样大的项目和预算,如果要真干塌了,不说整个部门要一起完蛋,至少这一条线的人员必然是要担责的,哪怕是主要领导也跑不掉。所以喽,关键的关键是,你必须找自己信的过的人,还要找确实有能力可以担当重任的人 —— 真正最后来干这个千卡训练的人,不但自己技术要过硬,更是团队的中坚力量,至少也要得到一两个大头目的支持,而且还要得到小头目支持。你再牛逼,没信任没大佬支持,这事情不说完全不可能,也是基本没可能。你再牛逼,要是真的小头目给你上眼药,比如,跟上面吹风,“好像看见你在看招聘网站”,你想大头目心里会不会有阴影?所以,别给我扯什么,老子有多少顶会顶刊,老子导师是谁谁谁,这在绝大多数情况,都不好使。所以,刚毕业的,或者做实习的,或者刚工作的,如果宣称自己有这个经验,就是一眼丁真。因为上面是绝对不会找不信任的人来这样重要的工作,这跟你有没有相关的工作经验无关。这同样意味着,真正干这些事情的人也很难流出来 —— 因为对于嫡系来说,加薪升职,在干好了的前提下,那还不都是so easy吗?所以,是你,你愿意出来吗?出来了,就算你牛逼,但是获取信任,成为嫡系也要一个时间,不是吗?

让 Python 拥有 C/C++ 一样的速度,编译神器 Codon

Python 的运行速度快吗?虽说不同场景不同定论,但整体而言,它没有 C、Java 快。这也导致 Python 凭借可读性、简单易上手、良好的生态系统横行 AI 领域时,一提到速度,就成为众多开发者头疼的问题。

为了解决这一难题,麻省理工学院的计算机科学家出手了,他们共同研发了一种名为 Codon 的 Python 编译器,可以将 Python 代码转化为本地机器代码,而不会对运行时的性能产生影响。

当前,Codon 已经在 GitHub 上开源:https://github.com/exaloop/codon Docs 

Codon 与 C/C++ 性能相当:

“在单线程上,比 Python 的典型速度提高了 10-100 倍或更多,”Codon repo写道,”Codon 的性能通常与 C/C++ 的性能相当(有时甚至更好)。”

与 Python 不同的是,Codon 支持本地多线程,这可以使速度提高许多倍。

Codon 最初是作为一个在 Python 中创建高性能特定领域语言(DSL,domain-specific language)的框架而开发的。DSL 是专注于特定目的的语言,而不是像 Python 或 C 这样的通用编程语言。

据官方 GitHub repo 透露,Codon 源于 Seq 项目,后者是一个用于生物信息学和遗传学的 DSL,现如今它已经成长为一个与 Python 3 基本兼容的语言编译器。

近期,外媒 The Register 通过该工具的研究团队内部最新分享了一个关于 Codon 的论文。本论文的作者包含了,MIT、维多利亚大学等多名研究人员,如 Ariya Shajii(Exaloop)、Gabriel Ramirez(MIT CSAIL)、Haris Smajlović(加拿大维多利亚大学)、Jessica Ray(MIT CSAIL)、Bonnie Berger(MIT CSAIL)、Saman Amarasinghe(MIT CSAIL)和 Ibrahim Numanagić(维多利亚大学)。

这篇论文指出,”与其他面向性能的 Python 实现(如 PyPy 或 Numba)不同,Codon 是作为一个独立的系统从头开始构建的,它可以提前编译为静态可执行文件,并且不与现有的 Python 运行时(如 CPython 或 RPython)绑定执行。因此,Codon 可以实现更好的性能,并克服运行时的特定问题,如全局解释器锁”。

在论文中,作者也讨论了各种基于 Codon 的高性能 DSL,这些 DSL 设计用于生物信息学、数据压缩和并行编程,也利用了 Codon 的编译器基础设施。但是 Codon 可以大幅加速标准的 Python 程序,尽管那些依赖外部库(如 Django 或 DocUtils)的程序必须依赖 CPython 桥接,这就限制了 CPython 的性能。

“Codon 不需要使用像 numpy 这样的 C 语言实现的库来重写程序,也不需要完全用 C 语言重写,而是可以使用相同的 Python 实现,并给出你用 C 语言重写的相同性能”,麻省理工学院教授和 CSAIL 首席研究员 Saman Amarasinghe说 道,”因此,我相信 Codon 是成功的 Python 应用程序的最简单的前进道路,这些应用程序由于缺乏性能而达到了一个极限。”

测试

那么 Codon 是否真的如说的那样快?在 Codon 论坛上,一位开发者进行了测试:

$ cat fib.py 
def fib(n): if n == 0: return 0 elif n == 1: return 1 else: return fib(n-1) + fib(n-2)
if __name__ == "__main__": import sys print(fib(int(sys.argv[1])))

CPython 3.1

$ python fib.py 40102334155# mem: 8'816_KB# time: 18.42_s

PyPy 7.3.9

$ pypy fib.py 40102334155# mem: 74'596_kB# time: 4.99_s# ~= 3.7x
Codon compiled
$ codon build -release fib.py$ ./fib 40102334155# mem: 5'612_kB# time: 0.26_s# ~= 70.8x

Codon with python interpreter

# in fibpy.py, we just add `@python` decorator to fib function$ codon build -release fibpy.pyexport CODON_PYTHON=/path/to/libpython3.11.so$ ./fibpy 40102334155# mem: 12'828# time: 18.49# ~= 1x

最终发现,一个简单的 Codon 编译的斐波那契脚本比 CPython 版本快 70 多倍。

除此之外,该研究团队也在 GitHub 上贴出了 Codon 基准测试套件的结果,比较了 Python、PyPy、C++ 和 Codon 在一系列任务和应用上的表现。

据 Codon 官方文档显示,虽然 Codon 的语法和语义与 Python 的几乎相同,但还是有一些值得一提的区别,如数据类型方面:

  • 整数。Codon 的 int 是一个 64 位有符号的整数,而 Python 的(在版本 3 之后)可以是任意大的。然而 Codon 通过 Int[N] 支持更大的整数,其中 N 是位宽。
  • 字符串。Codon 目前使用 ASCII 字符串,与 Python 的 unicode 字符串不同。
  • 字典。Codon 的字典类型不保留插入顺序,与 Python 3.6 的不同。

此外,Codon 和 Python 在类型检查、数值运算、模块等维度还有些许的不同,更详细的内容可参考:https://docs.exaloop.io/codon/general/differences据悉,Codon 已经被商业化地应用在金融和生物信息学、深度学习等领域。

参考来源:

https://www.theregister.com/2023/03/11/python_codon_compiler/

GitHub 地址:https://github.com/exaloop/codon

PYTHON — 多进程和多线程

进程与线程的概念,以及为什么要有进程线程,其中有什么区别?

1. 基本概念:

进程是对运行时程序的封装,是系统进行资源调度和分配的的基本单位,实现了操作系统的并发

线程是进程的子任务,是CPU调度和分派的基本单位用于保证程序的实时性,实现进程内部的并发;线程是操作系统可识别的最小执行和调度单位。每个线程都独自占用一个虚拟处理器:独自的寄存器组指令计数器和处理器状态。每个线程完成不同的任务,但是共享同一地址空间(也就是同样的动态内存,映射文件,目标代码等等),打开的文件队列和其他内核资源

2. 区别:
  1. 一个线程只能属于一个进程,而一个进程可以有多个线程,但至少有一个线程。线程依赖于进程而存在。
  2. 进程在执行过程中拥有独立的内存单元,而多个线程共享进程的内存。(资源分配给进程,同一进程的所有线程共享该进程的所有资源。同一进程中的多个线程共享代码段(代码和常量),数据段(全局变量和静态变量),扩展段(堆存储)。但是每个线程拥有自己的栈段,栈段又叫运行时段,用来存放所有局部变量和临时变量。)
  3. 进程是资源分配的最小单位,线程是CPU调度的最小单位
  4. 系统开销: 由于在创建或撤消进程时,系统都要为之分配或回收资源,如内存空间、I/o设备等。因此,操作系统所付出的开销将显着地大于在创建或撤消线程时的开销。类似地,在进行进程切换时,涉及到整个当前进程CPU环境的保存以及新被调度运行的进程的CPU环境的设置。而线程切换只须保存和设置少量寄存器的内容,并不涉及存储器管理方面的操作。可见,进程切换的开销也远大于线程切换的开销
  5. 通信:由于同一进程中的多个线程具有相同的地址空间,致使它们之间的同步和通信的实现,也变得比较容易。进程间通信IPC,线程间可以直接读写进程数据段(如全局变量)来进行通信——需要进程同步和互斥手段的辅助,以保证数据的一致性。在有的系统中,线程的切换、同步和通信都无须操作系统内核的干预
  6. 进程编程调试简单可靠性高,但是创建销毁开销大;线程正相反,开销小,切换速度快,但是编程调试相对复杂
  7. 进程间不会相互影响 ;线程一个线程挂掉将导致整个进程挂掉
  8. 进程适应于多核、多机分布;线程适用于多核

进程与线程的一个简单解释:

计算机的核心是CPU,它承担了所有的计算任务。它就像一座工厂,时刻在运行。假定工厂的电力有限,一次只能供给一个车间使用。也就是说,一个车间开工的时候,其他车间都必须停工。背后的含义就是,单个CPU一次只能运行一个任务。进程就好比工厂的车间,它代表CPU所能处理的单个任务。任一时刻,CPU总是运行一个进程,其他进程处于非运行状态。

一个车间里,可以有很多工人。他们协同完成一个任务。线程就好比车间里的工人。一个进程可以包括多个线程。

车间的空间是工人们共享的,比如许多房间是每个工人都可以进出的。这象征一个进程的内存空间是共享的,每个线程都可以使用这些共享内存。可是,每间房间的大小不同,有些房间最多只能容纳一个人,比如厕所。里面有人的时候,其他人就不能进去了。这代表一个线程使用某些共享内存时,其他线程必须等它结束,才能使用这一块内存。

一个防止他人进入的简单方法,就是门口加一把锁。先到的人锁上门,后到的人看到上锁,就在门口排队,等锁打开再进去。这就叫“互斥锁”(Mutual exclusion,缩写 Mutex),防止多个线程同时读写某一块内存区域。

还有些房间,可以同时容纳n个人,比如厨房。也就是说,如果人数大于n,多出来的人只能在外面等着。这好比某些内存区域,只能供给固定数目的线程使用。

这时的解决方法,就是在门口挂n把钥匙。进去的人就取一把钥匙,出来时再把钥匙挂回原处。后到的人发现钥匙架空了,就知道必须在门口排队等着了。这种做法叫做“信号量”(Semaphore),用来保证多个线程不会互相冲突。

不难看出,mutex是semaphore的一种特殊情况(n=1时)。也就是说,完全可以用后者替代前者。但是,因为mutex较为简单,且效率高,所以在必须保证资源独占的情况下,还是采用这种设计。

操作系统的设计,因此可以归结为三点:

(1)以多进程形式,允许多个任务同时运行;

(2)以多线程形式,允许单个任务分成不同的部分运行;

(3)提供协调机制,一方面防止进程之间和线程之间产生冲突,另一方面允许进程之间和线程之间共享资源。

Python中的进程和线程:

python中的的multiprocess和threading模块用于进行多线程和多进程编程。

Python的多进程编程与multiprocess模块

python的多进程编程主要依靠multiprocess模块。我们先对比两段代码,看看多进程编程的优势。我们模拟了一个非常耗时的任务,计算8的20次方,为了使这个任务显得更耗时,我们还让它sleep 2秒。第一段代码是单进程计算(代码如下所示),我们按顺序执行代码,重复计算2次,并打印出总共耗时。

import time
import os

def long_time_task():
    print('当前进程: {}'.format(os.getpid()))
    time.sleep(2)
    print("结果: {}".format(8 ** 20))

if __name__ == "__main__":
    print('当前母进程: {}'.format(os.getpid()))
    start = time.time()
    for i in range(2):
        long_time_task()

    end = time.time()
    print("用时{}秒".format((end-start)))

输出结果如下,总共耗时4秒,至始至终只有一个进程14236。看来电脑计算8的20次方基本不费时。

当前母进程: 14236
当前进程: 14236
结果: 1152921504606846976
当前进程: 14236
结果: 1152921504606846976
用时4.01080060005188秒

第2段代码是多进程计算代码。我们利用multiprocess模块的Process方法创建了两个新的进程p1和p2来进行并行计算。Process方法接收两个参数, 第一个是target,一般指向函数名,第二个时args,需要向函数传递的参数。对于创建的新进程,调用start()方法即可让其开始。我们可以使用os.getpid()打印出当前进程的名字。

from multiprocessing import Process
import os
import time


def long_time_task(i):
    print('子进程: {} - 任务{}'.format(os.getpid(), i))
    time.sleep(2)
    print("结果: {}".format(8 ** 20))


if __name__=='__main__':
    print('当前母进程: {}'.format(os.getpid()))
    start = time.time()
    p1 = Process(target=long_time_task, args=(1,))
    p2 = Process(target=long_time_task, args=(2,))
    print('等待所有子进程完成。')
    p1.start()
    p2.start()
    p1.join()
    p2.join()
    end = time.time()
    print("总共用时{}秒".format((end - start)))

输出结果如下所示,耗时变为2秒,时间减了一半,可见并发执行的时间明显比顺序执行要快很多。你还可以看到尽管我们只创建了两个进程,可实际运行中却包含里1个母进程和2个子进程。之所以我们使用join()方法就是为了让母进程阻塞,等待子进程都完成后才打印出总共耗时,否则输出时间只是母进程执行的时间。

当前母进程: 6920
等待所有子进程完成。
子进程: 17020 - 任务1
子进程: 5904 - 任务2
结果: 1152921504606846976
结果: 1152921504606846976
总共用时2.131091356277466秒

知识点:

  • 新创建的进程与进程的切换都是要耗资源的,所以平时工作中进程数不能开太大。
  • 同时可以运行的进程数一般受制于CPU的核数。
  • 除了使用Process方法,我们还可以使用Pool类创建多进程。

利用multiprocess模块的Pool类创建多进程

很多时候系统都需要创建多个进程以提高CPU的利用率,当数量较少时,可以手动生成一个个Process实例。当进程数量很多时,或许可以利用循环,但是这需要程序员手动管理系统中并发进程的数量,有时会很麻烦。这时进程池Pool就可以发挥其功效了。可以通过传递参数限制并发进程的数量,默认值为CPU的核数。

Pool类可以提供指定数量的进程供用户调用,当有新的请求提交到Pool中时,如果进程池还没有满,就会创建一个新的进程来执行请求。如果池满,请求就会告知先等待,直到池中有进程结束,才会创建新的进程来执行这些请求。

下面介绍一下multiprocessing 模块下的Pool类的几个方法:

1.apply_async

函数原型:apply_async(func[, args=()[, kwds={}[, callback=None]]])

其作用是向进程池提交需要执行的函数及参数, 各个进程采用非阻塞(异步)的调用方式,即每个子进程只管运行自己的,不管其它进程是否已经完成。这是默认方式。

2.map()

函数原型:map(func, iterable[, chunksize=None])

Pool类中的map方法,与内置的map函数用法行为基本一致,它会使进程阻塞直到结果返回。 注意:虽然第二个参数是一个迭代器,但在实际使用中,必须在整个队列都就绪后,程序才会运行子进程。

3.map_async()

函数原型:map_async(func, iterable[, chunksize[, callback]])
与map用法一致,但是它是非阻塞的。其有关事项见apply_async。

4.close()

关闭进程池(pool),使其不在接受新的任务。

5. terminate()

结束工作进程,不在处理未处理的任务。

6.join()

主进程阻塞等待子进程的退出, join方法要在close或terminate之后使用。

下例是一个简单的multiprocessing.Pool类的实例。因为小编我的CPU是4核的,一次最多可以同时运行4个进程,所以我开启了一个容量为4的进程池。4个进程需要计算5次,你可以想象4个进程并行4次计算任务后,还剩一次计算任务(任务4)没有完成,系统会等待4个进程完成后重新安排一个进程来计算。

from multiprocessing import Pool, cpu_count
import os
import time


def long_time_task(i):
    print('子进程: {} - 任务{}'.format(os.getpid(), i))
    time.sleep(2)
    print("结果: {}".format(8 ** 20))


if __name__=='__main__':
    print("CPU内核数:{}".format(cpu_count()))
    print('当前母进程: {}'.format(os.getpid()))
    start = time.time()
    p = Pool(4)
    for i in range(5):
        p.apply_async(long_time_task, args=(i,))
    print('等待所有子进程完成。')
    p.close()
    p.join()
    end = time.time()
    print("总共用时{}秒".format((end - start)))

知识点:

  • 对Pool对象调用join()方法会等待所有子进程执行完毕,调用join()之前必须先调用close()或terminate()方法,让其不再接受新的Process了。

输出结果如下所示,5个任务(每个任务大约耗时2秒)使用多进程并行计算只需4.37秒,, 耗时减少了60%,可见并行计算优势还是很明显的。

CPU内核数:4
当前母进程: 2556
等待所有子进程完成。
子进程: 16480 - 任务0
子进程: 15216 - 任务1
子进程: 15764 - 任务2
子进程: 10176 - 任务3
结果: 1152921504606846976
结果: 1152921504606846976
子进程: 15216 - 任务4
结果: 1152921504606846976
结果: 1152921504606846976
结果: 1152921504606846976
总共用时4.377134561538696秒

相信大家都知道python解释器中存在GIL(全局解释器锁), 它的作用就是保证同一时刻只有一个线程可以执行代码。由于GIL的存在,很多人认为python中的多线程其实并不是真正的多线程,如果想要充分地使用多核CPU的资源,在python中大部分情况需要使用多进程。然而这并意味着python多线程编程没有意义哦,请继续阅读下文。

多进程间的数据共享与通信

通常,进程之间是相互独立的,每个进程都有独立的内存。通过共享内存(nmap模块),进程之间可以共享对象,使多个进程可以访问同一个变量(地址相同,变量名可能不同)。多进程共享资源必然会导致进程间相互竞争,所以应该尽最大可能防止使用共享状态。还有一种方式就是使用队列queue来实现不同进程间的通信或数据共享,这一点和多线程编程类似。

下例这段代码中中创建了2个独立进程,一个负责写(pw), 一个负责读(pr), 实现了共享一个队列queue。

from multiprocessing import Process, Queue
import os, time, random

# 写数据进程执行的代码:
def write(q):
    print('Process to write: {}'.format(os.getpid()))
    for value in ['A', 'B', 'C']:
        print('Put %s to queue...' % value)
        q.put(value)
        time.sleep(random.random())

# 读数据进程执行的代码:
def read(q):
    print('Process to read:{}'.format(os.getpid()))
    while True:
        value = q.get(True)
        print('Get %s from queue.' % value)

if __name__=='__main__':
    # 父进程创建Queue,并传给各个子进程:
    q = Queue()
    pw = Process(target=write, args=(q,))
    pr = Process(target=read, args=(q,))
    # 启动子进程pw,写入:
    pw.start()
    # 启动子进程pr,读取:
    pr.start()
    # 等待pw结束:
    pw.join()
    # pr进程里是死循环,无法等待其结束,只能强行终止:
    pr.terminate()

输出结果如下所示:

Process to write: 3036
Put A to queue...
Process to read:9408
Get A from queue.
Put B to queue...
Get B from queue.
Put C to queue...
Get C from queue.

python中的多线程(伪多线程):

我们知道 Python 之所以灵活和强大,是因为它是一个解释性语言,边解释边执行,实现这种特性的标准实现叫作 CPython。

它分两步来运行 Python 程序:

  • 首先解析源代码文本,并将其编译为字节码(bytecode)[1]
  • 然后采用基于栈的解释器来运行字节码
  • 不断循环这个过程,直到程序结束或者被终止

灵活性有了,但是为了保证程序执行的稳定性,也付出了巨大的代价:

引入了 全局解释器锁 GIL(global interpreter lock)[2]

以保证同一时间只有一个字节码在运行,这样就不会因为没用事先编译,而引发资源争夺和状态混乱的问题了。

看似 “十全十美” ,但,这样做,就意味着多线程执行时,会被 GIL 变为单线程,无法充分利用硬件资源。

戴着镣铐跳舞

难道 Python 里的多线程真的没用吗?

其实也并不是,虽然了因为 GIL,无法实现真正意义上的多线程,但,多线程机制,还是为我们提供了两个重要的特性。

一:多线程写法可以让某些程序更好写

怎么理解呢?

如果要解决一个需要同时维护多种状态的程序,用单线程是实现是很困难的。

比如要检索一个文本文件中的数据,为了提高检索效率,可以将文件分成小段的来处理,最先在那段中找到了,就结束处理过程。

用单线程的话,很难实现同时兼顾多个分段的情况,只能顺序,或者用二分法执行检索任务。

而采用多线程,可以将每个分段交给每个线程,会轮流执行,相当于同时推荐检索任务,处理起来,效率会比顺序查找大大提高。

二:处理阻塞型 I/O 任务效率更高

阻塞型 I/O 的意思是,当系统需要与文件系统(也包括网络和终端显示)交互时,由于文件系统相比于 CPU 的处理速度慢得多,所以程序会被设置为阻塞状态,即,不再被分配计算资源。

直到文件系统的结果返回,才会被激活,将有机会再次被分配计算资源。

也就是说,处于阻塞状态的程序,会一直等着。

那么如果一个程序是需要不断地从文件系统读取数据,处理后在写入,单线程的话就需要等等读取后,才能处理,等待处理完才能写入,于是处理过程就成了一个个的等待。

而用多线程,当一个处理过程被阻塞之后,就会立即被 GIL 切走,将计算资源分配给其他可以执行的过程,从而提示执行效率。

有了这两个特性,就说明 Python 的多线程并非一无是处,如果能根据情况编写好,效率会大大提高,只不过对于计算密集型的任务,多线程特性爱莫能助。

自强不息

了解到 Python 多线程的问题和解决方案,对于钟爱 Python 的我们,何去何从呢?

有句话用在这里很合适:

求人不如求己

哪怕再怎么厉害的工具或者武器,都无法解决所有的问题,而问题之所以能被解决,主要是因为我们的主观能动性。

对情况进行分析判断,选择合适的解决方案,不就是需要我们做的么?

对于 Python 中 多线程的诟病,我们更多的是看到它阳光和美的一面,而对于需要提升速度的地方,采取合适的方式。这里简单总结一下:

  1. I/O 密集型的任务,采用 Python 的多线程完全没用问题,可以大幅度提高执行效率
  2. 对于计算密集型任务,要看数据依赖性是否低,如果低,采用 ProcessPoolExecutor 代替多线程处理,可以充分利用硬件资源
  3. 如果数据依赖性高,可以考虑将关键的地方该用 C 来实现,一方面 C 本身比 Python 更快,另一方面,C 可以之间使用更底层的多线程机制,而完全不用担心受 GIL 的影响
  4. 大部分情况下,对于只能用多线程处理的任务,不用太多考虑,之间利用 Python 的多线程机制就好了,不用考虑太多

Python的多线程编程与threading模块

python 3中的多进程编程主要依靠threading模块。创建新线程与创建新进程的方法非常类似。threading.Thread方法可以接收两个参数, 第一个是target,一般指向函数名,第二个时args,需要向函数传递的参数。对于创建的新线程,调用start()方法即可让其开始。我们还可以使用current_thread().name打印出当前线程的名字。 下例中我们使用多线程技术重构之前的计算代码。

import threading
import time


def long_time_task(i):
    print('当前子线程: {} - 任务{}'.format(threading.current_thread().name, i))
    time.sleep(2)
    print("结果: {}".format(8 ** 20))


if __name__=='__main__':
    start = time.time()
    print('这是主线程:{}'.format(threading.current_thread().name))
    t1 = threading.Thread(target=long_time_task, args=(1,))
    t2 = threading.Thread(target=long_time_task, args=(2,))
    t1.start()
    t2.start()

    end = time.time()
    print("总共用时{}秒".format((end - start)))

下面是输出结果。为什么总耗时居然是0秒? 我们可以明显看到主线程和子线程其实是独立运行的,主线程根本没有等子线程完成,而是自己结束后就打印了消耗时间。主线程结束后,子线程仍在独立运行,这显然不是我们想要的。

这是主线程:MainThread
当前子线程: Thread-1 - 任务1
当前子线程: Thread-2 - 任务2
总共用时0.0017192363739013672秒
结果: 1152921504606846976
结果: 1152921504606846976

如果要实现主线程和子线程的同步,我们必需使用join方法(代码如下所示)。

import threading
import time


def long_time_task(i):
    print('当前子线程: {} 任务{}'.format(threading.current_thread().name, i))
    time.sleep(2)
    print("结果: {}".format(8 ** 20))


if __name__=='__main__':
    start = time.time()
    print('这是主线程:{}'.format(threading.current_thread().name))
    thread_list = []
    for i in range(1, 3):
        t = threading.Thread(target=long_time_task, args=(i, ))
        thread_list.append(t)

    for t in thread_list:
        t.start()

    for t in thread_list:
        t.join()

    end = time.time()
    print("总共用时{}秒".format((end - start)))

修改代码后的输出如下所示。这时你可以看到主线程在等子线程完成后才答应出总消耗时间(2秒),比正常顺序执行代码(4秒)还是节省了不少时间。

这是主线程:MainThread
当前子线程: Thread - 1 任务1
当前子线程: Thread - 2 任务2
结果: 1152921504606846976
结果: 1152921504606846976
总共用时2.0166890621185303秒

当我们设置多线程时,主线程会创建多个子线程,在python中,默认情况下主线程和子线程独立运行互不干涉。如果希望让主线程等待子线程实现线程的同步,我们需要使用join()方法。如果我们希望一个主线程结束时不再执行子线程,我们应该怎么办呢? 我们可以使用t.setDaemon(True),代码如下所示。

import threading
import time


def long_time_task():
    print('当子线程: {}'.format(threading.current_thread().name))
    time.sleep(2)
    print("结果: {}".format(8 ** 20))


if __name__=='__main__':
    start = time.time()
    print('这是主线程:{}'.format(threading.current_thread().name))
    for i in range(5):
        t = threading.Thread(target=long_time_task, args=())
        t.setDaemon(True)
        t.start()

    end = time.time()
    print("总共用时{}秒".format((end - start)))

通过继承Thread类重写run方法创建新进程

除了使用Thread()方法创建新的线程外,我们还可以通过继承Thread类重写run方法创建新的线程,这种方法更灵活。下例中我们自定义的类为MyThread, 随后我们通过该类的实例化创建了2个子线程。

#-*- encoding:utf-8 -*-
import threading
import time


def long_time_task(i):
    time.sleep(2)
    return 8**20


class MyThread(threading.Thread):
    def __init__(self, func, args , name='', ):
        threading.Thread.__init__(self)
        self.func = func
        self.args = args
        self.name = name
        self.result = None

    def run(self):
        print('开始子进程{}'.format(self.name))
        self.result = self.func(self.args[0],)
        print("结果: {}".format(self.result))
        print('结束子进程{}'.format(self.name))


if __name__=='__main__':
    start = time.time()
    threads = []
    for i in range(1, 3):
        t = MyThread(long_time_task, (i,), str(i))
        threads.append(t)

    for t in threads:
        t.start()
    for t in threads:
        t.join()

    end = time.time()
    print("总共用时{}秒".format((end - start)))

输出结果如下所示:

开始子进程1
开始子进程2
结果: 1152921504606846976
结果: 1152921504606846976
结束子进程1
结束子进程2
总共用时2.005445718765259秒

不同线程间的数据共享

一个进程所含的不同线程间共享内存,这就意味着任何一个变量都可以被任何一个线程修改,因此线程之间共享数据最大的危险在于多个线程同时改一个变量,把内容给改乱了。如果不同线程间有共享的变量,其中一个方法就是在修改前给其上一把锁lock,确保一次只有一个线程能修改它。threading.lock()方法可以轻易实现对一个共享变量的锁定,修改完后release供其它线程使用。比如下例中账户余额balance是一个共享变量,使用lock可以使其不被改乱。

# -*- coding: utf-8 -*

import threading


class Account:
    def __init__(self):
        self.balance = 0

    def add(self, lock):
        # 获得锁
        lock.acquire()
        for i in range(0, 100000):
            self.balance += 1
        # 释放锁
        lock.release()

    def delete(self, lock):
        # 获得锁
        lock.acquire()
        for i in range(0, 100000):
            self.balance -= 1
            # 释放锁
        lock.release()


if __name__ == "__main__":
    account = Account()
    lock = threading.Lock()
    # 创建线程
   thread_add = threading.Thread(target=account.add, args=(lock,), name='Add')
    thread_delete = threading.Thread(target=account.delete, args=(lock,), name='Delete')

    # 启动线程
   thread_add.start()
    thread_delete.start()

    # 等待线程结束
   thread_add.join()
    thread_delete.join()

    print('The final balance is: {}'.format(account.balance))

另一种实现不同线程间数据共享的方法就是使用消息队列queue。不像列表,queue是线程安全的,可以放心使用,见下文。

使用queue队列通信-经典的生产者和消费者模型

下例中创建了两个线程,一个负责生成,一个负责消费,所生成的产品存放在queue里,实现了不同线程间沟通。

from queue import Queue
import random, threading, time


# 生产者类
class Producer(threading.Thread):
    def __init__(self, name, queue):
        threading.Thread.__init__(self, name=name)
        self.queue = queue

    def run(self):
        for i in range(1, 5):
            print("{} is producing {} to the queue!".format(self.getName(), i))
            self.queue.put(i)
            time.sleep(random.randrange(10) / 5)
        print("%s finished!" % self.getName())


# 消费者类
class Consumer(threading.Thread):
    def __init__(self, name, queue):
        threading.Thread.__init__(self, name=name)
        self.queue = queue

    def run(self):
        for i in range(1, 5):
            val = self.queue.get()
            print("{} is consuming {} in the queue.".format(self.getName(), val))
            time.sleep(random.randrange(10))
        print("%s finished!" % self.getName())


def main():
    queue = Queue()
    producer = Producer('Producer', queue)
    consumer = Consumer('Consumer', queue)

    producer.start()
    consumer.start()

    producer.join()
    consumer.join()
    print('All threads finished!')


if __name__ == '__main__':
    main()

队列queue的put方法可以将一个对象obj放入队列中。如果队列已满,此方法将阻塞至队列有空间可用为止。queue的get方法一次返回队列中的一个成员。如果队列为空,此方法将阻塞至队列中有成员可用为止。queue同时还自带emtpy(), full()等方法来判断一个队列是否为空或已满,但是这些方法并不可靠,因为多线程和多进程,在返回结果和使用结果之间,队列中可能添加/删除了成员。

Python多进程和多线程哪个快?

由于GIL的存在,很多人认为Python多进程编程更快,针对多核CPU,理论上来说也是采用多进程更能有效利用资源。网上很多人已做过比较,我直接告诉你结论吧。

  • 对CPU密集型代码(比如循环计算) – 多进程效率更高
  • 对IO密集型代码(比如文件操作,网络爬虫) – 多线程效率更高。

为什么是这样呢?其实也不难理解。对于IO密集型操作,大部分消耗时间其实是等待时间,在等待时间中CPU是不需要工作的,那你在此期间提供双CPU资源也是利用不上的,相反对于CPU密集型代码,2个CPU干活肯定比一个CPU快很多。那么为什么多线程会对IO密集型代码有用呢?这时因为python碰到等待会释放GIL供新的线程使用,实现了线程间的切换。

TensorRT – 使用trtexec工具转换模型、运行模型、测试网络性能

转换模型将onnx转换为TensorRT:

方法一、trtexec

trtexec是在tensorrt包中自带的转换程序,该程序位于bin目录下,用起来比较方便,也是最简单的trt模型转换方式,在使用之前需要系统安装好cuda和cudnn,否则无法正常运行。使用示例如下:

首先将pytorch模型先转换成onnx模型,示例代码如下:

def torch2onnx(model_path,onnx_path):
model = load_model(model_path)
test_arr = torch.randn(1,3,32,448)
input_names = ['input']
output_names = ['output']
tr_onnx.export(
model,
test_arr,
onnx_path,
verbose=False,
opset_version=11,
input_names=input_names,
output_names=output_names,
dynamic_axes={"input":{3:"width"}} #动态推理W纬度,若需其他动态纬度可以自行修改,不需要动态推理的话可以注释这行
)
print('->>模型转换成功!')

trtexec转换命令如下:

固定尺寸模型转换:将ONNX模型转换为静态batchsize的TensorRT模型,启动所有精度以达到最佳性能,工作区大小设置为1024M

./trtexec --onnx=repvgg_a1.onnx --saveEngine=repvgg_a1.engine --workspace=1024  --fp16 --verbose

动态尺寸模型转换:将ONNX模型转换为动态batchsize的TensorRT模型,启动所有精度以达到最佳性能,工作区大小设置为1024M

./trtexec --onnx=repvgg_a1.onnx --saveEngine=repvgg_a1.engine --workspace=1024 --minShapes=input:1x3x32x32 --optShapes=input:1x3x32x320 --maxShapes=input:1x3x32x640 --fp16

注意:
–minShapes,–optShapes ,–maxShapes必须全部设置,设置的形式为:batchsize x 通道数 x 输入尺寸x x 输入尺寸y

例如:
--minShapes=input:1x3x416x416
--optShapes=input:8x3x416x416
--maxShapes=input:8x3x416x416

参看命名详解: ./trtexec –help, -h

trtexec的参数使用说明

1.1 Model Option 模型选项

–uff : UFF模型文件名–onnx : ONNX模型文件名–model : Caffe模型文件名,模式时无模型,使用随机权重–deploy : Caffe prototxt 文件名–output : 输出名称(可多次指定);UFF和Caffe至少需要一个输出–uffInput : 输入blob名称及其维度(X、Y、Z=C、H、W),可以多次指定;UFF型号至少需要一个–uffNHWC : 设置输入是否在NHWC布局中而不是NCHW中(在–uffInput中使用X、Y、Z=H、W、C顺序)

1.2 Build Options 构建选项

–maxBatch : 设置最大批处理大小并构建隐式批处理引擎(默认值=1)–explicitBatch :构建引擎时使用显式批量大小(默认 = 隐式)–minShapes=spec : 使用提供的最小形状的配置文件构建动态形状–optShapes=spec : 使用提供的 opt 形状的配置文件构建动态形状–maxShapes=spec : 使用提供的最大形状的配置文件构建动态形状–minShapesCalib=spec : 使用提供的最小形状的配置文件校准动态形状–optShapesCalib=spec : 使用提供的 opt 形状的配置文件校准动态形状–maxShapesCalib=spec :使用提供的最大形状的配置文件校准动态形状注意:必须提供所有三个 min、opt 和 max 形状。但是,如果只提供了 opt 形状,那么它将被扩展,以便将最小形状和最大形状设置为与 opt 形状相同的值。此外,使用 动态形状意味着显式批处理。 输入名称可以用转义单引号括起来(例如:‘Input:0’)。示例输入形状规范:input0:1x3x256x256,input1:1x3x128x128 每个输入形状都作为键值对提供,其中 key 是输入名称 值是用于该输入的维度(包括批次维度)。 每个键值对都使用冒号 (😃 分隔键和值。 可以通过逗号分隔的键值对提供多个输入形状。–inputIOFormats=spec : 每个输入张量的类型和格式(默认所有输入为fp32:chw)注意:如果指定此选项,请按照与网络输入ID相同的顺序为所有输入设置逗号分隔的类型和格式(即使只有一个输入需要指定IO格式)或设置一次类型和格式以进行广播。–outputIOFormats=spec : 每个输出张量的类型和格式(默认所有输入为fp32:chw)注意:如果指定此选项,请按照与网络输出ID相同的顺序为所有输出设置逗号分隔的类型和格式(即使只有一个输出需要指定IO格式)或设置一次类型和格式以进行广播。–workspace=N : 以M为单位设置工作区大小(默认值 = 16)–noBuilderCache : 在构建器中禁用时序缓存(默认是启用时序缓存)–nvtxMode=mode : 指定 NVTX 注释详细程度。 mode ::= default|verbose|none–minTiming=M : 设置内核选择中使用的最小迭代次数(默认值 = 1)–avgTiming=M : 为内核选择设置每次迭代的平均次数(默认值 = 8)–noTF32 : 禁用 tf32 精度(默认是启用 tf32,除了 fp32)–refit : 将引擎标记为可改装。这将允许检查引擎内的可改装层和重量。–fp16 : 除 fp32 外,启用 fp16 精度(默认 = 禁用)–int8 : 除 fp32 外,启用 int8 精度(默认 = 禁用)–best : 启用所有精度以达到最佳性能(默认 = 禁用)–calib= : 读取INT8校准缓存文件–safe : 仅测试安全受限流中可用的功能–saveEngine= : 保存序列化模型的文件名–loadEngine= : 加载序列化模型的文件名–tacticSources=tactics : 通过从默认策略源(默认 = 所有可用策略)中添加 (+) 或删除 (-) 策略来指定要使用的策略。

1.3 Inference Options 推理选项

–batch=N : 为隐式批处理引擎设置批处理大小(默认值 = 1)–shapes=spec : 为动态形状推理输入设置输入形状。注意:使用动态形状意味着显式批处理。 输入名称可以用转义的单引号括起来(例如:‘Input:0’)。 示例输入形状规范:input0:1x3x256x256, input1:1x3x128x128 每个输入形状都作为键值对提供,其中键是输入名称,值是用于该输入的维度(包括批次维度)。 每个键值对都使用冒号 (😃 分隔键和值。 可以通过逗号分隔的键值对提供多个输入形状。–loadInputs=spec :从文件加载输入值(默认 = 生成随机输入)。 输入名称可以用单引号括起来(例如:‘Input:0’)–iterations=N : 至少运行 N 次推理迭代(默认值 = 10)–warmUp=N : 在测量性能之前运行 N 毫秒以预热(默认值 = 200)–duration=N : 运行至少 N 秒挂钟时间的性能测量(默认值 = 3)–sleepTime=N : 延迟推理以启动和计算之间的 N 毫秒间隔开始(默认 = 0)–streams=N : 实例化 N 个引擎以同时使用(默认值 = 1)–exposeDMA : 串行化进出设备的 DMA 传输。 (默认 = 禁用)–noDataTransfers : 在推理过程中,请勿将数据传入和传出设备。 (默认 = 禁用)–useSpinWait : 主动同步 GPU 事件。 此选项可能会减少同步时间,但会增加 CPU 使用率和功率(默认 = 禁用)–threads : 启用多线程以驱动具有独立线程的引擎(默认 = 禁用)–useCudaGraph : 使用 cuda 图捕获引擎执行,然后启动推理(默认 = 禁用)–separateProfileRun : 不要在基准测试中附加分析器; 如果启用分析,将执行第二次分析运行(默认 = 禁用)–buildOnly : 跳过推理性能测量(默认 = 禁用)

1.4 Build and Inference Batch Options 构建和推理批处理选项
使用隐式批处理时,引擎的最大批处理大小(如果未指定)设置为推理批处理大小; 使用显式批处理时,如果仅指定形状用于推理,它们也将在构建配置文件中用作 min/opt/max; 如果只为构建指定了形状,则 opt 形状也将用于推理; 如果两者都被指定,它们必须是兼容的; 如果启用了显式批处理但都未指定,则模型必须为所有输入提供完整的静态维度,包括批处理大小

1.5 Reporting Options 报告选项

–verbose : 使用详细日志记录(默认值 = false)–avgRuns=N : 报告 N 次连续迭代的平均性能测量值(默认值 = 10)–percentile=P : 报告 P 百分比的性能(0<=P<=100,0 代表最大性能,100 代表最小性能;(默认 = 99%)–dumpRefit : 从可改装引擎打印可改装层和重量–dumpOutput : 打印最后一次推理迭代的输出张量(默认 = 禁用)–dumpProfile : 每层打印配置文件信息(默认 = 禁用)–exportTimes= : 将计时结果写入 json 文件(默认 = 禁用)–exportOutput= : 将输出张量写入 json 文件(默认 = 禁用)–exportProfile= : 将每层的配置文件信息写入 json 文件(默认 = 禁用)

1.6 System Options 系统选项

–device=N :选择 cuda 设备 N(默认 = 0)–useDLACore=N : 为支持 DLA 的层选择 DLA 核心 N(默认 = 无)–allowGPUFallback : 启用 DLA 后,允许 GPU 回退不受支持的层(默认 = 禁用)–plugins : 要加载的插件库 (.so)(可以多次指定)

1.7 Help 帮助
–help, -h : 打印以上帮助信息

方法2、使用python脚本

参考官方给到的demo写一个脚本转:官方脚本位于下载的目录:TensorRT-7.2.3.4/samples/python/yolov3_onnx/onnx_to_tensorrt.py

import os 
import tensorrt as trt
os.environ["CUDA_VISIBLE_DEVICES"]='0'
TRT_LOGGER = trt.Logger()
onnx_file_path = 'Unet375-simple.onnx'
engine_file_path = 'Unet337.trt'

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
            builder.max_workspace_size = 1 << 28 # 256MiB
            builder.max_batch_size = 1
            # Parse model file
            if not os.path.exists(onnx_file_path):
                print('ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.'.format(onnx_file_path))
                exit(0)
            print('Loading ONNX file from path {}...'.format(onnx_file_path))
            with open(onnx_file_path, 'rb') as model:
                print('Beginning ONNX file parsing')
                if not parser.parse(model.read()):
                    print ('ERROR: Failed to parse the ONNX file.')
                    for error in range(parser.num_errors):
                        print (parser.get_error(error))

            network.get_input(0).shape = [1, 3, 300, 400]
            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
            #network.mark_output(network.get_layer(network.num_layers-1).get_output(0))
            engine = builder.build_cuda_engine(network)
            print("Completed creating Engine")
            with open(engine_file_path, "wb") as f:
                f.write(engine.serialize())

运行ONNX模型

  • 在具有静态输入形状的全维模式下运行 ONNX 模型
trtexec --onnx=model.onnx
  • 使用给定的输入形状在全维模式下运行 ONNX 模型

trtexec –onnx=model.onnx –shapes=input:32x3x244x244

  • 使用一系列可能的输入形状对 ONNX 模型进行基准测试
trtexec --onnx=model.onnx --minShapes=input:1x3x244x244 --optShapes=input:16x3x244x244 --maxShapes=input:32x3x244x244 --shapes=input:5x3x244x244



trtexec --onnx=depth_feat_model.onnx --minShapes=input:1x4x128x128 --maxShapes=input:1x4x896x896 --shapes=input:1x4x512x512 --saveEngine=depth_feat_model.engine --verbose --workspace=1024 --fp32

网络性能测试

  • 加载转换后的TensorRT模型进行性能测试,指定batch大小
trtexec --loadEngine=mnist16.trt --batch=1

打印输出:
trtexec会打印出很多时间,这里需要对每个时间的含义进行解释,然后大家各取所需,进行评测。总的打印如下:
[09/06/2021-13:50:34] [I] Average on 10 runs - GPU latency: 2.74553 ms - Host latency: 3.74192 ms (end to end 4.93066 ms, enqueue 0.624805 ms)  # 跑了10次,GPU latency: GPU计算耗时, Host latency:GPU输入+计算+输出耗时,end to end:GPU端到端的耗时,eventout - eventin,enqueue:CPU异步耗时
[09/06/2021-13:50:34] [I] Host Latency
[09/06/2021-13:50:34] [I] min: 3.65332 ms (end to end 3.67603 ms)
[09/06/2021-13:50:34] [I] max: 5.95093 ms (end to end 6.88892 ms)
[09/06/2021-13:50:34] [I] mean: 3.71375 ms (end to end 5.30082 ms)
[09/06/2021-13:50:34] [I] median: 3.70032 ms (end to end 5.32935 ms)
[09/06/2021-13:50:34] [I] percentile: 4.10571 ms at 99% (end to end 6.11792 ms at 99%)
[09/06/2021-13:50:34] [I] throughput: 356.786 qps
[09/06/2021-13:50:34] [I] walltime: 3.00741 s
[09/06/2021-13:50:34] [I] Enqueue Time
[09/06/2021-13:50:34] [I] min: 0.248474 ms
[09/06/2021-13:50:34] [I] max: 2.12134 ms
[09/06/2021-13:50:34] [I] median: 0.273987 ms
[09/06/2021-13:50:34] [I] GPU Compute
[09/06/2021-13:50:34] [I] min: 2.69702 ms
[09/06/2021-13:50:34] [I] max: 4.99219 ms
[09/06/2021-13:50:34] [I] mean: 2.73299 ms
[09/06/2021-13:50:34] [I] median: 2.71875 ms
[09/06/2021-13:50:34] [I] percentile: 3.10791 ms at 99%
[09/06/2021-13:50:34] [I] total compute time: 2.93249 s

Host Latency gpu: 输入+计算+输出 三部分的耗时
Enqueue Time:CPU异步的时间(该时间不具有参考意义,因为GPU的计算可能还没有完成)
GPU Compute:GPU计算的耗时
综上,去了Enqueue Time时间都是有意义的
  • 收集和打印时序跟踪信息
trtexec --deploy=data/AlexNet/AlexNet_N2.prototxt --output=prob --exportTimes=trace.json
  • 使用多流调整吞吐量

调整吞吐量可能需要运行多个并发执行流。例如,当实现的延迟完全在所需阈值内时,我们可以增加吞吐量,即使以一些延迟为代价。例如,为批量大小 1 和 2 保存引擎并假设两者都在 2ms 内执行,延迟阈值:

trtexec --deploy=GoogleNet_N2.prototxt --output=prob --batch=1 --saveEngine=g1.trt --int8 --buildOnly
trtexec --deploy=GoogleNet_N2.prototxt --output=prob --batch=2 --saveEngine=g2.trt --int8 --buildOnly
  • 保存的引擎可以尝试找到低于 2 ms 的组合批次/流,以最大化吞吐量:
trtexec --loadEngine=g1.trt --batch=1 --streams=2
trtexec --loadEngine=g1.trt --batch=1 --streams=3
trtexec --loadEngine=g1.trt --batch=1 --streams=4
trtexec --loadEngine=g2.trt --batch=2 --streams=2

python调用 TensorRT模型的推理

推理依旧分为动态尺寸的和固定尺寸的,动态推理这一块C++版本的资料比较多,python接口的比较少,固定尺寸的推理官方也有demo,分为异步同步推理。

python推理接收numpy格式的数据输入。

动态推断

import tensorrt as trt
import pycuda.driver as cuda
#import pycuda.driver as cuda2
import pycuda.autoinit
import numpy as np
import cv2
def load_engine(engine_path):
    #TRT_LOGGER = trt.Logger(trt.Logger.WARNING)  # INFO
    TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())
 
path ='/home/caidou/trt_python/model_1_-1_-1_3.engine'
#这里不以某个具体模型做为推断例子.
 
# 1. 建立模型,构建上下文管理器
engine = load_engine(path)
context = engine.create_execution_context()
context.active_optimization_profile = 0
 
#2. 读取数据,数据处理为可以和网络结构输入对应起来的的shape,数据可增加预处理
imgpath = '/home/caidou/test/aaa.jpg'
image = cv2.imread(imgpath)
image = np.expand_dims(image, 0)  # Add batch dimension.  
 
 
#3.分配内存空间,并进行数据cpu到gpu的拷贝
#动态尺寸,每次都要set一下模型输入的shape,0代表的就是输入,输出根据具体的网络结构而定,可以是0,1,2,3...其中的某个头。
context.set_binding_shape(0, image.shape)
d_input = cuda.mem_alloc(image.nbytes)  #分配输入的内存。
 
 
output_shape = context.get_binding_shape(1) 
buffer = np.empty(output_shape, dtype=np.float32)
d_output = cuda.mem_alloc(buffer.nbytes)    #分配输出内存。
cuda.memcpy_htod(d_input,image)
bindings = [d_input ,d_output]
 
#4.进行推理,并将结果从gpu拷贝到cpu。
context.execute_v2(bindings)  #可异步和同步
cuda.memcpy_dtoh(buffer,d_output)  
output = buffer.reshape(output_shape)
 
#5.对推理结果进行后处理。这里只是举了一个简单例子,可以结合官方静态的yolov3案例完善。

静态推断:

静态推断和动态推断差不多,只不过不需要每次都分配输入和输出的内存空间。

import tensorrt as trt
import pycuda.driver as cuda
#import pycuda.driver as cuda2
import pycuda.autoinit
import numpy as np
import cv2
path ='/home/caidou/trt_python/model_1_4_256_256.engine'
engine = load_engine(path)
imgpath = 'aaa.jpg'
context = engine.create_execution_context()
image1 = cv2.write(imgpath)
image1 = cv2.resize(image1,(256,256))
image2 = image1.copy()
image3 = image1.copy()
image4 = image1.copy()
image = np.concatenate((image1,image2,image3,image4))
image = image.reshape(-1,256,256)
 
# image = np.expand_dims(image, axis=1)
image = image.astype(np.float32)
 
image = image.ravel()#数据平铺
outshape= context.get_binding_shape(1) 
output = np.empty((outshape), dtype=np.float32)
d_input = cuda.mem_alloc(1 * image.size * image.dtype.itemsize)
d_output = cuda.mem_alloc(1*output.size * output.dtype.itemsize)
bindings = [int(d_input), int(d_output)]
stream = cuda.Stream()
for i in tqdm.tqdm(range(600)):
    cuda.memcpy_htod(d_input,image)
    context.execute_v2(bindings)
    cuda.memcpy_dtoh(output, d_output)

更新:

TorchScript—模型部署

摘自:https://zhuanlan.zhihu.com/p/486914187

官网:https://pytorch.org/docs/stable/jit.html

PyTorch 无疑是现在最成功的深度学习训练框架之一,是各种顶会顶刊论文实验的大热门。比起其他的框架,PyTorch 最大的卖点是它对动态网络的支持,比其他需要构建静态网络的框架拥有更低的学习成本。PyTorch 源码 Readme 中还专门为此做了一张动态图:


对研究员而言, PyTorch 能极大地提高想 idea、做实验、发论文的效率,是训练框架中的豪杰,但是它不适合部署。动态建图带来的优势对于性能要求更高的应用场景而言更像是缺点,非固定的网络结构给网络结构分析并进行优化带来了困难,多数参数都能以 Tensor 形式传输也让资源分配变成一件闹心的事。另外由于图是由 python 代码构建的,一方面部署要依赖 python 环境,另一方面模型也毫无保密性可言。

而 TorchScript 就是为了解决这个问题而诞生的工具。包括代码的追踪及解析、中间表示的生成、模型优化、序列化等各种功能,可以说是覆盖了模型部署的方方面面。

TorchScript

动态图模型通过牺牲一些高级特性来换取易用性,那到底 JIT 有哪些特性,在什么情况下不得不用到 JIT 呢?下面主要通过介绍 TorchScript(PyTorch 的 JIT 实现)来分析 JIT 到底带来了哪些好处。

  1. 模型部署

PyTorch 的 1.0 版本发布的最核心的两个新特性就是 JIT 和 C++ API,这两个特性一起发布不是没有道理的,JIT 是 Python 和 C++ 的桥梁,我们可以使用 Python 训练模型,然后通过 JIT 将模型转为语言无关的模块,从而让 C++ 可以非常方便得调用,从此「使用 Python 训练模型,使用 C++ 将模型部署到生产环境」对 PyTorch 来说成为了一件很容易的事。而因为使用了 C++,我们现在几乎可以把 PyTorch 模型部署到任意平台和设备上:树莓派、iOS、Android 等等…

2. 性能提升

既然是为部署生产所提供的特性,那免不了在性能上面做了极大的优化,如果推断的场景对性能要求高,则可以考虑将模型(torch.nn.Module)转换为 TorchScript Module,再进行推断。

3. 模型可视化

TensorFlow 或 Keras 对模型可视化工具(TensorBoard等)非常友好,因为本身就是静态图的编程模型,在模型定义好后整个模型的结构和正向逻辑就已经清楚了;但 PyTorch 本身是不支持的,所以 PyTorch 模型在可视化上一直表现得不好,但 JIT 改善了这一情况。现在可以使用 JIT 的 trace 功能来得到 PyTorch 模型针对某一输入的正向逻辑,通过正向逻辑可以得到模型大致的结构,但如果在 `forward` 方法中有很多条件控制语句,这依然不是一个好的方法,所以 PyTorch JIT 还提供了 Scripting 的方式。

TorchScript Module 的两种生成方式

1. 编码(Scripting)

可以直接使用 TorchScript Language 来定义一个 PyTorch JIT Module,然后用 torch.jit.script 来将他转换成 TorchScript Module 并保存成文件。而 TorchScript Language 本身也是 Python 代码,所以可以直接写在 Python 文件中。

使用 TorchScript Language 就如同使用 TensorFlow 一样,需要前定义好完整的图。对于 TensorFlow 我们知道不能直接使用 Python 中的 if 等语句来做条件控制,而是需要用 tf.cond,但对于 TorchScript 我们依然能够直接使用 if 和 for 等条件控制语句,所以即使是在静态图上,PyTorch 依然秉承了「易用」的特性。TorchScript Language 是静态类型的 Python 子集,静态类型也是用了 Python 3 的 typing 模块来实现,所以写 TorchScript Language 的体验也跟 Python 一模一样,只是某些 Python 特性无法使用(因为是子集),可以通过 TorchScript Language Reference 来查看和原生 Python 的异同。

理论上,使用 Scripting 的方式定义的 TorchScript Module 对模型可视化工具非常友好,因为已经提前定义了整个图结构。

2. 追踪(Tracing)

使用 TorchScript Module 的更简单的办法是使用 Tracing,Tracing 可以直接将 PyTorch 模型(torch.nn.Module)转换成 TorchScript Module。「追踪」顾名思义,就是需要提供一个「输入」来让模型 forward 一遍,以通过该输入的流转路径,获得图的结构。这种方式对于 forward 逻辑简单的模型来说非常实用,但如果 forward 里面本身夹杂了很多流程控制语句,则可能会有问题,因为同一个输入不可能遍历到所有的逻辑分枝。

此外,还可以混合使用上面两种方式。

模型转换

作为模型部署的一个范式,通常我们都需要生成一个模型的中间表示(IR),这个 IR 拥有相对固定的图结构,所以更容易优化,让我们看一个例子:

import torch 
from torchvision.models import resnet18 
 
# 使用PyTorch model zoo中的resnet18作为例子 
model = resnet18() 
model.eval() 
 
# 通过trace的方法生成IR需要一个输入样例 
dummy_input = torch.rand(1, 3, 224, 224) 
 
# IR生成 
with torch.no_grad(): 
    jit_model = torch.jit.trace(model, dummy_input) 

JIT 是一种概念,全称是 Just In Time Compilation,中文译为「即时编译」,是一种程序优化的方法

到这里就将 PyTorch 的模型转换成了 TorchScript 的 IR。这里我们使用了 trace 模式来生成 IR,所谓 trace 指的是进行一次模型推理,在推理的过程中记录所有经过的计算,将这些记录整合成计算图。

那么这个 IR 中到底都有些什么呢?我们可以可视化一下其中的 layer1 看看:

jit_layer1 = jit_model.layer1 
print(jit_layer1.graph) 
 
# graph(%self.6 : __torch__.torch.nn.modules.container.Sequential, 
#       %4 : Float(1, 64, 56, 56, strides=[200704, 3136, 56, 1], requires_grad=0, device=cpu)): 
#   %1 : __torch__.torchvision.models.resnet.___torch_mangle_10.BasicBlock = prim::GetAttr[name="1"](%self.6) 
#   %2 : __torch__.torchvision.models.resnet.BasicBlock = prim::GetAttr[name="0"](%self.6) 
#   %6 : Tensor = prim::CallMethod[name="forward"](%2, %4) 
#   %7 : Tensor = prim::CallMethod[name="forward"](%1, %6) 
#   return (%7) 

是不是有点摸不着头脑?TorchScript 有它自己对于 Graph 以及其中元素的定义,对于第一次接触的人来说可能比较陌生,但是没关系,我们还有另一种可视化方式:

print(jit_layer1.code) 
 
# def forward(self, 
#     argument_1: Tensor) -> Tensor: 
#   _0 = getattr(self, "1") 
#   _1 = (getattr(self, "0")).forward(argument_1, ) 
#   return (_0).forward(_1, ) 

没错,就是代码!TorchScript 的 IR 是可以还原成 python 代码的,如果你生成了一个 TorchScript 模型并且想知道它的内容对不对,那么可以通过这样的方式来做一些简单的检查。

刚才的例子中我们使用 trace 的方法生成IR。除了 trace 之外,PyTorch 还提供了另一种生成 TorchScript 模型的方法:script。这种方式会直接解析网络定义的 python 代码,生成抽象语法树 AST,因此这种方法可以解决一些 trace 无法解决的问题,比如对 branch/loop 等数据流控制语句的建图。script方式的建图有很多有趣的特性,会在未来的分享中做专题分析,敬请期待。

模型优化

聪明的同学可能发现了,上面的可视化中只有resnet18forward的部分,其中的子模块信息是不是丢失了呢?如果没有丢失,那么怎么样才能确定子模块的内容是否正确呢?别担心,还记得我们说过 TorchScript 支持对网络的优化吗,这里我们就可以用一个pass解决这个问题:

# 调用inline pass,对graph做变换 
torch._C._jit_pass_inline(jit_layer1.graph) 
print(jit_layer1.code) 
 
# def forward(self, 
#     argument_1: Tensor) -> Tensor: 
#   _0 = getattr(self, "1") 
#   _1 = getattr(self, "0") 
#   _2 = _1.bn2 
#   _3 = _1.conv2 
#   _4 = _1.bn1 
#   input = torch._convolution(argument_1, _1.conv1.weight, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1, False, False, True, True) 
#   _5 = _4.running_var 
#   _6 = _4.running_mean 
#   _7 = _4.bias 
#   input0 = torch.batch_norm(input, _4.weight, _7, _6, _5, False, 0.10000000000000001, 1.0000000000000001e-05, True) 
#   input1 = torch.relu_(input0) 
#   input2 = torch._convolution(input1, _3.weight, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1, False, False, True, True) 
#   _8 = _2.running_var 
#   _9 = _2.running_mean 
#   _10 = _2.bias 
#   out = torch.batch_norm(input2, _2.weight, _10, _9, _8, False, 0.10000000000000001, 1.0000000000000001e-05, True) 
#   input3 = torch.add_(out, argument_1, alpha=1) 
#   input4 = torch.relu_(input3) 
#   _11 = _0.bn2 
#   _12 = _0.conv2 
#   _13 = _0.bn1 
#   input5 = torch._convolution(input4, _0.conv1.weight, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1, False, False, True, True) 
#   _14 = _13.running_var 
#   _15 = _13.running_mean 
#   _16 = _13.bias 
#   input6 = torch.batch_norm(input5, _13.weight, _16, _15, _14, False, 0.10000000000000001, 1.0000000000000001e-05, True) 
#   input7 = torch.relu_(input6) 
#   input8 = torch._convolution(input7, _12.weight, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1, False, False, True, True) 
#   _17 = _11.running_var 
#   _18 = _11.running_mean 
#   _19 = _11.bias 
#   out0 = torch.batch_norm(input8, _11.weight, _19, _18, _17, False, 0.10000000000000001, 1.0000000000000001e-05, True) 
#   input9 = torch.add_(out0, input4, alpha=1) 
#   return torch.relu_(input9) 

这里我们就能看到卷积、batch_norm、relu等熟悉的算子了。

上面代码中我们使用了一个名为inlinepass,将所有子模块进行内联,这样我们就能看见更完整的推理代码。pass是一个来源于编译原理的概念,一个 TorchScript 的 pass 会接收一个图,遍历图中所有元素进行某种变换,生成一个新的图。我们这里用到的inline起到的作用就是将模块调用展开,尽管这样做并不能直接影响执行效率,但是它其实是很多其他pass的基础。PyTorch 中定义了非常多的 pass 来解决各种优化任务,未来我们会做一些更详细的介绍。

序列化

不管是哪种方法创建的 TorchScript 都可以进行序列化,比如:

# 将模型序列化 
jit_model.save('jit_model.pth') 
# 加载序列化后的模型 
jit_model = torch.jit.load('jit_model.pth') 

序列化后的模型不再与 python 相关,可以被部署到各种平台上。

PyTorch 提供了可以用于 TorchScript 模型推理的 c++ API,序列化后的模型终于可以不依赖 python 进行推理了:

// 加载生成的torchscript模型 
auto module = torch::jit::load('jit_model.pth'); 
// 根据任务需求读取数据 
std::vector<torch::jit::IValue> inputs = ...; 
// 计算推理结果 
auto output = module.forward(inputs).toTensor(); 

与 torch.onnx 的关系:ONNX 是业界广泛使用的一种神经网络中间表示,PyTorch 自然也对 ONNX 提供了支持。torch.onnx.export函数可以帮助我们把 PyTorch 模型转换成 ONNX 模型,这个函数会使用 trace 的方式记录 PyTorch 的推理过程。聪明的同学可能已经想到了,没错,ONNX 的导出,使用的正是 TorchScript 的 trace 工具。具体步骤如下:

  1. 使用 trace 的方式先生成一个 TorchScipt 模型,如果你转换的本身就是 TorchScript 模型,则可以跳过这一步。
  2. 使用许多 pass 对 1 中生成的模型进行变换,其中对 ONNX 导出最重要的一个 pass 就是ToONNX,这个 pass 会进行一个映射,将 TorchScript 中primaten空间下的算子映射到onnx空间下的算子。
  3. 使用 ONNX 的 proto 格式对模型进行序列化,完成 ONNX 的导出。

ONNX 模型的修改与调试

转自:模型部署入门教程(五):ONNX 模型的修改与调试

一直以来,我们都是通过 PyTorch 来导出 ONNX 模型的,基本没有单独探究过 ONNX 模型的构造知识。
不知道大家会不会有这样一些疑问:ONNX 模型在底层是用什么格式存储的?如何不依赖深度学习框架,只用 ONNX 的 API 来构造一个 ONNX 模型?如果没有源代码,只有一个 ONNX 模型,该如何对这个模型进行调试?别急,今天我们就来为大家一一揭晓。
在这期教程里,我们将围绕 ONNX 这一套神经网络定义标准本身,探究 ONNX 模型的构造、读取、子模型提取、调试。首先,我们会学习 ONNX 的底层表示方式。之后,我们会用 ONNX API 构造和读取模型。最后,我们会利用 ONNX 提供的子模型提取功能,学习如何调试 ONNX 模型。

ONNX 的底层实现

ONNX 的存储格式

ONNX 在底层是用 Protobuf 定义的。Protobuf,全称 Protocol Buffer,是 Google 提出的一套表示和序列化数据的机制。使用 Protobuf 时,用户需要先写一份数据定义文件,再根据这份定义文件把数据存储进一份二进制文件。可以说,数据定义文件就是数据类,二进制文件就是数据类的实例。
这里给出一个 Protobuf 数据定义文件的例子:

message Person { 
  required string name = 1; 
  required int32 id = 2; 
  optional string email = 3; 
} 

这段定义表示在 Person 这种数据类型中,必须包含 nameid 这两个字段,选择性包含 email字段。根据这份定义文件,用户就可以选择一种编程语言,定义一个含有成员变量 nameidemail 的 Person 类,把这个类的某个实例用 Protobuf 存储成二进制文件;反之,用户也可以用二进制文件和对应的数据定义文件,读取出一个 Person 类的实例。
而对于 ONNX ,Protobuf 的数据定义文件在其开源库,这些文件定义了神经网络中模型、节点、张量的数据类型规范;而二进制文件就是我们熟悉的“.onnx”文件,每一个 onnx 文件按照数据定义规范,存储了一个神经网络的所有相关数据。直接用 Protobuf 生成 ONNX 模型还是比较麻烦的。幸运的是,ONNX 提供了很多实用 API,我们可以在完全不了解 Protobuf 的前提下,构造和读取 ONNX 模型。


ONNX 的结构定义


在用 API 对 ONNX 模型进行操作之前,我们还需要先了解一下 ONNX 的结构定义规则,学习一下 ONNX 在 Protobuf 定义文件里是怎样描述一个神经网络的。
回想一下,神经网络本质上是一个计算图。计算图的节点是算子,边是参与运算的张量。而通过可视化 ONNX 模型,我们知道 ONNX 记录了所有算子节点的属性信息,并把参与运算的张量信息存储在算子节点的输入输出信息中。事实上,ONNX 模型的结构可以用类图大致表示如下:

如图所示,一个 ONNX 模型可以用 ModelProto 类表示。ModelProto 包含了版本、创建者等日志信息,还包含了存储计算图结构的 graphGraphProto 类则由输入张量信息、输出张量信息、节点信息组成。张量信息 ValueInfoProto 类包括张量名、基本数据类型、形状。节点信息 NodeProto 类包含了算子名、算子输入张量名、算子输出张量名。
让我们来看一个具体的例子。假如我们有一个描述 output=a*x+b 的 ONNX 模型 model,用 print(model) 可以输出以下内容:

ir_version: 8 
graph { 
  node { 
    input: "a" 
    input: "x" 
    output: "c" 
    op_type: "Mul" 
  } 
  node { 
    input: "c" 
    input: "b" 
    output: "output" 
    op_type: "Add" 
  } 
  name: "linear_func" 
  input { 
    name: "a" 
    type { 
      tensor_type { 
        elem_type: 1 
        shape { 
          dim {dim_value: 10} 
          dim {dim_value: 10} 
        } 
      } 
    } 
  } 
  input { 
    name: "x" 
    type { 
      tensor_type { 
        elem_type: 1 
        shape { 
          dim {dim_value: 10} 
          dim {dim_value: 10} 
        } 
      } 
    } 
  } 
  input { 
    name: "b" 
    type { 
      tensor_type { 
        elem_type: 1 
        shape { 
          dim {dim_value: 10} 
          dim {dim_value: 10} 
        } 
      } 
    } 
  } 
  output { 
    name: "output" 
    type { 
      tensor_type { 
        elem_type: 1 
        shape { 
          dim { dim_value: 10} 
          dim { dim_value: 10} 
        } 
      } 
    } 
  } 
} 
opset_import {version: 15} 

对应上文中的类图,这个模型的信息由 ir_versionopset_import 等全局信息和 graph 图信息组成。而 graph 包含一个乘法节点、一个加法节点、三个输入张量 a, x, b 以及一个输出张量 output。在下一节里,我们会用 API 构造出这个模型,并输出这段结果。

读写 ONNX 模型

构造 ONNX 模型


在上一小节中,我们知道了 ONNX 模型是按以下的结构组织起来的:

  • ModelProto
    • GraphProto
      • NodeProto
      • ValueInfoProto

现在,让我们抛开 PyTorch,尝试完全用 ONNX 的 Python API 构造一个描述线性函数 output=a*x+b 的 ONNX 模型。我们将根据上面的结构,自底向上地构造这个模型。
首先,我们可以用 helper.make_tensor_value_info 构造出一个描述张量信息的 ValueInfoProto 对象。如前面的类图所示,我们要传入张量名、张量的基本数据类型、张量形状这三个信息。在 ONNX 中,不管是输入张量还是输出张量,它们的表示方式都是一样的。因此,这里我们用类似的方式为三个输入 a, x, b 和一个输出 output 构造 ValueInfoProto 对象。如下面的代码所示:

import onnx 
from onnx import helper 
from onnx import TensorProto 
 
a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [10, 10]) 
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 10]) 
b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [10, 10]) 
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [10, 10]) 

之后,我们要构造算子节点信息 NodeProto,这可以通过在 helper.make_node 中传入算子类型、输入算子名、输出算子名这三个信息来实现。我们这里先构造了描述 c=a*x 的乘法节点,再构造了 output=c+b 的加法节点。如下面的代码所示:

mul = helper.make_node('Mul', ['a', 'x'], ['c']) 
add = helper.make_node('Add', ['c', 'b'], ['output']) 

在计算机中,图一般是用一个节点集和一个边集表示的。而 ONNX 巧妙地把边的信息保存在了节点信息里,省去了保存边集的步骤。在 ONNX 中,如果某节点的输入名和之前某节点的输出名相同,就默认这两个节点是相连的。如上面的例子所示:Mul 节点定义了输出 cAdd 节点定义了输入 c,则 Mul 节点和 Add 节点是相连的。
正是因为有这种边的隐式定义规则,所以 ONNX 对节点的输入有一定的要求:一个节点的输入,要么是整个模型的输入,要么是之前某个节点的输出。如果我们把 a, x, b 中的某个输入节点从计算图中拿出(这个操作会在之后的代码中介绍),或者把 Mul 的输出从 c 改成 d,则最终的 ONNX 模型都是不满足标准的。

一个不满足标准的 ONNX 模型可能无法被推理引擎正确识别。ONNX 提供了 API onnx.checker.check_model 来判断一个 ONNX 模型是否满足标准。

接下来,我们用 helper.make_graph 来构造计算图 GraphProtohelper.make_graph 函数需要传入节点、图名称、输入张量信息、输出张量信息这 4 个参数。如下面的代码所示,我们把之前构造出来的 NodeProto 对象和 ValueInfoProto 对象按照顺序传入即可。

graph = helper.make_graph([mul, add], 'linear_func', [a, x, b], [output]) 

这里 make_graph 的节点参数有一个要求:计算图的节点必须以拓扑序给出。

拓扑序是与有向图的相关的数学概念。如果按拓扑序遍历所有节点的话,能保证每个节点的输入都能在之前节点的输出里找到(对于 ONNX 模型,我们把计算图的输入张量也看成“之前的输出”)。

如果对这个概念不熟也没有关系,我们以刚刚构造出来的这个计算图为研究对象,通过下图展示的两个例子来直观理解拓扑序。

这里我们只关注 Mul 和 Add 节点以及它们之间的边 c。在情况 1 中:如果我们的节点以 [Mul, Add] 顺序给出,那么遍历到 Add 时,它的输入 c 可以在之前的Mul的输出中找到。但是,如情况 2 所示:如果我们的节点以 [Add, Mul] 的顺序给出,那么 Add 就找不到输入边,计算图也无法成功构造出来了。这里的 [Mul, Add] 就是符合有向图的拓扑序的,而 [Add, Mul] 则不满足。

最后,我们用 helper.make_model 把计算图 GraphProto 封装进模型 ModelProto 里,一个 ONNX 模型就构造完成了。make_model 函数中还可以添加模型制作者、版本等信息,为了简单起见,我们没有添加额外的信息。如下面的代码所示:

model = helper.make_model(graph) 

构造完模型之后,我们用下面这三行代码来检查模型正确性、把模型以文本形式输出、存储到一个 “.onnx” 文件里。这里用 onnx.checker.check_model 来检查模型是否满足 ONNX 标准是必要的,因为无论模型是否满足标准,ONNX 都允许我们用 onnx.save 存储模型。我们肯定不希望生成一个不满足标准的模型。

onnx.checker.check_model(model) 
print(model) 
onnx.save(model, 'linear_func.onnx') 

成功执行这些代码的话,程序会以文本格式输出模型的信息,其内容应该和我们在上一节展示的输出一样。
整理一下,用 ONNX Python API 构造模型的代码如下:

import onnx 
from onnx import helper 
from onnx import TensorProto 
 
# input and output 
a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [10, 10]) 
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 10]) 
b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [10, 10]) 
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [10, 10]) 
 
# Mul 
mul = helper.make_node('Mul', ['a', 'x'], ['c']) 
 
# Add 
add = helper.make_node('Add', ['c', 'b'], ['output']) 
 
# graph and model 
graph = helper.make_graph([mul, add], 'linear_func', [a, x, b], [output]) 
model = helper.make_model(graph) 
 
# save model 
onnx.checker.check_model(model) 
print(model) 
onnx.save(model, 'linear_func.onnx') 

老规矩,我们可以用 ONNX Runtime 运行模型,来看看模型是否正确:

import onnxruntime 
import numpy as np 
 
sess = onnxruntime.InferenceSession('linear_func.onnx') 
a = np.random.rand(10, 10).astype(np.float32) 
b = np.random.rand(10, 10).astype(np.float32) 
x = np.random.rand(10, 10).astype(np.float32) 
 
output = sess.run(['output'], {'a': a, 'b': b, 'x': x})[0] 
 
assert np.allclose(output, a * x + b) 

一切顺利的话,这段代码不会有任何报错信息。这说明我们的模型等价于执行 a * x + b 这个计算。


读取并修改 ONNX 模型

通过用 API 构造 ONNX 模型,我们已经彻底搞懂了 ONNX 由哪些模块组成。现在,让我们看看该如何读取现有的”.onnx”文件并从中提取模型信息。
首先,我们可以用下面的代码读取一个 ONNX 模型:

import onnx 
model = onnx.load('linear_func.onnx') 
print(model) 

之前在输出模型时,我们传给 onnx.save 的是一个 ModelProto 的对象。同理,用上面的 onnx.load 读取 ONNX 模型时,我们收获的也是一个 ModelProto 的对象。输出这个对象后,我们应该得到和之前完全相同的输出。
接下来,我们来看看怎么把图 GraphProto、节点 NodeProto、张量信息 ValueInfoProto 读取出来:

graph = model.graph 
node = graph.node 
input = graph.input 
output = graph.output 
print(node) 
print(input) 
print(output) 

使用如上这些代码,我们可以分别访问模型的图、节点、张量信息。这里大家或许会有疑问:该怎样找出 graph.node,graph.input 中 node, input 这些属性名称呢?其实,属性的名称就写在每个对象的输出里。我们以 print(node) 的输出为例:

[input: "a" 
input: "x" 
output: "c" 
op_type: "Mul" 
, input: "c" 
input: "b" 
output: "output" 
op_type: "Add" 
] 

在这段输出中,我们能看出 node 其实就是一个列表,列表中的对象有属性 input, output, op_type(这里 input 也是一个列表,它包含的两个元素都显示出来了)。我们可以用下面的代码来获取 node 里第一个节点 Mul 的属性:

node_0 = node[0] 
node_0_inputs = node_0.input 
node_0_outputs = node_0.output 
input_0 = node_0_inputs[0] 
input_1 = node_0_inputs[1] 
output = node_0_outputs[0] 
op_type = node_0.op_type 
 
print(input_0) 
print(input_1) 
print(output) 
print(op_type) 
 
# Output 
""" 
a 
x 
c 
Mul 
""" 

当我们想知道 ONNX 模型某数据对象有哪些属性时,我们不必去翻 ONNX 文档,只需要先把数据对象输出一下,然后在输出结果找出属性名即可。
读取 ONNX 模型的信息后,修改 ONNX 模型就是一件很轻松的事了。我们既可以按照上一小节的模型构造方法,新建节点和张量信息,与原有模型组合成一个新的模型,也可以在不违反 ONNX 规范的前提下直接修改某个数据对象的属性。
这里我们来看一个直接修改模型属性的例子:

import onnx 
model = onnx.load('linear_func.onnx') 
 
node = model.graph.node 
node[1].op_type = 'Sub' 
 
onnx.checker.check_model(model) 
onnx.save(model, 'linear_func_2.onnx') 

在读入之前的 linear_func.onnx 模型后,我们可以直接修改第二个节点的类型 node[1].op_type,把加法变成减法。这样,我们的模型描述的是 a * x - b 这个线性函数。大家感兴趣的话,可以用 ONNX Runtime 运行新模型 linear_func_2.onnx,来验证一下它和 a * x - b 是否等价。

调试 ONNX 模型

在实际部署中,如果用深度学习框架导出的 ONNX 模型出了问题,一般要通过修改框架的代码来解决,而不会从 ONNX 入手,我们把 ONNX 模型当成一个不可修改的黑盒看待。
现在,我们已经深入学习了 ONNX 的原理,可以尝试对 ONNX 模型本身进行调试了。在这一节里,让我们看看该如何巧妙利用 ONNX 提供的子模型提取功能,对 ONNX 模型进行调试。

子模型提取

ONNX 官方为开发者提供了子模型提取(extract)的功能。子模型提取,顾名思义,就是从一个给定的 ONNX 模型中,拿出一个子模型。这个子模型的节点集、边集都是原模型中对应集合的子集。让我们来用 PyTorch 导出一个复杂一点的 ONNX 模型,并在它的基础上执行提取操作:

import torch 
 
class Model(torch.nn.Module): 
 
    def __init__(self): 
        super().__init__() 
        self.convs1 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3), 
                                          torch.nn.Conv2d(3, 3, 3), 
                                          torch.nn.Conv2d(3, 3, 3)) 
        self.convs2 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3), 
                                          torch.nn.Conv2d(3, 3, 3)) 
        self.convs3 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3), 
                                          torch.nn.Conv2d(3, 3, 3)) 
        self.convs4 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3), 
                                          torch.nn.Conv2d(3, 3, 3), 
                                          torch.nn.Conv2d(3, 3, 3)) 
    def forward(self, x): 
        x = self.convs1(x) 
        x1 = self.convs2(x) 
        x2 = self.convs3(x) 
        x = x1 + x2 
        x = self.convs4(x) 
        return x 
 
model = Model() 
input = torch.randn(1, 3, 20, 20) 
 
torch.onnx.export(model, input, 'whole_model.onnx') 


这个模型的可视化结果如下图所示(提取子模型需要输入边的序号,为了大家方面阅读,这幅图标出了之后要用到的边的序号):

在前面的章节中,我们学过,ONNX 的边用同名张量表示的。也就是说,这里的边序号,实际上是前一个节点的输出张量序号和后一个节点的输入张量序号。由于这个模型是用 PyTorch 导出的,这些张量序号都是 PyTorch 自动生成的。


接着,我们可以下面的代码提取出一个子模型:

import onnx  
 
onnx.utils.extract_model('whole_model.onnx', 'partial_model.onnx', ['22'], ['28']) 

子模型的可视化结果如下图所示:

通过观察代码和输出图,应该不难猜出这段代码的作用是把原计算图从边 22 到边 28 的子图提取出来,并组成一个子模型。onnx.utils.extract_model 就是完成子模型提取的函数,它的参数分别是原模型路径、输出模型路径、子模型的输入边(输入张量)、子模型的输出边(输出张量)。
直观地来看,子模型提取就是把输入边到输出边之间的全部节点都取出来。那么,这个功能在使用上有什么限制呢?基于 whole_model.onnx, 我们来看一看三个子模型提取的示例。

添加额外输出

我们在提取时新设定了一个输出张量,如下面的代码所示:

onnx.utils.extract_model('whole_model.onnx', 'submodel_1.onnx', ['22'], ['27', '31']) 

我们可以看到子模型会添加一条把张量输出的新边,如下图所示:

添加冗余输入

如果我们还是像开始一样提取边 22 到边 28 之间的子模型,但是多添加了一个输入 input.1,那么提取出的子模型会有一个冗余的输入 input.1,如下面的代码所示:

onnx.utils.extract_model('whole_model.onnx', 'submodel_2.onnx', ['22', 'input.1'], ['28']) 

从下图可以看到:无论给这个输入传入什么值,都不会影响子模型的输出。可以认为如果只用子模型的部分输入就能得到输出,那么那些”较早“的多出来的输入就是冗余的。

输入信息不足

这次,我们尝试提取的子模型输入是边 24,输出是边 28。如下面的代码和图所示:

# Error 
onnx.utils.extract_model('whole_model.onnx', 'submodel_3.onnx', ['24'], ['28']) 

从图中可以看出,想通过边 24 计算边 28 的结果,至少还需要输入边 26,或者更上面的边。仅凭借边 24 是无法计算出边 28 的结果的,因此这样提取子模型会报错。

通过上面几个使用示例,我们可以整理出子模型提取的实现原理:新建一个模型,把给定的输入和输出填入。之后把图的所有有向边反向,从输出边开始遍历节点,碰到输入边则停止,把这样遍历得到的节点做为子模型的节点。
如果还没有彻底弄懂这个提取原理,没关系,我们只要尽量保证在填写子模型的输入输出时,让输出恰好可以由输入决定即可。

输出 ONNX 中间节点的值

在使用 ONNX 模型时,最常见的一个需求是能够用推理引擎输出中间节点的值。这多见于深度学习框架模型和 ONNX 模型的精度对齐中,因为只要能够输出中间节点的值,就能定位到精度出现偏差的算子。我们来看看如何用子模型提取实现这一任务。
在刚刚的第一个子模型提取示例中,我们添加了一条原来模型中不存在的输出边。用同样的原理,我们可以在保持原有输入输出不变的同时,新增加一些输出,提取出一个能输出中间节点的”子模型“。例如:

 onnx.utils.extract_model('whole_model.onnx', 'more_output_model.onnx', ['input.1'], ['31', '23', '25', '27'])

在这个子模型中,我们在保持原有的输入 input.1,输出 31 的同时,把其他几个边加入了输出中。如下图所示:

这样,用 ONNX Runtime 运行 more_output_model.onnx 这个模型时,我们就能得到更多的输出了。
为了方便调试,我们还可以把原模型拆分成多个互不相交的子模型。这样,在每次调试时,可以只对原模型的部分子模块调试。比如:

onnx.utils.extract_model('whole_model.onnx', 'debug_model_1.onnx', ['input.1'], ['23']) 
onnx.utils.extract_model('whole_model.onnx', 'debug_model_2.onnx', ['23'], ['25']) 
onnx.utils.extract_model('whole_model.onnx', 'debug_model_3.onnx', ['23'], ['27']) 
onnx.utils.extract_model('whole_model.onnx', 'debug_model_4.onnx', ['25', '27'], ['31']) 

在这个例子中,我们把原来较为复杂的模型拆成了四个较为简单的子模型,如下图所示。在调试时,我们可以先调试顶层的子模型,确认顶层子模型无误后,把它的输出做为后面子模型的输入。
比如对于这些子模型,我们可以先调试第一个子模型,并存储输出 23。之后把张量 23 做为第二个和第三个子模型的输入,调试这两个模型。最后用同样方法调试第四个子模型。可以说,有了子模型提取功能,哪怕是面对一个庞大的模型,我们也能够从中提取出有问题的子模块,细致地只对这个子模块调试。

子模型提取固然是一个便利的 ONNX 调试工具。但是,在实际的情况中,我们一般是用 PyTorch 等框架导出 ONNX 模型。这里有两个问题:

  1. 一旦 PyTorch 模型改变,ONNX 模型的边序号也会改变。这样每次提取同样的子模块时都要重新去 ONNX 模型里查序号,如此繁琐的调试方法是不会在实践中采用的。
  2. 即使我们能保证 ONNX 的边序号不发生改变,我们也难以把 PyTorch 代码和 ONNX 节点对应起来——当模型结构变得十分复杂时,要识别 ONNX 中每个节点的含义是不可能的。

PyTorch 中支持更多 ONNX 算子

学习了 PyTorch 转 ONNX 的方法,可以发现 PyTorch 对 ONNX 的支持还不错。但在实际的部署过程中,难免碰到模型无法用原生 PyTorch 算子表示的情况。这个时候,我们就得考虑扩充 PyTorch,即在 PyTorch 中支持更多 ONNX 算子。

而要使 PyTorch 算子顺利转换到 ONNX ,我们需要保证以下三个环节都不出错:

  • 算子在 PyTorch 中有实现
  • 有把该 PyTorch 算子映射成一个或多个 ONNX 算子的方法
  • ONNX 有相应的算子

可在实际部署中,这三部分的内容都可能有所缺失。其中最坏的情况是:我们定义了一个全新的算子,它不仅缺少 PyTorch 实现,还缺少 PyTorch 到 ONNX 的映射关系。但所谓车到山前必有路,对于这三个环节,我们也分别都有以下的添加支持的方法:

  • PyTorch 算子
    • 组合现有算子
    • 添加 TorchScript 算子
    • 添加普通 C++ 拓展算子
  • 映射方法
    • 为 ATen 算子添加符号函数
    • 为 TorchScript 算子添加符号函数
    • 封装成 torch.autograd.Function 并添加符号函数
  • ONNX 算子
    • 使用现有 ONNX 算子
    • 定义新 ONNX 算子

那么面对不同的情况时,就需要我们灵活地选用和组合这些方法。听起来是不是很复杂?别担心,本篇文章中,我们将围绕着三种算子映射方法,学习三个添加算子支持的实例,来理清如何合适地为 PyTorch 算子转 ONNX 算子的三个环节添加支持。

 支持 ATen 算子

实际的部署过程中,我们都有可能会碰到一个最简单的算子缺失问题: 算子在 ATen 中已经实现了,ONNX 中也有相关算子的定义,但是相关算子映射成 ONNX 的规则没有写。在这种情况下,我们只需要为 ATen 算子补充描述映射规则的符号函数就行了。

ATen 是 PyTorch 内置的 C++ 张量计算库,PyTorch 算子在底层绝大多数计算都是用 ATen 实现的。

上期习题中,我们曾经提到了 ONNX 的 Asinh 算子。这个算子在 ATen 中有实现,却缺少了映射到 ONNX 算子的符号函数。在这里,我们来尝试为它补充符号函数,并导出一个包含这个算子的 ONNX 模型。

获取 ATen 中算子接口定义

为了编写符号函数,我们需要获得 asinh 推理接口的输入参数定义。这时,我们要去 torch/_C/_VariableFunctions.pyi 和 torch/nn/functional.pyi 这两个文件中搜索我们刚刚得到的这个算子名。这两个文件是编译 PyTorch 时本地自动生成的文件,里面包含了 ATen 算子的 PyTorch 调用接口。通过搜索,我们可以知道 asinh 在文件 torch/_C/_VariableFunctions.pyi 中,其接口定义为:

def asinh(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ... 

经过这些步骤,我们确认了缺失的算子名为 asinh,它是一个有实现的 ATen 算子。我们还记下了 asinh 的调用接口。接下来,我们要为它补充符号函数,使它在转换成 ONNX 模型时不再报错。

添加符号函数

到目前为止,我们已经多次接触了定义 PyTorch 到 ONNX 映射规则的符号函数了。现在,我们向大家正式介绍一下符号函数。

符号函数,可以看成是 PyTorch 算子类的一个静态方法。在把 PyTorch 模型转换成 ONNX 模型时,各个 PyTorch 算子的符号函数会被依次调用,以完成 PyTorch 算子到 ONNX 算子的转换。符号函数的定义一般如下:

def symbolic(g: torch._C.Graph, input_0: torch._C.Value, input_1: torch._C.Value, ...): 

其中,torch._C.Graph 和 torch._C.Value 都对应 PyTorch 的 C++ 实现里的一些类。我们在这篇文章不深究它们的细节(感兴趣的话可以参考我们的 TorchScript 系列文章中对 trace 机制的解读),只需要知道第一个参数就固定叫 g,它表示和计算图相关的内容;后面的每个参数都表示算子的输入,需要和算子的前向推理接口的输入相同。对于 ATen 算子来说,它们的前向推理接口就是上述两个 .pyi 文件里的函数接口。

g 有一个方法 op。在把 PyTorch 算子转换成 ONNX 算子时,需要在符号函数中调用此方法来为最终的计算图添加一个 ONNX 算子。其定义如下:

def op(name: str, input_0: torch._C.Value, input_1: torch._C.Value, ...) 

其中,第一个参数是算子名称。如果该算子是普通的 ONNX 算子,只需要把它在 ONNX 官方文档里的名称填进去即可(我们稍后再讲其他情况)。

在最简单的情况下,我们只要把 PyTorch 算子的输入用g.op()一一对应到 ONNX 算子上即可,并把g.op()的返回值作为符号函数的返回值。在情况更复杂时,我们转换一个 PyTorch 算子可能要新建若干个 ONNX 算子。

补充完了背景知识,让我们回到 asinh 算子上,来为它编写符号函数。我们先去翻阅一下 ONNX 算子文档,学习一下我们在符号函数里的映射关系 g.op() 里应该怎么写。Asinh 的文档写道:该算子有一个输入 input,一个输出 output,二者的类型都为张量。

到这里,我们已经完成了信息收集环节。我们在上一小节得知了 asinh 的推理接口定义,在这一小节里收集了 ONNX 算子 Asinh 的定义。现在,我们可以用代码来补充这二者的映射关系了。在刚刚导出 asinh 算子的代码中,我们添加以下内容:

from torch.onnx.symbolic_registry import register_op 
 
def asinh_symbolic(g, input, *, out=None): 
    return g.op("Asinh", input) 
 
register_op('asinh', asinh_symbolic, '', 9)  

这里的asinh_symbolic就是asinh的符号函数。从除g以外的第二个输入参数开始,其输入参数应该严格对应它在 ATen 中的定义:

def asinh(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ... 

在符号函数的函数体中,g.op("Asinh", input)则完成了 ONNX 算子的定义。其中,第一个参数"Asinh"是算子在 ONNX 中的名称。至于第二个参数 input,如我们刚刚在文档里所见,这个算子只有一个输入,因此我们只要把符号函数的输入参数 input 对应过去就行。ONNX 的 Asinh 的输出和 ATen 的 asinh 的输出是一致的,因此我们直接把 g.op() 的结果返回即可。

定义完符号函数后,我们要把这个符号函数和原来的 ATen 算子“绑定”起来。这里,我们要用到 register_op 这个 PyTorch API 来完成绑定。如示例所示,只需要一行简单的代码即可把符号函数 asinh_symbolic 绑定到算子 asinh 上:

register_op('asinh', asinh_symbolic, '', 9) 

register_op的第一个参数是目标 ATen 算子名,第二个是要注册的符号函数,这两个参数很好理解。第三个参数是算子的“域”,对于普通 ONNX 算子,直接填空字符串即可。第四个参数表示向哪个算子集版本注册。我们遵照 ONNX 标准,向第 9 号算子集注册。值得注意的是,这里向第 9 号算子集注册,不代表较新的算子集(第 10 号、第 11 号……)都得到了注册。在示例中,我们先只向第 9 号算子集注册。

整理一下,我们最终的代码如下:

import torch 
 
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
 
    def forward(self, x): 
        return torch.asinh(x) 
 
from torch.onnx.symbolic_registry import register_op 
 
def asinh_symbolic(g, input, *, out=None): 
    return g.op("Asinh", input) 
 
register_op('asinh', asinh_symbolic, '', 9) 
 
model = Model() 
input = torch.rand(1, 3, 10, 10) 
torch.onnx.export(model, input, 'asinh.onnx') 
 

成功导出的话,asinh.onnx 应该长这个样子:

测试算子

在完成了一份自定义算子后,我们一定要测试一下算子的正确性。一般我们要用 PyTorch 运行一遍原算子,再用推理引擎(比如 ONNX Runtime)运行一下 ONNX 算子,最后比对两次的运行结果。对于我们刚刚得到的 asinh.onnx,可以用如下代码来验证:

import onnxruntime 
import torch 
import numpy as np 
 
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
 
    def forward(self, x): 
        return torch.asinh(x) 
 
model = Model() 
input = torch.rand(1, 3, 10, 10) 
torch_output = model(input).detach().numpy() 
 
sess = onnxruntime.InferenceSession('asinh.onnx') 
ort_output = sess.run(None, {'0': input.numpy()})[0] 
 
assert np.allclose(torch_output, ort_output) 

在这份代码里,我们用 PyTorch 做了一遍推理,并把结果转成了 numpy 格式。之后,我们又用 ONNX Runtime 对 onnx 文件做了一次推理。

忘了 ONNX Runtime 的调用方法的话,欢迎回顾第一篇教程~

最后,我们使用 np.allclose 来保证两个结果张量的误差在一个可以允许的范围内。一切正常的话,运行这段代码后,assert 所在行不会报错,程序应该没有任何输出。

支持 TorchScript 算子

对于一些比较复杂的运算,仅使用 PyTorch 原生算子是无法实现的。这个时候,就要考虑自定义一个 PyTorch 算子,再把它转换到 ONNX 中了。新增 PyTorch 算子的方法有很多,PyTorch 官方比较推荐的一种做法是添加 TorchScript 算子 。

由于添加算子的方法较繁琐,我们今天跳过新增 TorchScript 算子的内容,以可变形卷积(Deformable Convolution)算子为例,介绍为现有 TorchScript 算子添加 ONNX 支持的方法。

可变形卷积(Deformable Convolution)是在 Torchvision 中实现的 TorchScript 算子,虽然尚未得到广泛支持,但是出现在许多模型中。

有了支持 ATen 算子的经验之后,我们可以知道为算子添加符号函数一般要经过以下几步:

  1. 获取原算子的前向推理接口。
  2. 获取目标 ONNX 算子的定义。
  3. 编写符号函数并绑定。

在为可变形卷积添加符号函数时,我们也可以尝试走一遍这个流程。

使用 TorchScript 算子

和之前一样,我们首先定义一个包含了算子的模型,为之后转换 ONNX 模型做准备。

import torch 
import torchvision 
 
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
        self.conv1 = torch.nn.Conv2d(3, 18, 3) 
        self.conv2 = torchvision.ops.DeformConv2d(3, 3, 3) 
 
    def forward(self, x): 
        return self.conv2(x, self.conv1(x)) 

其中,torchvision.ops.DeformConv2d 就是 Torchvision 中的可变形卷积层。相比于普通卷积,可变形卷积的其他参数都大致相同,唯一的区别就是在推理时需要多输入一个表示偏移量的张量。

然后,我们查询算子的前向推理接口。DeformConv2d 层最终会调用 deform_conv2d 这个算子。我们可以在 torchvision/csrc/ops/deform_conv2d.cpp 中查到该算子的调用接口:

m.def(TORCH_SELECTIVE_SCHEMA( 
      "torchvision::deform_conv2d(Tensor input,  
      Tensor weight,  
      Tensor offset,  
      ...... 
      bool use_mask) -> Tensor")); 

那么接下来,根据之前的经验,我们就是要去 ONNX 官方文档中查找算子的定义了。

自定义 ONNX 算子

很遗憾的是,如果我们去 ONNX 的官方算子页面搜索 “deform”,将搜不出任何内容。目前,ONNX 还没有提供可变形卷积的算子,我们要自己定义一个 ONNX 算子了。

我们在前面讲过,g.op() 是用来定义 ONNX 算子的函数。对于 ONNX 官方定义的算子,g.op() 的第一个参数就是该算子的名称。而对于一个自定义算子,g.op() 的第一个参数是一个带命名空间的算子名,比如:

g.op("custom::deform_conv2d, ...) 

其中,”::”前面的内容就是我们的命名空间。该概念和 C++ 的命名空间类似,是为了防止命名冲突而设定的。如果在 g.op() 里不加前面的命名空间,则算子会被默认成 ONNX 的官方算子。

PyTorch 在运行 g.op() 时会对官方的算子做检查,如果算子名有误,或者算子的输入类型不正确, g.op() 就会报错。为了让我们随心所欲地定义新 ONNX 算子,我们必须设定一个命名空间,给算子取个名,再定义自己的算子。

我们在第一篇教程讲过:ONNX 是一套标准,本身不包括实现。在这里,我们就简略地定义一个 ONNX 可变形卷积算子,而不去写它在某个推理引擎上的实现。在后续的文章中,我们再介绍在各个推理引擎中添加新 ONNX 算子支持的方法。此处,我们只关心如何导出一个包含新 ONNX 算子节点的 onnx 文件。因此,我们可以为新算子编写如下简单的符号函数:

@parse_args("v", "v", "v", "v", "v", "i", "i", "i", "i", "i", "i", "i", "i", "none") 
def symbolic(g,  
        input, 
        weight, 
        offset, 
        mask, 
        bias, 
        stride_h, stride_w, 
        pad_h, pad_w, 
        dil_h, dil_w, 
        n_weight_grps, 
        n_offset_grps, 
        use_mask): 
    return g.op("custom::deform_conv2d", input, offset) 
 

在这个符号函数中,我们以刚刚搜索到的算子输入参数作为符号函数的输入参数,并只用 input 和 offset 来构造一个简单的 ONNX 算子。

这段代码中,最令人疑惑的就是装饰器 @parse_args 了。简单来说,TorchScript 算子的符号函数要求标注出每一个输入参数的类型。比如”v”表示 Torch 库里的 value 类型,一般用于标注张量,而”i”表示 int 类型,”f”表示 float 类型,”none”表示该参数为空。具体的类型含义可以在 torch.onnx.symbolic_helper.py (https://github.com/pytorch/pytorch/blob/master/torch/onnx/symbolic_helper.py)中查看。这里输入参数中的 input, weight, offset, mask, bias 都是张量,所以用”v”表示。后面的其他参数同理。我们不必纠结于 @parse_args 的原理,根据实际情况对符号函数的参数标注类型即可。

有了符号函数后,我们通过如下的方式注册符号函数:

register_custom_op_symbolic("torchvision::deform_conv2d", symbolic, 9) 

和前面的 register_op 类似,注册符号函数时,我们要输入算子名、符号函数、算子集版本。与前面不同的是,这里的算子集版本是最早生效版本,在这里设定版本 9,意味着之后的第 10 号、第 11 号……版本集都能使用这个新算子。

最后,我们完整的模型导出代码如下:

import torch 
import torchvision 
 
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
        self.conv1 = torch.nn.Conv2d(3, 18, 3) 
        self.conv2 = torchvision.ops.DeformConv2d(3, 3, 3) 
 
    def forward(self, x): 
        return self.conv2(x, self.conv1(x)) 
 
from torch.onnx import register_custom_op_symbolic 
from torch.onnx.symbolic_helper import parse_args 
 
@parse_args("v", "v", "v", "v", "v", "i", "i", "i", "i", "i", "i", "i", "i", "none") 
def symbolic(g,  
        input, 
        weight, 
        offset, 
        mask, 
        bias, 
        stride_h, stride_w, 
        pad_h, pad_w, 
        dil_h, dil_w, 
        n_weight_grps, 
        n_offset_grps, 
        use_mask): 
    return g.op("custom::deform_conv2d", input, offset) 
 
register_custom_op_symbolic("torchvision::deform_conv2d", symbolic, 9) 
 
model = Model() 
input = torch.rand(1, 3, 10, 10) 
torch.onnx.export(model, input, 'dcn.onnx') 
 

代码成功运行的话,我们应该能得到如下的 ONNX 模型:

可以看到,我们自定义的 ONNX 算子 deform_conv2d 包含了两个输入,一个输出,和我们预想得一样。

使用 torch.autograd.Function

最后,我们来学习一种简单的为 PyTorch 添加 C++ 算子实现的方法,来代替较为复杂的新增 TorchScript 算子。同时,我们会用 torch.autograd.Function 封装这个新算子。torch.autograd.Function 能完成算子实现和算子调用的隔离。不管算子是怎么实现的,它封装后的使用体验以及 ONNX 导出方法会和原生的 PyTorch 算子一样。这是我们比较推荐的为算子添加 ONNX 支持的方法。

为了应对更复杂的情况,我们来自定义一个奇怪的 my_add 算子。这个算子的输入张量 a, b ,输出 2a + b 的值。我们会先把它在 PyTorch 中实现,再把它导出到 ONNX 中。

为 PyTorch 添加 C++ 拓展

为 PyTorch 添加简单的 C++ 拓展还是很方便的。对于我们定义的 my_add 算子,可以用以下的 C++ 源文件来实现。我们把该文件命名为 “my_add.cpp”:

// my_add.cpp 
 
#include <torch/torch.h> 
 
torch::Tensor my_add(torch::Tensor a, torch::Tensor b) 
{ 
    return 2 * a + b; 
} 
 
PYBIND11_MODULE(my_lib, m) 
{ 
    m.def("my_add", my_add); 
} 

由于在 PyTorch 中添加 C++ 拓展和模型部署关系不大,这里我们仅给出这个简单的示例,并不对其原理做过多讲解。

在这段代码中,torch::Tensor 就是 C++ 中 torch 的张量类型,它的加法和乘法等运算符均已重载。因此,我们可以像对普通标量一样对张量做加法和乘法。

轻松地完成了算子的实现后,我们用 PYBIND11_MODULE 来为 C++ 函数提供 Python 调用接口。这里的 my_lib 是我们未来要在 Python 里导入的模块名。双引号中的 my_add 是 Python 调用接口的名称,这里我们对齐 C++ 函数的名称,依然用 “my_add”这个名字。

之后,我们可以编写如下的 Python 代码并命名为 “setup.py”,来编译刚刚的 C++ 文件:

from setuptools import setup 
from torch.utils import cpp_extension 
 
setup(name='my_add', 
      ext_modules=[cpp_extension.CppExtension('my_lib', ['my_add.cpp'])], 
      cmdclass={'build_ext': cpp_extension.BuildExtension}) 

这段代码使用了 Python 的 setuptools 编译功能和 PyTorch 的 C++ 拓展工具函数,可以编译包含了 torch 库的 C++ 源文件。这里我们需要填写的只有模块名和模块中的源文件名。我们刚刚把模块命名为 my_lib,而源文件只有一个 my_add.cpp,因此拓展模块那一行要写成 ext_modules=[cpp_extension.CppExtension('my_lib', ['my_add.cpp'])],

之后,像处理普通的 Python 包一样执行安装命令,我们的 C++ 代码就会自动编译了。

python setup.py develop 

用 torch.autograd.Function 封装

直接用 Python 接口调用 C++ 函数不太“美观”,一种比较优雅的做法是把这个调用接口封装起来。这里我们用 torch.autograd.Function 来封装算子的底层调用:

import torch 
import my_lib 
class MyAddFunction(torch.autograd.Function): 
 
    @staticmethod 
    def forward(ctx, a, b): 
        return my_lib.my_add(a, b) 
 
    @staticmethod 
    def symbolic(g, a, b): 
        two = g.op("Constant", value_t=torch.tensor([2])) 
        a = g.op('Mul', a, two) 
        return g.op('Add', a, b) 

我们在前面的教程中已经见过 torch.autograd.Function,这里我们正式地对其做一个介绍。Function 类本身表示 PyTorch 的一个可导函数,只要为其定义了前向推理和反向传播的实现,我们就可以把它当成一个普通 PyTorch 函数来使用。

PyTorch 会自动调度该函数,合适地执行前向和反向计算。对模型部署来说,Function 类有一个很好的性质:如果它定义了 symbolic 静态方法,该 Function 在执行 torch.onnx.export() 时就可以根据 symbolic 中定义的规则转换成 ONNX 算子。这个 symbolic 就是前面提到的符号函数,只是它的名称必须是 symbolic 而已。

在 forward 函数中,我们用 my_lib.my_add(a, b) 就可以调用之前写的C++函数了。这里 my_lib 是库名,my_add 是函数名,这两个名字是在前面C++的 PYBIND11_MODULE 中定义的。

在 symbolic 函数中,我们用 g.op() 定义了三个算子:常量、乘法、加法。这里乘法和加法的用法和前面提到的 asinh 一样,只需要根据 ONNX 算子定义规则把输入参数填入即可。而在定义常量算子时,我们要把 PyTorch 张量的值传入 value_t 参数中。

在 ONNX 中,我们需要把新建常量当成一个算子来看待,尽管这个算子并不会以节点的形式出现在 ONNX 模型的可视化结果里。

把算子封装成 Function 后,我们可以把 my_add算子用起来了。

my_add = MyAddFunction.apply 
 
class MyAdd(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
 
    def forward(self, a, b): 
        return my_add(a, b) 

在这份代码里,我们先用 my_add = MyAddFunction.apply 获取了一个奇怪的变量。这个变量是用来做什么的呢?其实,applytorch.autograd.Function 的一个方法,这个方法完成了 Function 在前向推理或者反向传播时的调度。我们在使用 Function 的派生类做推理时,不应该显式地调用 forward(),而应该调用其 apply 方法。

这里我们使用 my_add = MyAddFunction.apply 把这个调用方法取了一个更简短的别名 my_add。以后在使用 my_add 算子时,我们应该忽略 MyAddFunction 的实现细节,而只通过 my_add 这个接口来访问算子。这里 my_add 的地位,和 PyTorch 的 asinhinterpolateconv2d等原生函数是类似的。

有了访问新算子的接口后,我们可以进一步把算子封装成一个神经网络中的计算层。我们定义一个叫做的 MyAdd 的 torch.nn.Module,它封装了my_add,就和封装了conv2d 的 torch.nn.Conv2d 一样。

测试算子

费了好大的功夫来“包装”我们的新算子后,我们终于可以来使用它了。和之前的测试流程一样,让我们用下面的代码来导出一个包含新算子的 ONNX 模型,并验证一下它是否正确。

model = MyAdd() 
input = torch.rand(1, 3, 10, 10) 
torch.onnx.export(model, (input, input), 'my_add.onnx') 
torch_output = model(input, input).detach().numpy() 
 
import onnxruntime 
import numpy as np 
sess = onnxruntime.InferenceSession('my_add.onnx') 
ort_output = sess.run(None, {'a': input.numpy(), 'b': input.numpy()})[0] 
 
assert np.allclose(torch_output, ort_output) 

在这份代码中,我们直接把 MyAdd 作为要导出的模型。我们计算了一个 PyTorch 模型的运行结果,又导出 ONNX 模型,计算了 ONNX 模型在 ONNX Runtime 上的运算结果。如果一切正常的话,这两个结果是一样的,这份代码不会报任何错误,没有任何输出。

可视化一下 my_add.onnx,可以看出,和我们设计得一样,my_add 算子被翻译成了两个 ONNX 算子节点(其中常量算子被放入了 Mul 的参数中)。

整理一下,整个流程的 Python 代码如下:

import torch 
import my_lib 
class MyAddFunction(torch.autograd.Function): 
 
    @staticmethod 
    def forward(ctx, a, b): 
        return my_lib.my_add(a, b) 
 
    @staticmethod 
    def symbolic(g, a, b): 
        two = g.op("Constant", value_t=torch.tensor([2])) 
        a = g.op('Mul', a, two) 
        return g.op('Add', a, b) 
 
my_add = MyAddFunction.apply 
 
class MyAdd(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
 
    def forward(self, a, b): 
        return my_add(a, b) 
 
model = MyAdd() 
input = torch.rand(1, 3, 10, 10) 
torch.onnx.export(model, (input, input), 'my_add.onnx') 
torch_output = model(input, input).detach().numpy() 
 
import onnxruntime 
import numpy as np 
sess = onnxruntime.InferenceSession('my_add.onnx') 
ort_output = sess.run(None, {'a': input.numpy(), 'b': input.numpy()})[0] 
 
assert np.allclose(torch_output, ort_output) 

总结

在这篇教程中,我们围绕“为 ATen 算子添加符号函数”、“为 TorchScript 算子添加符号函数”、“封装成 torch.autograd.Function 并添加符号函数”这三种添加映射关系的方法,讲解了 3 个为 PyTorch 和 ONNX 添加支持的实例。在这个过程中,我们学到了很多零散的知识,来总结一下吧。

  • ATen 是 PyTorch 的 C++ 张量运算库。通过查询 torch/_C/_VariableFunctions.pyi 和 torch/nn/functional.pyi,我们可以知道 ATen 算子的 Python 接口定义。
  • 用 register_op 可以为 ATen 算子补充注册符号函数
  • 用 register_custom_op_symbolic 可以为 TorchScript 算子补充注册符号函数
  • 如何在 PyTorch 里添加 C++ 拓展
  • 如何用 torch.autograd.Function 封装一个自定义 PyTorch 算子
  • 如何编写符号函数 symbolic(g, ...)
  • 如何用 g.op() 把一个 PyTorch 算子映射成一个或多个 ONNX 算子,或者是自定义的 ONNX 算子。

PyTorch 转 ONNX 详解

转自:模型部署入门教程(三):PyTorch 转 ONNX 详解

模型转换工具 https://convertmodel.com/

ONNX 是目前模型部署中最重要的中间表示之一。学懂了 ONNX 的技术细节,就能规避大量的模型部署问题。
在把 PyTorch 模型转换成 ONNX 模型时,我们往往只需要轻松地调用一句torch.onnx.export就行了。这个函数的接口看上去简单,但它在使用上还有着诸多的“潜规则”。在这篇教程中,我们会详细介绍 PyTorch 模型转 ONNX 模型的原理及注意事项。除此之外,我们还会介绍 PyTorch 与 ONNX 的算子对应关系,以教会大家如何处理 PyTorch 模型转换时可能会遇到的算子支持问题。

torch.onnx.export 细解


在这一节里,我们将详细介绍 PyTorch 到 ONNX 的转换函数—— torch.onnx.export。我们希望大家能够更加灵活地使用这个模型转换接口,并通过了解它的实现原理来更好地应对该函数的报错(由于模型部署的兼容性问题,部署复杂模型时该函数时常会报错)。


计算图导出方法

TorchScript 是一种序列化和优化 PyTorch 模型的格式,在优化过程中,一个torch.nn.Module模型会被转换成 TorchScript 的 torch.jit.ScriptModule模型。现在, TorchScript 也被常当成一种中间表示使用。我们在其他文章中对 TorchScript 有详细的介绍(https://zhuanlan.zhihu.com/p/486914187),这里介绍 TorchScript 仅用于说明 PyTorch 模型转 ONNX的原理。
torch.onnx.export中需要的模型实际上是一个torch.jit.ScriptModule。而要把普通 PyTorch 模型转一个这样的 TorchScript 模型,有跟踪(trace)和记录(script)两种导出计算图的方法。如果给torch.onnx.export传入了一个普通 PyTorch 模型(torch.nn.Module),那么这个模型会默认使用跟踪的方法导出。这一过程如下图所示:

回忆一下我们第一篇教程知识:跟踪法只能通过实际运行一遍模型的方法导出模型的静态图,即无法识别出模型中的控制流(如循环);记录法则能通过解析模型来正确记录所有的控制流。我们以下面这段代码为例来看一看这两种转换方法的区别:

import torch 
 
class Model(torch.nn.Module): 
    def __init__(self, n): 
        super().__init__() 
        self.n = n 
        self.conv = torch.nn.Conv2d(3, 3, 3) 
 
    def forward(self, x): 
        for i in range(self.n): 
            x = self.conv(x) 
        return x 
 
 
models = [Model(2), Model(3)] 
model_names = ['model_2', 'model_3'] 
 
for model, model_name in zip(models, model_names): 
    dummy_input = torch.rand(1, 3, 10, 10) 
    dummy_output = model(dummy_input) 
    model_trace = torch.jit.trace(model, dummy_input) 
    model_script = torch.jit.script(model) 
 
    # 跟踪法与直接 torch.onnx.export(model, ...)等价 
    torch.onnx.export(model_trace, dummy_input, f'{model_name}_trace.onnx', example_outputs=dummy_output) 
    # 记录法必须先调用 torch.jit.sciprt 
    torch.onnx.export(model_script, dummy_input, f'{model_name}_script.onnx', example_outputs=dummy_output) 

在这段代码里,我们定义了一个带循环的模型,模型通过参数n来控制输入张量被卷积的次数。之后,我们各创建了一个n=2n=3的模型。我们把这两个模型分别用跟踪和记录的方法进行导出。
值得一提的是,由于这里的两个模型(model_tracemodel_script)是 TorchScript 模型,export函数已经不需要再运行一遍模型了。(如果模型是用跟踪法得到的,那么在执行torch.jit.trace的时候就运行过一遍了;而用记录法导出时,模型不需要实际运行)参数中的dummy_inputdummy_output`仅仅是为了获取输入和输出张量的类型和形状。
运行上面的代码,我们把得到的 4 个 onnx 文件用 Netron 可视化:

首先看跟踪法得到的 ONNX 模型结构。可以看出来,对于不同的 n,ONNX 模型的结构是不一样的。

而用记录法的话,最终的 ONNX 模型用 Loop 节点来表示循环。这样哪怕对于不同的 n,ONNX 模型也有同样的结构。

本文使用的 PyTorch 版本是 1.8.2。据反馈,其他版本的 PyTorch 可能会得到不一样的结果。

由于推理引擎对静态图的支持更好,通常我们在模型部署时不需要显式地把 PyTorch 模型转成 TorchScript 模型,直接把 PyTorch 模型用 torch.onnx.export 跟踪导出即可。了解这部分的知识主要是为了在模型转换报错时能够更好地定位问题是否发生在 PyTorch 转 TorchScript 阶段。

参数讲解

了解完转换函数的原理后,我们来详细介绍一下该函数的主要参数的作用。我们主要会从应用的角度来介绍每个参数在不同的模型部署场景中应该如何设置,而不会去列出每个参数的所有设置方法。该函数详细的 API 文档可参考: torch.onnx ‒ PyTorch 1.11.0 documentation
torch.onnx.export 在 torch.onnx.__init__.py文件中的定义如下:

def export(model, args, f, export_params=True, verbose=False, training=TrainingMode.EVAL, 
           input_names=None, output_names=None, aten=False, export_raw_ir=False, 
           operator_export_type=None, opset_version=None, _retain_param_name=True, 
           do_constant_folding=True, example_outputs=None, strip_doc_string=True, 
           dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None, 
           enable_onnx_checker=True, use_external_data_format=False): 

前三个必选参数为模型、模型输入、导出的 onnx 文件名,我们对这几个参数已经很熟悉了。我们来着重看一下后面的一些常用可选参数。

export_params

模型中是否存储模型权重。一般中间表示包含两大类信息:模型结构和模型权重,这两类信息可以在同一个文件里存储,也可以分文件存储。ONNX 是用同一个文件表示记录模型的结构和权重的。
我们部署时一般都默认这个参数为 True。如果 onnx 文件是用来在不同框架间传递模型(比如 PyTorch 到 Tensorflow)而不是用于部署,则可以令这个参数为 False。

input_names, output_names

设置输入和输出张量的名称。如果不设置的话,会自动分配一些简单的名字(如数字)。
ONNX 模型的每个输入和输出张量都有一个名字。很多推理引擎在运行 ONNX 文件时,都需要以“名称-张量值”的数据对来输入数据,并根据输出张量的名称来获取输出数据。在进行跟张量有关的设置(比如添加动态维度)时,也需要知道张量的名字。
在实际的部署流水线中,我们都需要设置输入和输出张量的名称,并保证 ONNX 和推理引擎中使用同一套名称。

opset_version

转换时参考哪个 ONNX 算子集版本,默认为 9。后文会详细介绍 PyTorch 与 ONNX 的算子对应关系。

dynamic_axes

指定输入输出张量的哪些维度是动态的。
为了追求效率,ONNX 默认所有参与运算的张量都是静态的(张量的形状不发生改变)。但在实际应用中,我们又希望模型的输入张量是动态的,尤其是本来就没有形状限制的全卷积模型。因此,我们需要显式地指明输入输出张量的哪几个维度的大小是可变的。
我们来看一个dynamic_axes的设置例子:

import torch 
 
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
        self.conv = torch.nn.Conv2d(3, 3, 3) 
 
    def forward(self, x): 
        x = self.conv(x) 
        return x 
 
 
model = Model() 
dummy_input = torch.rand(1, 3, 10, 10) 
model_names = ['model_static.onnx',  
'model_dynamic_0.onnx',  
'model_dynamic_23.onnx'] 
 
dynamic_axes_0 = { 
    'in' : [0], 
    'out' : [0] 
} 
dynamic_axes_23 = { 
    'in' : [2, 3], 
    'out' : [2, 3] 
} 
 
torch.onnx.export(model, dummy_input, model_names[0],  
input_names=['in'], output_names=['out']) 
torch.onnx.export(model, dummy_input, model_names[1],  
input_names=['in'], output_names=['out'], dynamic_axes=dynamic_axes_0) 
torch.onnx.export(model, dummy_input, model_names[2],  
input_names=['in'], output_names=['out'], dynamic_axes=dynamic_axes_23) 

首先,我们导出 3 个 ONNX 模型,分别为没有动态维度、第 0 维动态、第 2 第 3 维动态的模型。
在这份代码里,我们是用列表的方式表示动态维度,例如:

dynamic_axes_0 = { 
    'in' : [0], 
    'out' : [0] 
} 


由于 ONNX 要求每个动态维度都有一个名字,这样写的话会引出一条 UserWarning,警告我们通过列表的方式设置动态维度的话系统会自动为它们分配名字。一种显式添加动态维度名字的方法如下:

dynamic_axes_0 = { 
    'in' : {0: 'batch'}, 
    'out' : {0: 'batch'} 
} 

由于在这份代码里我们没有更多的对动态维度的操作,因此简单地用列表指定动态维度即可。
之后,我们用下面的代码来看一看动态维度的作用:

import onnxruntime 
import numpy as np 
 
origin_tensor = np.random.rand(1, 3, 10, 10).astype(np.float32) 
mult_batch_tensor = np.random.rand(2, 3, 10, 10).astype(np.float32) 
big_tensor = np.random.rand(1, 3, 20, 20).astype(np.float32) 
 
inputs = [origin_tensor, mult_batch_tensor, big_tensor] 
exceptions = dict() 
 
for model_name in model_names: 
    for i, input in enumerate(inputs): 
        try: 
            ort_session = onnxruntime.InferenceSession(model_name) 
            ort_inputs = {'in': input} 
            ort_session.run(['out'], ort_inputs) 
        except Exception as e: 
            exceptions[(i, model_name)] = e 
            print(f'Input[{i}] on model {model_name} error.') 
        else: 
            print(f'Input[{i}] on model {model_name} succeed.') 

我们在模型导出计算图时用的是一个形状为(1, 3, 10, 10)的张量。现在,我们来尝试以形状分别是(1, 3, 10, 10), (2, 3, 10, 10), (1, 3, 20, 20)为输入,用ONNX Runtime运行一下这几个模型,看看哪些情况下会报错,并保存对应的报错信息。得到的输出信息应该如下:

Input[0] on model model_static.onnx succeed. 
Input[1] on model model_static.onnx error. 
Input[2] on model model_static.onnx error. 
Input[0] on model model_dynamic_0.onnx succeed. 
Input[1] on model model_dynamic_0.onnx succeed. 
Input[2] on model model_dynamic_0.onnx error. 
Input[0] on model model_dynamic_23.onnx succeed. 
Input[1] on model model_dynamic_23.onnx error. 
Input[2] on model model_dynamic_23.onnx succeed. 

可以看出,形状相同的(1, 3, 10, 10)的输入在所有模型上都没有出错。而对于batch(第 0 维)或者长宽(第 2、3维)不同的输入,只有在设置了对应的动态维度后才不会出错。我们可以错误信息中找出是哪些维度出了问题。比如我们可以用以下代码查看input[1]model_static.onnx中的报错信息:

print(exceptions[(1, 'model_static.onnx')]) 
 
# output 
# [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: in for the following indices index: 0 Got: 2 Expected: 1 Please fix either the inputs or the model. 

这段报错告诉我们名字叫in的输入的第 0 维不匹配。本来该维的长度应该为 1,但我们的输入是 2。实际部署中,如果我们碰到了类似的报错,就可以通过设置动态维度来解决问题。

使用提示

通过学习之前的知识,我们基本掌握了 torch.onnx.export函数的部分实现原理和参数设置方法,足以完成简单模型的转换了。但在实际应用中,使用该函数还会踩很多坑。这里我们模型部署团队把在实战中积累的一些经验分享给大家。

使模型在 ONNX 转换时有不同的行为

有些时候,我们希望模型在导出至 ONNX 时有一些不同的行为模型在直接用 PyTorch 推理时有一套逻辑,而在导出的ONNX模型中有另一套逻辑。比如,我们可以把一些后处理的逻辑放在模型里,以简化除运行模型之外的其他代码。torch.onnx.is_in_onnx_export()可以实现这一任务,该函数仅在执行 torch.onnx.export()时为真。以下是一个例子:

import torch 
 
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
        self.conv = torch.nn.Conv2d(3, 3, 3) 
 
    def forward(self, x): 
        x = self.conv(x) 
        if torch.onnx.is_in_onnx_export(): 
            x = torch.clip(x, 0, 1) 
        return x 


这里,我们仅在模型导出时把输出张量的数值限制在[0, 1]之间。使用 is_in_onnx_export确实能让我们方便地在代码中添加和模型部署相关的逻辑。但是,这些代码对只关心模型训练的开发者和用户来说很不友好,突兀的部署逻辑会降低代码整体的可读性。同时,is_in_onnx_export只能在每个需要添加部署逻辑的地方都“打补丁”,难以进行统一的管理。我们之后会介绍如何使用 MMDeploy 的重写机制来规避这些问题。

利用中断张量跟踪的操作

PyTorch 转 ONNX 的跟踪导出法是不是万能的。如果我们在模型中做了一些很“出格”的操作,跟踪法会把某些取决于输入的中间结果变成常量,从而使导出的 ONNX 模型和原来的模型有出入。以下是一个会造成这种“跟踪中断”的例子:

class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
 
    def forward(self, x): 
        x = x * x[0].item() 
        return x, torch.Tensor([i for i in x]) 
 
model = Model()       
dummy_input = torch.rand(10) 
torch.onnx.export(model, dummy_input, 'a.onnx') 

如果你尝试去导出这个模型,会得到一大堆 warning,告诉你转换出来的模型可能不正确。这也难怪,我们在这个模型里使用了.item()把 torch 中的张量转换成了普通的 Python 变量,还尝试遍历 torch 张量,并用一个列表新建一个 torch 张量。这些涉及张量与普通变量转换的逻辑都会导致最终的 ONNX 模型不太正确。
另一方面,我们也可以利用这个性质,在保证正确性的前提下令模型的中间结果变成常量。这个技巧常常用于模型的静态化上,即令模型中所有的张量形状都变成常量。在未来的教程中,我们会在部署实例中详细介绍这些“高级”操作。

使用张量为输入(PyTorch版本 < 1.9.0)

正如我们第一篇教程所展示的,在较旧(< 1.9.0)的 PyTorch 中把 Python 数值作为 torch.onnx.export()的模型输入时会报错。出于兼容性的考虑,我们还是推荐以张量为模型转换时的模型输入。

PyTorch 对 ONNX 的算子支持

在确保torch.onnx.export()的调用方法无误后,PyTorch 转 ONNX 时最容易出现的问题就是算子不兼容了。这里我们会介绍如何判断某个 PyTorch 算子在 ONNX 中是否兼容,以助大家在碰到报错时能更好地把错误归类。而具体添加算子的方法我们会在之后的文章里介绍。
在转换普通的torch.nn.Module模型时,PyTorch 一方面会用跟踪法执行前向推理,把遇到的算子整合成计算图;另一方面,PyTorch 还会把遇到的每个算子翻译成 ONNX 中定义的算子。在这个翻译过程中,可能会碰到以下情况:

  • 该算子可以一对一地翻译成一个 ONNX 算子。
  • 该算子在 ONNX 中没有直接对应的算子,会翻译成一至多个 ONNX 算子。
  • 该算子没有定义翻译成 ONNX 的规则,报错。

那么,该如何查看 PyTorch 算子与 ONNX 算子的对应情况呢?由于 PyTorch 算子是向 ONNX 对齐的,这里我们先看一下 ONNX 算子的定义情况,再看一下 PyTorch 定义的算子映射关系。

ONNX 算子文档

ONNX 算子的定义情况,都可以在官方的算子文档中查看。这份文档十分重要,我们碰到任何和 ONNX 算子有关的问题都得来”请教“这份文档

这份文档中最重要的开头的这个算子变更表格。表格的第一列是算子名,第二列是该算子发生变动的算子集版本号,也就是我们之前在torch.onnx.export中提到的opset_version表示的算子集版本号。通过查看算子第一次发生变动的版本号,我们可以知道某个算子是从哪个版本开始支持的;通过查看某算子小于等于opset_version的第一个改动记录,我们可以知道当前算子集版本中该算子的定义规则。

通过点击表格中的链接,我们可以查看某个算子的输入、输出参数规定及使用示例。比如上图是 Relu 在 ONNX 中的定义规则,这份定义表明 Relu 应该有一个输入和一个输入,输入输出的类型相同,均为 tensor。

PyTorch 对 ONNX 算子的映射

在 PyTorch 中,和 ONNX 有关的定义全部放在 torch.onnx目录中,如下图所示:

其中,symbolic_opset{n}.py(符号表文件)即表示 PyTorch 在支持第 n 版 ONNX 算子集时新加入的内容。我们之前讲过, bicubic 插值是在第 11 个版本开始支持的。我们以它为例来看看如何查找算子的映射情况。
首先,使用搜索功能,在torch/onnx文件夹搜索”bicubic”,可以发现这个这个插值在第 11 个版本的定义文件中:

之后,我们按照代码的调用逻辑,逐步跳转直到最底层的 ONNX 映射函数:

upsample_bicubic2d = _interpolate("upsample_bicubic2d", 4, "cubic") 
 
-> 
 
def _interpolate(name, dim, interpolate_mode): 
    return sym_help._interpolate_helper(name, dim, interpolate_mode) 
 
-> 
 
def _interpolate_helper(name, dim, interpolate_mode): 
    def symbolic_fn(g, input, output_size, *args): 
        ... 
 
    return symbolic_fn 

最后,在symbolic_fn中,我们可以看到插值算子是怎么样被映射成多个 ONNX 算子的。其中,每一个g.op就是一个 ONNX 的定义。比如其中的 Resize 算子就是这样写的:

return g.op("Resize", 
                input, 
                empty_roi, 
                empty_scales, 
                output_size, 
                coordinate_transformation_mode_s=coordinate_transformation_mode, 
                cubic_coeff_a_f=-0.75,  # only valid when mode="cubic" 
                mode_s=interpolate_mode,  # nearest, linear, or cubic 
                nearest_mode_s="floor")  # only valid when mode="nearest" 

通过在前面提到的ONNX 算子文档中查找 Resize 算子的定义,我们就可以知道这每一个参数的含义了。用类似的方法,我们可以去查询其他 ONNX 算子的参数含义,进而知道 PyTorch 中的参数是怎样一步一步传入到每个 ONNX 算子中的。
掌握了如何查询 PyTorch 映射到 ONNX 的关系后,我们在实际应用时就可以在 torch.onnx.export()opset_version中先预设一个版本号,碰到了问题就去对应的 PyTorch 符号表文件里去查。如果某算子确实不存在,或者算子的映射关系不满足我们的要求,我们就可能得用其他的算子绕过去,或者自定义算子了。

总结

在这篇教程中,我们系统地介绍了 PyTorch 转 ONNX 的原理。我们先是着重讲解了使用最频繁的 torch.onnx.export函数,又给出了查询 PyTorch 对 ONNX 算子支持情况的方法。通过本文,我们希望大家能够成功转换出大部分不需要添加新算子的 ONNX 模型,并在碰到算子问题时能够有效定位问题原因。具体而言,大家读完本文后应该了解以下的知识:

  • 跟踪法和记录法在导出带控制语句的计算图时有什么区别。
  • torch.onnx.export()中该如何设置 input_names, output_names, dynamic_axes
  • 使用 torch.onnx.is_in_onnx_export()来使模型在转换到 ONNX 时有不同的行为。
  • 如何查询 ONNX 算子文档(https://github.com/onnx/onnx/blob/main/docs/Operators.md)。
  • 如何查询 PyTorch 对某个 ONNX 版本的新特性支持情况。
  • 如何判断 PyTorch 对某个 ONNX 算子是否支持,支持的方法是怎样的。

模型部署:解决模型部署中的难题

转:模型部署入门教程(二):解决模型部署中的难题

我们部署了一个简单的超分辨率模型,一切都十分顺利。但是,上一个模型还有一些缺陷——图片的放大倍数固定是 4,我们无法让图片放大任意的倍数。现在,我们来尝试部署一个支持动态放大倍数的模型,体验一下在模型部署中可能会碰到的困难。

模型部署中常见的难题

在之前的学习中,我们在模型部署上顺风顺水,没有碰到任何问题。这是因为 SRCNN 模型只包含几个简单的算子,而这些卷积、插值算子已经在各个中间表示和推理引擎上得到了完美支持。如果模型的操作稍微复杂一点,我们可能就要为兼容模型而付出大量的功夫了。实际上,模型部署时一般会碰到以下几类困难:

  • 模型的动态化。出于性能的考虑,各推理框架都默认模型的输入形状、输出形状、结构是静态的。而为了让模型的泛用性更强,部署时需要在尽可能不影响原有逻辑的前提下,让模型的输入输出或是结构动态化。
  • 新算子的实现。深度学习技术日新月异,提出新算子的速度往往快于 ONNX 维护者支持的速度。为了部署最新的模型,部署工程师往往需要自己在 ONNX 和推理引擎中支持新算子。
  • 中间表示与推理引擎的兼容问题。由于各推理引擎的实现不同,对 ONNX 难以形成统一的支持。为了确保模型在不同的推理引擎中有同样的运行效果,部署工程师往往得为某个推理引擎定制模型代码,这为模型部署引入了许多工作量。

现在,让我们对原来的 SRCNN 模型做一些小的修改,体验一下模型动态化对模型部署造成的困难,并学习解决该问题的一种方法。

实现动态放大的超分辨率模型

在原来的 SRCNN 中,图片的放大比例是写死在模型里的:

class SuperResolutionNet(nn.Module): 
    def __init__(self, upscale_factor): 
        super().__init__() 
        self.upscale_factor = upscale_factor 
        self.img_upsampler = nn.Upsample( 
            scale_factor=self.upscale_factor, 
            mode='bicubic', 
            align_corners=False) 
 
... 
 
def init_torch_model(): 
    torch_model = SuperResolutionNet(upscale_factor=3) 
 

我们使用 upscale_factor 来控制模型的放大比例。初始化模型的时候,我们默认令 upscale_factor 为 3,生成了一个放大 3 倍的 PyTorch 模型。这个 PyTorch 模型最终被转换成了 ONNX 格式的模型。如果我们需要一个放大 4 倍的模型,需要重新生成一遍模型,再做一次到 ONNX 的转换。

现在,假设我们要做一个超分辨率的应用。我们的用户希望图片的放大倍数能够自由设置。而我们交给用户的,只有一个 .onnx 文件和运行超分辨率模型的应用程序。我们在不修改 .onnx 文件的前提下改变放大倍数。

因此,我们必须修改原来的模型,令模型的放大倍数变成推理时的输入。在上一篇文章中的 Python 脚本的基础上,我们做一些修改,得到这样的脚本:

import torch 
from torch import nn 
from torch.nn.functional import interpolate 
import torch.onnx 
import cv2 
import numpy as np 
 
 
class SuperResolutionNet(nn.Module): 
 
    def __init__(self): 
        super().__init__() 
 
        self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4) 
        self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0) 
        self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2) 
 
        self.relu = nn.ReLU() 
 
    def forward(self, x, upscale_factor): 
        x = interpolate(x, 
                        scale_factor=upscale_factor, 
                        mode='bicubic', 
                        align_corners=False) 
        out = self.relu(self.conv1(x)) 
        out = self.relu(self.conv2(out)) 
        out = self.conv3(out) 
        return out 
 
 
def init_torch_model(): 
    torch_model = SuperResolutionNet() 
 
    state_dict = torch.load('srcnn.pth')['state_dict'] 
 
    # Adapt the checkpoint 
    for old_key in list(state_dict.keys()): 
        new_key = '.'.join(old_key.split('.')[1:]) 
        state_dict[new_key] = state_dict.pop(old_key) 
 
    torch_model.load_state_dict(state_dict) 
    torch_model.eval() 
    return torch_model 
 
 
model = init_torch_model() 
 
input_img = cv2.imread('face.png').astype(np.float32) 
 
# HWC to NCHW 
input_img = np.transpose(input_img, [2, 0, 1]) 
input_img = np.expand_dims(input_img, 0) 
 
# Inference 
torch_output = model(torch.from_numpy(input_img), 3).detach().numpy() 
 
# NCHW to HWC 
torch_output = np.squeeze(torch_output, 0) 
torch_output = np.clip(torch_output, 0, 255) 
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8) 
 
# Show image 
cv2.imwrite("face_torch_2.png", torch_output) 

SuperResolutionNet 未修改之前,nn.Upsample 在初始化阶段固化了放大倍数,而 PyTorch 的 interpolate 插值算子可以在运行阶段选择放大倍数。因此,我们在新脚本中使用 interpolate 代替 nn.Upsample,从而让模型支持动态放大倍数的超分。 在第 55 行使用模型推理时,我们把放大倍数设置为 3。最后,图片保存在文件 “face_torch_2.png” 中。一切正常的话,”face_torch_2.png” 和 “face_torch.png” 的内容一模一样。

通过简单的修改,PyTorch 模型已经支持了动态分辨率。现在我们来尝试一下导出模型:

x = torch.randn(1, 3, 256, 256) 
 
with torch.no_grad(): 
    torch.onnx.export(model, (x, 3), 
                      "srcnn2.onnx", 
                      opset_version=11, 
                      input_names=['input', 'factor'], 
                      output_names=['output']) 
 

运行这些脚本时,会报一长串错误。没办法,我们碰到了模型部署中的兼容性问题。

解决方法:自定义算子

直接使用 PyTorch 模型的话,我们修改几行代码就能实现模型输入的动态化。但在模型部署中,我们要花数倍的时间来设法解决这一问题。现在,让我们顺着解决问题的思路,体验一下模型部署的困难,并学习使用自定义算子的方式,解决超分辨率模型的动态化问题。

刚刚的报错是因为 PyTorch 模型在导出到 ONNX 模型时,模型的输入参数的类型必须全部是 torch.Tensor。而实际上我们传入的第二个参数” 3 “是一个整形变量。这不符合 PyTorch 转 ONNX 的规定。我们必须要修改一下原来的模型的输入。为了保证输入的所有参数都是 torch.Tensor 类型的,我们做如下修改

... 
 
class SuperResolutionNet(nn.Module): 
 
    def forward(self, x, upscale_factor): 
        x = interpolate(x, 
                        scale_factor=upscale_factor.item(), 
                        mode='bicubic', 
                        align_corners=False) 
 
... 
 
# Inference 
# Note that the second input is torch.tensor(3) 
torch_output = model(torch.from_numpy(input_img), torch.tensor(3)).detach().numpy() 
 
... 
 
with torch.no_grad(): 
    torch.onnx.export(model, (x, torch.tensor(3)), 
                      "srcnn2.onnx", 
                      opset_version=11, 
                      input_names=['input', 'factor'], 
                      output_names=['output']) 

由于 PyTorch 中 interpolate 的 scale_factor 参数必须是一个数值,我们使用 torch.Tensor.item() 来把只有一个元素的 torch.Tensor 转换成数值。之后,在模型推理时,我们使用 torch.tensor(3) 代替 3,以使得我们的所有输入都满足要求。现在运行脚本的话,无论是直接运行模型,还是导出 ONNX 模型,都不会报错了。

但是,导出 ONNX 时却报了一条 TraceWarning 的警告。这条警告说有一些量可能会追踪失败。这是怎么回事呢?让我们把生成的 srcnn2.onnx 用 Netron 可视化一下:

可以发现,虽然我们把模型推理的输入设置为了两个,但 ONNX 模型还是长得和原来一模一样,只有一个叫 ” input ” 的输入。这是由于我们使用了 torch.Tensor.item() 把数据从 Tensor 里取出来,而导出 ONNX 模型时这个操作是无法被记录的,只好报了一条 TraceWarning。这导致 interpolate 插值函数的放大倍数还是被设置成了” 3 “这个固定值,我们导出的” srcnn2.onnx “和最开始的” srcnn.onnx “完全相同。

直接修改原来的模型似乎行不通,我们得从 PyTorch 转 ONNX 的原理入手,强行令 ONNX 模型明白我们的想法了。

仔细观察 Netron 上可视化出的 ONNX 模型,可以发现在 PyTorch 中无论是使用最早的 nn.Upsample,还是后来的 interpolate,PyTorch 里的插值操作最后都会转换成 ONNX 定义的 Resize 操作。也就是说,所谓 PyTorch 转 ONNX,实际上就是把每个 PyTorch 的操作映射成了 ONNX 定义的算子。

点击该算子,可以看到它的详细参数如下:

其中,展开 scales,可以看到 scales 是一个长度为 4 的一维张量,其内容为 [1, 1, 3, 3], 表示 Resize 操作每一个维度的缩放系数;其类型为 Initializer,表示这个值是根据常量直接初始化出来的。如果我们能够自己生成一个 ONNX 的 Resize 算子,让 scales 成为一个可变量而不是常量,就像它上面的 X 一样,那这个超分辨率模型就能动态缩放了。

现有实现插值的 PyTorch 算子有一套规定好的映射到 ONNX Resize 算子的方法,这些映射出的 Resize 算子的 scales 只能是常量,无法满足我们的需求。我们得自己定义一个实现插值的 PyTorch 算子,然后让它映射到一个我们期望的 ONNX Resize 算子上。

下面的脚本定义了一个 PyTorch 插值算子,并在模型里使用了它。我们先通过运行模型来验证该算子的正确性:

import torch 
from torch import nn 
from torch.nn.functional import interpolate 
import torch.onnx 
import cv2 
import numpy as np 
 
 
class NewInterpolate(torch.autograd.Function): 
 
    @staticmethod 
    def symbolic(g, input, scales): 
        return g.op("Resize", 
                    input, 
                    g.op("Constant", 
                         value_t=torch.tensor([], dtype=torch.float32)), 
                    scales, 
                    coordinate_transformation_mode_s="pytorch_half_pixel", 
                    cubic_coeff_a_f=-0.75, 
                    mode_s='cubic', 
                    nearest_mode_s="floor") 
 
    @staticmethod 
    def forward(ctx, input, scales): 
        scales = scales.tolist()[-2:] 
        return interpolate(input, 
                           scale_factor=scales, 
                           mode='bicubic', 
                           align_corners=False) 
 
 
class StrangeSuperResolutionNet(nn.Module): 
 
    def __init__(self): 
        super().__init__() 
 
        self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4) 
        self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0) 
        self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2) 
 
        self.relu = nn.ReLU() 
 
    def forward(self, x, upscale_factor): 
        x = NewInterpolate.apply(x, upscale_factor) 
        out = self.relu(self.conv1(x)) 
        out = self.relu(self.conv2(out)) 
        out = self.conv3(out) 
        return out 
 
 
def init_torch_model(): 
    torch_model = StrangeSuperResolutionNet() 
 
    state_dict = torch.load('srcnn.pth')['state_dict'] 
 
    # Adapt the checkpoint 
    for old_key in list(state_dict.keys()): 
        new_key = '.'.join(old_key.split('.')[1:]) 
        state_dict[new_key] = state_dict.pop(old_key) 
 
    torch_model.load_state_dict(state_dict) 
    torch_model.eval() 
    return torch_model 
 
 
model = init_torch_model() 
factor = torch.tensor([1, 1, 3, 3], dtype=torch.float) 
 
input_img = cv2.imread('face.png').astype(np.float32) 
 
# HWC to NCHW 
input_img = np.transpose(input_img, [2, 0, 1]) 
input_img = np.expand_dims(input_img, 0) 
 
# Inference 
torch_output = model(torch.from_numpy(input_img), factor).detach().numpy() 
 
# NCHW to HWC 
torch_output = np.squeeze(torch_output, 0) 
torch_output = np.clip(torch_output, 0, 255) 
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8) 
 
# Show image 
cv2.imwrite("face_torch_3.png", torch_output) 

模型运行正常的话,一幅放大3倍的超分辨率图片会保存在”face_torch_3.png”中,其内容和”face_torch.png”完全相同。

在刚刚那个脚本中,我们定义 PyTorch 插值算子的代码如下:

class NewInterpolate(torch.autograd.Function): 
 
    @staticmethod 
    def symbolic(g, input, scales): 
        return g.op("Resize", 
                    input, 
                    g.op("Constant", 
                         value_t=torch.tensor([], dtype=torch.float32)), 
                    scales, 
                    coordinate_transformation_mode_s="pytorch_half_pixel", 
                    cubic_coeff_a_f=-0.75, 
                    mode_s='cubic', 
                    nearest_mode_s="floor") 
 
    @staticmethod 
    def forward(ctx, input, scales): 
        scales = scales.tolist()[-2:] 
        return interpolate(input, 
                           scale_factor=scales, 
                           mode='bicubic', 
                           align_corners=False) 

在具体介绍这个算子的实现前,让我们先理清一下思路。我们希望新的插值算子有两个输入,一个是被用于操作的图像,一个是图像的放缩比例。前面讲到,为了对接 ONNX 中 Resize 算子的 scales 参数,这个放缩比例是一个 [1, 1, x, x] 的张量,其中 x 为放大倍数。在之前放大3倍的模型中,这个参数被固定成了[1, 1, 3, 3]。因此,在插值算子中,我们希望模型的第二个输入是一个 [1, 1, w, h] 的张量,其中 w 和 h 分别是图片宽和高的放大倍数。

搞清楚了插值算子的输入,再看一看算子的具体实现。算子的推理行为由算子的 foward 方法决定。该方法的第一个参数必须为 ctx,后面的参数为算子的自定义输入,我们设置两个输入,分别为被操作的图像和放缩比例。为保证推理正确,需要把 [1, 1, w, h] 格式的输入对接到原来的 interpolate 函数上。我们的做法是截取输入张量的后两个元素,把这两个元素以 list 的格式传入 interpolate 的 scale_factor 参数。

接下来,我们要决定新算子映射到 ONNX 算子的方法。映射到 ONNX 的方法由一个算子的 symbolic 方法决定。symbolic 方法第一个参数必须是g,之后的参数是算子的自定义输入,和 forward 函数一样。ONNX 算子的具体定义由 g.op 实现。g.op 的每个参数都可以映射到 ONNX 中的算子属性:

对于其他参数,我们可以照着现在的 Resize 算子填。而要注意的是,我们现在希望 scales 参数是由输入动态决定的。因此,在填入 ONNX 的 scales 时,我们要把 symbolic 方法的输入参数中的 scales 填入。

接着,让我们把新模型导出成 ONNX 模型:

x = torch.randn(1, 3, 256, 256) 
 
with torch.no_grad(): 
    torch.onnx.export(model, (x, factor), 
                      "srcnn3.onnx", 
                      opset_version=11, 
                      input_names=['input', 'factor'], 
                      output_names=['output']) 

把导出的 ” srcnn3.onnx ” 进行可视化:

可以看到,正如我们所期望的,导出的 ONNX 模型有了两个输入!第二个输入表示图像的放缩比例。

之前在验证 PyTorch 模型和导出 ONNX 模型时,我们宽高的缩放比例设置成了 3×3。现在,在用 ONNX Runtime 推理时,我们尝试使用 4×4 的缩放比例:

import onnxruntime 
 
input_factor = np.array([1, 1, 4, 4], dtype=np.float32) 
ort_session = onnxruntime.InferenceSession("srcnn3.onnx") 
ort_inputs = {'input': input_img, 'factor': input_factor} 
ort_output = ort_session.run(None, ort_inputs)[0] 
 
ort_output = np.squeeze(ort_output, 0) 
ort_output = np.clip(ort_output, 0, 255) 
ort_output = np.transpose(ort_output, [1, 2, 0]).astype(np.uint8) 
cv2.imwrite("face_ort_3.png", ort_output) 

运行上面的代码,可以得到一个边长放大4倍的超分辨率图片 “face_ort_3.png”。动态的超分辨率模型生成成功了!只要修改 input_factor,我们就可以自由地控制图片的缩放比例。

我们刚刚的工作,实际上是绕过 PyTorch 本身的限制,凭空“捏”出了一个 ONNX 算子。事实上,我们不仅可以创建现有的 ONNX 算子,还可以定义新的 ONNX 算子以拓展 ONNX 的表达能力。后续教程中我们将介绍自定义新 ONNX 算子的方法。

总结:

  • 模型部署中常见的几类困难有:模型的动态化;新算子的实现;框架间的兼容。
  • PyTorch 转 ONNX,实际上就是把每一个操作转化成 ONNX 定义的某一个算子。比如对于 PyTorch 中的 Upsample 和 interpolate,在转 ONNX 后最终都会成为 ONNX 的 Resize 算子。
  • 通过修改继承自 torch.autograd.Function 的算子的 symbolic 方法,可以改变该算子映射到 ONNX 算子的行为。