torchrun-torch多机多卡分布式训练命令

https://pytorch.org/docs/stable/elastic/run.html

torchrun 提供了作为 torch.distributed.launch 的功能的超集,并具有以下附加功能:

  1. Worker failures are handled gracefully by restarting all workers.
    通过重新启动所有工作进程来优雅地处理工作进程故障。
  2. Worker RANK and WORLD_SIZE are assigned automatically.
    工作器 RANK 和 WORLD_SIZE 自动分配。
  3. Number of nodes is allowed to change between minimum and maximum sizes (elasticity).
    允许节点数量在最小和最大大小之间变化(弹性)。

torchrun 是setup.py中 entry_points 配置中声明的主模块torch.distributed.run的Python控制台脚本。它相当于调用 python -m torch.distributed.run 。

torchrun启动单机多卡DDP并行训练:

启动方式:

  • 使用 torchrun 命令来启动程序
  • torchrun –standalone –nproc_per_node=gpu XXX.py
  1. --standalone 代表单机运行
  2. --nproc_per_node=gpu 代表使用所有可用GPU。等于号后也可写gpu数量n,这样会使用前n个GPU

如果想要进一步指定要运行的 GPU,可以通过 CUDA_VISIBLE_DEVICES 设置GPU可见性,比如:

CUDA_VISIBLE_DEVICES=2,3 torchrun –standalone –nproc_per_node=gpu multi_gpu_torchrun.py

多机多卡

torchrun –nproc_per_node=4 –nnodes=3 –node_rank=0 –master_addr=192.168.0.101 –master_port=29500 test_mpi.py

1.指定每个节点(机器)上的进程数,这里是4个。意味着每个机器将启动4个进程来参与分布式训练。 –nproc_per_node=4 【一般设置为为节点GPU数量】

2.指定总共的节点数,这里是3个。意味着总共有3个机器参与分布式训练。

--nnodes=3

3.指定当前节点(机器)的排名,这里是0。排名从0开始,用于在分布式环境中区分不同的节点。

--node_rank=0 【0代表主节点】

4.指定主节点的IP地址,这里是192.168.0.101(更根据实际修改)。主节点用于协调分布式训练过程。

--master_addr=192.168.0.101

5.指定主节点的端口号,这里是29500。主节点使用指定的端口来与其他节点进行通信。

–master_port=29500

6.单机运行

--standalone

torch.cuda.empty_cache 清空显存

empty_cache()不会增加Pytorch可用的GPU内存量。但是,在某些情况下,它可能有助于减少GPU记忆的碎片化。有关GPU内存管理的更多详细信息,请参见内存管理。

因为PyTorch是有缓存区的设置的,意思就是一个Tensor就算被释放了,进程也不会把空闲出来的显存还给GPU,而是等待下一个Tensor来填入这一片被释放的空间。所以我们用nvidia-smi/gpustat 看到的显存占用不会减少

用torch.cuda.empty_cache可以清空缓冲区:

在程序中加上这句会使速度变慢一些,但是有些情况下会有用,例如程序之前test的时候总是爆显存,然后在循环中加上了这句就不爆了

for i, data in enumerate(data_loader):
    torch.cuda.empty_cache()
    img_meta = data['img_meta'][0].data[0]
    img_name = img_meta[0]['filename'].split('/')[-1]
    with torch.no_grad():
        result = model(return_loss=False, rescale=not show, **data)

如果显存资源比较紧缺,可以在每个epoch开始时释放下不用的显存资源。

torch.cuda.empty_cache() # 释放显存
————————————————

torchaudio

官网:https://pytorch.org/audio/stable/torchaudio.html

Torchaudio is a library for audio and signal processing with PyTorch. It provides I/O, signal and data processing functions, datasets, model implementations and application components.

读取音频:

使用 torchaudio.load 加载音频数据。torchaudio.load 支持类路径对象和类文件对象。返回值是波形(tensor)和采样率(int)的元组。默认情况下,生成的 tensor 对象的类型为 torch.float32,其值在[−1.0,1.0][−1.0,1.0]内标准化。
waveform, sr = torchaudio.load(filepath, frame_offset=0 , num_frames=-1, normalize=True, channels_first=True)
参数:

filepath (str): 原始音频文件路径;
frame_offset (int): 在此之后开始读取,默认为0,以帧为单位;
num_frames (int): 读取的最大帧数。默认是-1,则表示从frame_offset直到末尾。如果给定文件中没有足够的帧,这个函数可能会返回实际剩余的帧数。
normalize (bool): 当为True时,该函数总是返回float32,并且所有的值被归一化到[-1,1]。如果输入文件是wav,且是整形,若为False时,则会输出int类型。需要注意的是,该参数仅对wav类型的文件起作用, 默认是True;
channels_first (bool)—当为True时,返回的Tensor的维度是[channel, time]。否则,维数为[time, channel], 默认是True。
返回:

waveform (torch.Tensor): 如果输入文件是int类型的wav,且normalization为False,则waveform的数据就为int类型的,否则是float32;如果channel_first=True,则waveform.shape=[channel, time]。
sr (int): 采样率
重采样
waveform = torchaudio.transforms.Resample(orig_freq=16000, new_freq=16000)(waveform)
参数:

orig_freq (int, optional): 原始采样率,默认:16000;
new_freq (int, optional): 转换后的采样率,默认:16000;
resampling_method (str, optional) – 重采样方法,默认: ‘sinc_interpolation’;
waveform (torch.Tensor): 输入音频维度可以是[channel,time],也可以是[time, channel];
返回:

waveform (torch.Tensor): 输出音频维度和输入音频相同,但由于重采样了,time的数值会不同;
保存音频
torchaudio.save(filepath, src, sample_rate, channels_first)
参数:

firepath (str or pathlib.Path): 保存路径;
src (torch.Tensor): 音频数据,必须是二维的;(注:需要转到cpu下的tensor)
sample_rate(int): 采样率;
channels_first (bool): If True, 维度必须是[channel, time],否则是[time, channel]。

The NSynth Dataset

A large-scale and high-quality dataset of annotated musical notes.( 一个大规模、高质量的注释音符数据集。)

下载地址:https://magenta.tensorflow.org/datasets/nsynth#files

Motivation

Recent breakthroughs in generative modeling of images have been predicated on the availability of high-quality and large-scale datasebts such as MNIST, CIFAR and ImageNet. We recognized the need for an audio dataset that was as approachable as those in the image domain.

Audio signals found in the wild contain multi-scale dependencies that prove particularly difficult to model, leading many previous efforts at data-driven audio synthesis to focus on more constrained domains such as texture synthesis or training small parametric models.

We encourage the broader community to use NSynth as a benchmark and entry point into audio machine learning. We also view NSynth as a building block for future datasets and envision a high-quality multi-note dataset for tasks like generation and transcription that involve learning complex language-like dependencies.

Description

NSynth is an audio dataset containing 305,979 musical notes, each with a unique pitch, timbre, and envelope. For 1,006 instruments from commercial sample libraries, we generated four second, monophonic 16kHz audio snippets, referred to as notes, by ranging over every pitch of a standard MIDI pian o (21-108) as well as five different velocities (25, 50, 75, 100, 127). The note was held for the first three seconds and allowed to decay for the final second.

Some instruments are not capable of producing all 88 pitches in this range, resulting in an average of 65.4 pitches per instrument. Furthermore, the commercial sample packs occasionally contain duplicate sounds across multiple velocities, leaving an average of 4.75 unique velocities per pitch.

We also annotated each of the notes with three additional pieces of information based on a combination of human evaluation and heuristic algorithms:

  • Source: The method of sound production for the note’s instrument. This can be one of acoustic or electronic for instruments that were recorded from acoustic or electronic instruments, respectively, or synthetic for synthesized instruments. See their frequencies below.
  • Family: The high-level family of which the note’s instrument is a member. Each instrument is a member of exactly one family. See the complete list and their frequencies below.
  • Qualities: Sonic qualities of the note. See the quality descriptions and their co-occurrences below. Each note is annotated with zero or more qualities.

Format

Files

The NSynth dataset can be download in two formats:

The full dataset is split into three sets:

  • Train [tfrecord | json/wav]: A training set with 289,205 examples. Instruments do not overlap with valid or test.
  • Valid [tfrecord | json/wav]: A validation set with 12,678 examples. Instruments do not overlap with train.
  • Test [tfrecord | json/wav]: A test set with 4,096 examples. Instruments do not overlap with train.

Below we detail how the note features are encoded in the Example protocol buffers and JSON files.

Example Features

Each Example contains the following features.

FeatureTypeDescription
noteint64A unique integer identifier for the note.
note_strbytesA unique string identifier for the note in the format <instrument_str>-<pitch>-<velocity>.
instrumentint64A unique, sequential identifier for the instrument the note was synthesized from.
instrument_strbytesA unique string identifier for the instrument this note was synthesized from in the format <instrument_family_str>-<instrument_production_str>-<instrument_name>.
pitchint64The 0-based MIDI pitch in the range [0, 127].
velocityint64The 0-based MIDI velocity in the range [0, 127].
sample_rateint64The samples per second for the audio feature.
audio*[float]A list of audio samples represented as floating point values in the range [-1,1].
qualities[int64]A binary vector representing which sonic qualities are present in this note.
qualities_str[bytes]A list IDs of which qualities are present in this note selected from the sonic qualities list.
instrument_familyint64The index of the instrument family this instrument is a member of.
instrument_family_strbytesThe ID of the instrument family this instrument is a member of.
instrument_sourceint64The index of the sonic source for this instrument.
instrument_source_strbytesThe ID of the sonic source for this instrument.

Note: the “audio” feature is ommited from the JSON-encoded examples since the audio data is stored separately in WAV files keyed by the “note_str”.

Vector Quantization 矢量量化 [VQVAE]

http://www.mqasem.net/vectorquantization/vq.html

VQ, 即Vector Quantization,矢量量化,在多个场景下使用,如图像压缩,声音压缩,语音识别等。

Github: https://github.com/lucidrains/vector-quantize-pytorch

矢量量化方法,即Vector Quantization,其具体定义为:将一个向量空间中的点用其中的一个有限子集来进行编码的过程。

什么是VQ?

作为示例,我们在不失一般性的情况下采用二维情况下的向量。 图 1 显示了空间中的一些向量。 与每个向量簇相关联的是一个代表性代码字。 每个代码字都位于其自己的 Voronoi 区域中。 为了说明,这些区域在图 1 中用假想线分隔。 给定一个输入向量,被选择来表示它的代码字是在同一个 Voronoi 区域中的码字。

相互欧几里德距离最近的点代表为码字

欧几里德距离定义为:

VQ如何在压缩中工作?

Vevtor quantizer由两个操作组成。 第一个是编码器,第二个是解码器。 编码器采用输入向量并输出提供最低失真的码字索引。 在这种情况下,通过评估输入向量与码本中每个码字之间的欧几里得距离,可以找到最低失真。 一旦找到最接近的码字,该码字的索引就会通过通道发送(该通道可以是计算机存储、通信通道等)。 当编码器接收到代码字的索引时,它用相关的代码字替换索引。 

在矢量量化编码中,关键是码本的建立和码字搜索算法,如果想对矢量量化有个整体的概览,强烈推荐《Handbook of Image and Video Processing》一书中Fundamentals of Vector Quantization章节。下面对矢量量化中两类典型的方法多阶段矢量量化、乘积量化以及乘积量化的改进做简单介绍。

codebook如何设计?

到目前为止,我们已经讨论了 VQ 的工作方式,但我们还没有讨论如何生成码本。 什么码字最能代表一组给定的输入向量? 应该选多少?

不幸的是,设计一个最能代表输入向量集的密码本是 NP 难的。 这意味着它需要在空间中穷尽搜索最佳可能的码字,并且随着码字数量的增加,搜索呈指数增长(如果你能在多项式时间内找到最佳解决方案,你的名字将永远载入史册)。 因此,我们求助于次优码本设计方案,第一个想到的是最简单的。 它以 Linde-Buzo-Gray 的名字命名为 LBG,Linde-Buzo-Gray 是这个想法的作者。 该算法类似于k-means算法。

算法如下,

  1. 确定码字数 N 或码本的大小。

2. 随机选择N个码字,将其作为初始码本。 可以从一组输入向量中随机选择初始码字。

3. 使用欧几里得距离度量将每个码字周围的向量聚类。 这是通过获取每个输入向量并找到它与每个码字之间的欧几里德距离来完成的。 输入向量属于产生最小距离的码字簇。

4. 计算新的码字集。 这是通过获取每个集群的平均值来完成的。 添加每个向量的分量并除以群集中的向量数。

重复2和3直到所有码字不再变化或者变化很小为止。

该算法是迄今为止最受欢迎的,这是由于它的简单性。 虽然它是局部最优的,但速度很慢。 它慢的原因是因为对于每次迭代,确定每个聚类需要将每个输入向量与码本中的所有码字进行比较。

典型的方法:

下面对矢量量化中两类典型的方法多阶段矢量量化、乘积量化以及乘积量化的改进做简单介绍。

1、多阶段矢量量化:

多阶段矢量量化(Multi-Stage Vector Quantization,MSVQ)也称为残差矢量量化(Residual Vector Quantization, RVQ),它是一种思想,即将编码任务分解为一系列级联过程。级联过程可以用下图直观的展示出来:

如上图所示,对于待量化的向量x,经过一级量化器quantizer1后,得到的量化残差为r1 = x – C1b1,其中C1为一级量化器的码本,b1为x经过一级量化器quantizer1后的表示结果,将一级量化误差r1作为二级量化器的输入,后面过程与此类似。通过这种级联量化的量化方式,当构建的量化器为无穷个时,x可以被这无穷个码本精确表示。上图右侧子图比较直观的描绘了x被多个码本逐步近似的过程。

上述 C1、C2、…、Ci、… 这些码本在构建的时候,可以采用KMeans等方式得到各个量化器的码本。以上面构建的4个级联的码本为例,当得到码本C1、C2、C3、C4后,x量化的结果即可用[b1, b2, b3, b4]表示。对于xq查询向量与x距离的计算,在计算xq与 C1、C2、…、Ci、… 之间的内积距离表后,可以通过查表的方式,获取到非对称距离。

这种多阶段级联的矢量量化方式,相比单阶段一次性量化,极大的降低了码本在训练过程中消耗的计算资源。举个例子,4个阶段的MSVQ,每阶段用KMeans只需构建构建256大小的码本,则对空间分割构建的cell数目为256256256256,效率是很高的,但是如果采用单阶段一次性量化构建4294967296大小的码本,这个码本根本没法用KMeans聚出来。此外在计算距离的时候,采用4阶段的MSVQ方式,只需计算4256次距离的计算构成距离表,然后采用查表方式计算距离,而单阶段一次性量化需要计算4294967296次的距离计算。MSVQ的进一步加速版本是倒排MSVQ,将一级码本视为倒排链,从而构建倒排结构,构建MSVQ倒排结构。

我们可以将MSVQ类比成“深度加深”的过程,下面介绍的非常经典的乘积量化方法,可以为“宽度加宽”的过程。

2、乘积量化:

乘积量化(Product Quantization,PQ)是Herve Jegou在2011年提出的一种非常经典实用的矢量量化索引方法,在工业界向量索引中已得到广泛的引用,并作为主要的向量索引方法,在Fasis有非常高效的实现。乘积量化的核心思想是分段(划分子空间)和聚类,或者说具体应用到ANN近似最近邻搜索上,KMeans是PQ乘积量化子空间数目为1的特例。PQ乘积量化生成码本和量化的过程可以用如下图示来说明:

在训练阶段,针对N个训练样本,假设样本维度为128维,我们将其切分为4个子空间,则每一个子空间的维度为32维,然后我们在每一个子空间中,对子向量采用K-Means对其进行聚类(图中示意聚成256类),这样每一个子空间都能得到一个码本。这样训练样本的每个子段,都可以用子空间的聚类中心来近似,对应的编码即为类中心的ID。如图所示,通过这样一种编码方式,训练样本仅使用的很短的一个编码得以表示,从而达到量化的目的。对于待编码的样本,将它进行相同的切分,然后在各个子空间里逐一找到距离它们最近的类中心,然后用类中心的id来表示它们,即完成了待编码样本的编码。

正如前面所说的,在矢量量化编码中,关键是码本的建立和码字的搜索算法,在上面,我们得到了建立的码本以及量化编码的方式。剩下的重点就是查询样本与dataset中的样本距离如何计算的问题了。

在查询阶段,PQ同样在计算查询样本与dataset中各个样本的距离,只不过这种距离的计算转化为间接近似的方法而获得。PQ乘积量化方法在计算距离的时候,有两种距离计算方式,一种是对称距离,另外一种是非对称距离。非对称距离的损失小(也就是更接近真实距离),实际中也经常采用这种距离计算方式。下面过程示意的是查询样本来到时,以非对称距离的方式(红框标识出来的部分)计算到dataset样本间的计算示意:

具体地,查询向量来到时,按训练样本生成码本的过程,将其同样分成相同的子段,然后在每个子空间中,计算子段到该子空间中所有聚类中心得距离,如图中所示,可以得到4*256个距离,这里为便于后面的理解说明,可以把这些算好的距离称作距离表。在计算库中某个样本到查询向量的距离时,比如编码为(124, 56, 132, 222)这个样本到查询向量的距离时,我们分别到距离表中取各个子段对应的距离即可,比如编码为124这个子段,在第1个算出的256个距离里面把编号为124的那个距离取出来就可,所有子段对应的距离取出来后,将这些子段的距离求和相加,即得到该样本到查询样本间的非对称距离。所有距离算好后,排序后即得到我们最终想要的结果。

从上面这个过程可以很清楚地看出PQ乘积量化能够加速索引的原理:即将全样本的距离计算,转化为到子空间类中心的距离计算。比如上面所举的例子,原本brute-force search的方式计算距离的次数随样本数目N成线性增长,但是经过PQ编码后,对于耗时的距离计算,只要计算4*256次,几乎可以忽略此时间的消耗。另外,从上图也可以看出,对特征进行编码后,可以用一个相对比较短的编码来表示样本,自然对于内存的消耗要大大小于brute-force search的方式。

在某些特殊的场合,我们总是希望获得精确的距离,而不是近似的距离,并且我们总是喜欢获取向量间的余弦相似度(余弦相似度距离范围在[-1,1]之间,便于设置固定的阈值),针对这种场景,可以针对PQ乘积量化得到的前top@K做一个brute-force search的排序。

3、倒排乘积量化

倒排PQ乘积量化(IVFPQ)是PQ乘积量化的更进一步加速版。其加速的本质逃不开在最前面强调的是加速原理:brute-force搜索的方式是在全空间进行搜索,为了加快查找的速度,几乎所有的ANN方法都是通过对全空间分割,将其分割成很多小的子空间,在搜索的时候,通过某种方式,快速锁定在某一(几)子空间,然后在该(几个)子空间里做遍历。在上一小节可以看出,PQ乘积量化计算距离的时候,距离虽然已经预先算好了,但是对于每个样本到查询样本的距离,还是得老老实实挨个去求和相加计算距离。但是,实际上我们感兴趣的是那些跟查询样本相近的样本(姑且称这样的区域为感兴趣区域),也就是说老老实实挨个相加其实做了很多的无用功,如果能够通过某种手段快速将全局遍历锁定为感兴趣区域,则可以舍去不必要的全局计算以及排序。倒排PQ乘积量化的”倒排“,正是这样一种思想的体现,在具体实施手段上,采用的是通过聚类的方式实现感兴趣区域的快速定位,在倒排PQ乘积量化中,聚类可以说应用得淋漓尽致。

倒排PQ乘积量化整个过程如下图所示:

在PQ乘积量化之前,增加了一个粗量化过程。具体地,先对N个训练样本采用KMeans进行聚类,这里聚类的数目一般设置得不应过大,一般设置为1024差不多,这种可以以比较快的速度完成聚类过程。得到了聚类中心后,针对每一个样本x_i,找到其距离最近的类中心c_i后,两者相减得到样本x_i的残差向量(x_i-c_i),后面剩下的过程,就是针对(x_i-c_i)的PQ乘积量化过程,此过程不再赘述。

在查询的时候,通过相同的粗量化,可以快速定位到查询向量属于哪个c_i(即在哪一个感兴趣区域),然后在该感兴趣区域按上面所述的PQ乘积量化距离计算方式计算距离。

4、最优乘积量化

最优乘积量化(Optimal Product Quantization, OPQ)是PQ的一种改进版本。其改进体现在,致力于在子空间分割时,对各子空间的方差进行均衡。在具体实现的时候,我们可以将Optimal的过程实现为一个组件。

通常,用于检索的原始特征维度较高,所以实际在使用PQ等方法构建索引的时候,常会对高维的特征使用PCA等降维方法对特征先做降维处理,这样降维预处理,可以达到两个目的:一是降低特征维度;二是在对向量进行子段切分的时候要求特征各个维度是不相关的,做完PCA之后,可以一定程度缓解这个问题。但是这么做了后,在切分子段的时候,采用顺序切分子段仍然存在一定的问题,这个问题可以借用ITQ中的一个二维平面的例子加以说明:

如上面a图所示,对于PCA降维后的二维空间,假设在做PQ的时候,将子段数目设置为2段,即切分成x和y两个子向量,然后分别在x和y上做聚类(假设聚类中心设置为2)。对a图和c图聚类的结果进行比较,可以明显的发现,a图在y方向上聚类的效果明显差于c图,而PQ又是采用聚类中心来近似原始向量(这里指降维后的向量),也就是c图是我们需要的结果。这个问题可以转化为数据方差来描述:在做PQ编码时,对于切分的各个子空间,我们应尽可能使得各个子空间的方差比较接近,最理想的情况是各个子空间的方差都相等。上图a图中,x和y各个方向的方差明显是差得比较大的,而对于c图,x和y方向各个方向的方差差不多是比较接近的。

为了在切分子段的时候,使得各个子空间的方差尽可能的一致,Herve Jegou在Aggregating local descriptors into a compact image representation中提出使用一个正交矩阵来对PCA降维后的数据再做一次变换,使得各个子空间的方差尽可能的一致。其对应的待优化目标函数见论文的第5页,由于优化该目标函数极其困难,Herve Jegou使用了Householder矩阵来得到该正交矩阵,但是得到的该正交矩阵并不能很好的均衡子空间的方差。

OPQ致力于解决的问题正是对各个子空间方差的均衡。具体到方法上,OPQ借鉴了ITQ的思想,在聚类的时候对聚类中心寻找对应的最优旋转矩阵,使得所有子空间中各个数据点到对应子空间的类中心的L2损失的求和最小。OPQ在具体求解的时候,分为非参求解方法和带参求解方法,具体为:

  • 非参求解方法。跟ITQ的求解过程一样。
  • 带参求解方法。带参求解方法假设数据服从高斯分布,在此条件下,最终可以将求解过程简化为数据经过PCA分解后,特征值如何分组的问题。在实际中,该解法更具备高实用性。

从上面可以看到,倒排乘积量化IVFPQ可以视为1阶段的MSVQ和PQ的结合版本,而OPQ是PQ对子空间方差均衡的改进。基于这样一种普适性的视角,可以构建一种矢量量化框架,MSVQ、PQ、OPQ中的O,都是该矢量量化框架中的基础组件,通过这些组件的组合,我们可以敏捷的得到上面介绍方法的各种实现。

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)

更新:

PyTorch3D:面向3D计算机视觉的PyTorch工具箱

PyTorch3D通过PyTorch为3D计算机视觉研究提供高效,可重复使用的组件。目前已基于此开发了:Mesh R-CNN、SynSin等模型。

Facebook开源了一个专门用于3D模型学习的库pytorch3d,说白了就是将3d中一些常用的操作封装起来了。那这个玩意到底有啥用呢?使用这个库又能完成什么事情呢?个人觉得这个库还是蛮有用的,它将一些常用的3D概念整理在一起,并且通过使用这个库可以完成一些基于3D的创作,对于学习入门3D的视觉生成、渲染、甚至是3d的目标检测、3维的姿态评估都大有裨益。

Pytorch3D_上手学习3D的AI模型

Accelerating 3D Deep Learning with PyTorch3D

文档:Welcome to PyTorch3D’s documentation!
项目链接:facebookresearch/pytorch3d
论文:https://arxiv.org/abs/2007.08501

PyTorch3D

主要功能包括:

  • 用于存储和操作 triangle meshes的数据结构
  • 在 triangle meshes上的有效操作(投影变换,图卷积,采样,损失函数)
  • 可微的mesh渲染器

PyTorch3D旨在与深度学习方法稳定集成,以预测和处理3D数据。 因此,PyTorch3D中的所有运算符:

  • 使用PyTorch张量实现
  • 可以处理小批量的异构数据
  • 可以differentiated
  • 可以利用GPU进行加速

深度学习已大大改善了2D图像识别。扩展到3D可能会推动许多新应用的发展,包括自动驾驶汽车,虚拟现实和增强现实,创作3D内容,甚至改善2D识别。然而,尽管兴趣日益浓厚,但3D深度学习仍相对未得到充分开发。我们认为,这种差异是由于3D深度学习所涉及的工程挑战所致,例如有效处理异构数据和将图形操作重构为可微的。

我们通过引入PyTorch3D来应对这些挑战,PyTorch3D是一个用于3D深度学习的模块化,高效且可微的运算符库。它包括一个用于网格和点云的快速,模块化,可微的渲染器,支持按合成进行分析的方法。

与其他可微的渲染器相比,PyTorch3D更具模块化和效率,允许用户更轻松地扩展它,同时还可以优雅地缩放到较大的网格和图像。我们将PyTorch3D运算符和渲染器与其他实现进行了比较,并展示了显著的速度和内存改进。我们还使用PyTorch3D改进了ShapeNet上2D图像的无监督3D网格和点云预测的最新技术。

PyTorch3D是开源的,我们希望它将有助于加速3D深度学习的研究。

实验结果

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 的导出。

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 算子的行为。