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

NVIDIA TensorRT—推理引擎加速深度学习推理

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

深度学习的工作流程,如下图所示,可分为训练和推理两个部分。

训练过程通过设定数据处理方式,并设计合适的网络模型结构以及损失函数和优化算法,在此基础上将数据集以小批量的方式(mini-batch)反复进行前向计算并计算损失,然后 反向计算梯度利用特定的优化函数来更新模型,来使得损失函数达到最优的结果。训练过程最重要的就是梯度的计算和反向传播。

而推理就是在训练好的模型结构和参数基础上,做一次前向传播得到模型输出的过程。相对于训练而言,推理不涉及梯度和损失优化。推理的最终目标是将训练好的模型部署生产环境中。

高性能推理引擎的工作项

虽然推理就是数据经过模型的一次前向计算,但是推理是面向不同的终端部署,一般推理需要满足:

  • 精度要求: 推理的精度需要和训练的精度保持一致,
  • 效率要求:性能尽可能的快
  • 异构的推理设备:生产环境因为场景不同,支持不同的设备如TPU,CPU,GPU, NPU等

所以推理框架一般包括模型优化和推理加速,以便于支持高性能的推理要求。

那么一个推理框架要做哪些事情呢?

首先,因为推理框架要支持现有流行的深度学习框架如TensorFlow和Pytorch等,而不同的深度学习框内在的不一致性,就要求推理框架需要有一种同一个表达形式,来统一外部的不一致性,这就需要推理框架外部模型解析和转换为内在形式的功能。

其次,为了追求性能的提升,需要能够对训练好的模型针对特定推理设备进行特定的优化,主要优化可以包括

  • 低精度优化:FP16低精度转换,INT8后训练量化
  • 算子编译优化
  • 内存优化
  • 计算图调度

低精度优化

一般模型训练过程中都是采用FP32或者FP64高精度的方式进行存储模型参数,主要是因为梯度计算更新的可能是很小的一个小数。高精度使得模型更大,并且计算很耗时。而在推理不需要梯度更新,所以通常如果精度从FP32降低到FP16,模型就会变小很多,并且计算量也下降,而相对于模型的推理效果几乎不会有任何的变化,一般都会做FP16的精度裁剪

而FP32如果转换到INT8,推理性能会提高很多,但是裁剪不是直接裁剪,参数变动很多,会影响模型的推理效果,需要做重新的训练,来尽可能保持模型的效果

算子编译优化

我们先来了解下计算图的概念,计算图是由算子和张量构建成一个数据计算流向图,通常深度学习网络都可以看成一个计算图。而推理可以理解成数据从计算图起点到终点的过程。

算子编译优化其中一项优化就是计算图的优化。计算图优化的目标是对计算图进行等价的组合变换,使得减少算子的读写操作提供效率。

最简单的情况,就是算子融合。比如常见Conv+ReLu的两个算子,因为Conv需要做大量卷积计算,需要密集的计算单元支持,而Relu几乎不需要计算,如果Relu算子单独运算,则不仅需要一个计算单元支持其实不需要怎么计算的算子,同时又要对前端的数据进行一次读操作,很浪费资源和增加I/O操作; 此时,可以将Conv和Relu合并融合成一个算子,可以节省I/O访问和带宽开销,也可以节省计算单元。

这种算子融合对于所有推理设备都是支持,是通用的硬件优化。有些是针对特定硬件优化,比如某些硬件的计算单元不支持过大算子输入,此时就需要对算子进行拆解。

计算图的优化可以总结为算子拆解、算子聚合、算子重建,以便达到在硬件设备上更好的性能。

算子编译优化的另一个优化就是数据排布优化。我们知道,在TensorFlow框架的输入格式NHWC,而pytorch是NCHW。这些格式是框架抽象出来的矩阵格式,实际在内存中的存储都是按照1维的形式存储。这就涉及物理存储和逻辑存储之间的映射关系,如何更好的布局数据能带来存储数据的访问是一个优化方向;另外在硬件层面,有些硬件在某种存储下有最佳的性能,通常可以根据硬件的读写特点进行优化。

内存优化

我们推理的时候都需要借助额外的硬件设备来达到高速推理,如GPU,NPU等,此时就需要再CPU和这些硬件设备进行交互;以GPU为例,推理时需要将CPU中的数据copy到GPU显存中,然后进行模型推理,推理完成后的数据是在GPU显存中,此时又需要将GPU显存中的数据copy回cpu中。

这个过程就涉及到存储设备的申请、释放以及内存对齐等操作,而这部分也是比较耗时的。

因此内存优化的方向,通常是减少频繁的设备内存空间的申请和尽量做到内存的复用。

一般的,可以根据张量生命周期来申请空间:

  • 静态内存分配:比如一些固定的算子在整个计算图中都会使用,此时需要再模型初始化时一次性申请完内存空间,在实际推理时不需要频繁申请操作,提高性能
  • 动态内存分配:对于中间临时的内存需求,可以进行临时申请和释放,节省内存使用,提高模型并发能力
  • 内存复用:对于同一类同一个大小的内存形式,又满足临时性,可以复用内存地址,减少内存申请。

计算图调度

在计算图中,存在某些算子是串行依赖,而某些算子是不依赖性;这些相互独立的子计算图,就可以进行并行计算,提高推理速度,这就是计算图的调度。

TensorRT

我们讲解了推理引擎的一般工作流程和优化思路,这一部分介绍一个具体的推理引擎框架:TensorRT。NVIDIA TensorRT 是一个用于深度学习推理的 SDK 。 TensorRT 提供了 API 和解析器,可以从所有主要的深度学习框架中导入经过训练的模型。然后,它生成可在数据中心以及汽车和嵌入式环境中部署的优化运行时引擎。TensorRT是NVIDIA出品的针对深度学习的高性能推理SDK。

目前,TensorRT只支持NVIDIA自家的设备的推理服务,如服务器GPUTesla v100、NVIDIA GeForce系列以及支持边缘的NVIDIA Jetson等。

TensorRT通过将现有深度学习框架如TensorFlow、mxnet、pytorch、caffe2以及theano等训练好的模型进行转换和优化,并生成TensorRT的运行时(Runtime Engine),利用TensorRT提供的推理接口(支持不同前端语言如c++/python等),部署不同的NVIDIA GPU设备上,提供高性能人工智能的服务。

在性能方面,TensorRT在自家的设备上提供了优越的性能:

对于TensorRT而言,主要优化如下:

  • 算子和张量的融合 Layer & Tensor Fusion

以上面Inception模块的计算图为例子,左边是未优化原始的结构图,右边是经过TensorRT优化过的计算图。优化的目标是减少GPU核数的使用,以便于减少GPU核计算需要的数据读写,提高GPU核数的计算效率

  • 首先是合并conv+bias+relu为一个CBR模块,减少2/3 核的使用
  • 然后是对于同一输入1x1conv,合并为一个大的CBR,输出保持不变,减少了2次的相同数据的读写
  • 有没有发现还少了一个concat层,这个是怎么做到的?concat操作可以理解为数据的合并,TensorRT采用预先先申请足够的缓存,直接把需要concat的数据放到相应的位置就可以达到concat的效果。

经过优化,使得整个模型层数更少,占用更少GPU核,运行效率更快。

  • 精度裁剪 Precision Calibration
    这个是所有推理引擎都有部分,TensorRT支持低精度FP16和INT8的模型精度裁剪,在尽量不降低模型性能的情况,通过裁剪精度,降低模型大小,提供推理速度。但需要注意的是:不一定FP16就一定比FP32的要快。这取决于设备的不同精度计算单元的数量,比如在GeForce 1080Ti设备上由于FP16的计算单元要远少于FP32的,裁剪后反而效率降低,而GeForce 2080Ti则相反。
  • Dynamic Tensor Memory: 这属于提高内存利用率
  • Multi-Stream Execution: 这属于内部执行进程控制,支持多路并行执行,提供效率
  • Auto-Tuning 可理解为TensorRT针对NVIDIA GPU核,设计有针对性的GPU核优化模型,如上面所说的算子编译优化。

 TensorRT安装

了解了TensorRT是什么和如何做优化,我们实际操作下TensorRT, 先来看看TensorRT的安装。

TensorRT是针对NVIDIA GPU的推理引擎,所以需要CUDA和cudnn的支持,需要注意版本的对应关系; 以TensorRT 7.1.3.4为例,需要至少CUDA10.2和cudnn 8.x。

本质上 TensorRT的安装包就是动态库文件(CUDA和cudnn也是如此),需要注意的是TensorRT提供的模型转换工具。

下载可参考

rpm -i cuda-repo-rhel7-10-2-local-10.2.89-440.33.01-1.0-1.x86_64.rpm
tar -zxvf cudnn-10.2-linux-x64-v8.0.1.13.tgz
# tar -xzvf TensorRT-${version}.Linux.${arch}-gnu.${cuda}.${cudnn}.tar.gz
tar -xzvf TensorRT-7.1.3.4.CentOS-7.6.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz

TensorRT也提供了python版本(底层还是c的动态库)

#1.创建虚拟环境 tensorrt
  conda create -n tensorrt python=3.6
  
  #安装其他需要的工具包, 按需包括深度学习框架
  pip install keras,opencv-python,numpy,tensorflow-gpu==1.14,pytorch,torchvision
 
#2. 安装pycuda
  #首先使用nvcc确认cuda版本是否满足要求: nvcc -V
  pip install 'pycuda>=2019.1.1'
       
#3. 安装TensorRT
  # 下载解压的tar包
  tar -xzvf TensorRT-7.1.3.4.CentOS-7.6.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz
  
  #解压得到 TensorRT-7.1.3.4的文件夹,将里面lib绝对路径添加到环境变量中
  export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/TensorRT-7.1.3.4/lib
  
  #安装TensorRT
  cd TensorRT-7.1.3.4/python
  pip install pip install tensorrt-7.1.3.4-cp36-none-linux_x86_64.whl
 
#4.安装UFF
  cd TensorRT-7.1.3.4/uff
  pip install uff-0.6.9-py2.py3-none-any.whl
 
#5. 安装graphsurgeon
  cd TensorRT-7.1.3.4/graphsurgeon
  pip install uff-0.6.9-py2.py3-none-any.whl
 
#6. 环境测试
  #进入python shell,导入相关包没有报错,则安装成功
  import tensorrt
  import uff

安装完成后,在该路径的samples/python给了很多使用tensorrt的python接口进行推理的例子(图像分类、目标检测等),以及如何使用不同的模型解析接口(uff,onnx,caffe)。

另外给了一个common.py文件,封装了tensorrt如何为engine分配显存,如何进行推理等操作,我们可以直接调用该文件内的相关函数进行tensorrt的推理工作。

TensorRT工作流程

在安装TensorRT之后,如何使用TensorRT呢?我们先来了解下TensorRT的工作流程

总体流程可以拆分成两块:

  • 模型转换
    TensorRT需要将不同训练框架训练出来的模型,转换为TensorRT支持的中间表达(IR),并做计算图的优化等,并序列化生成plan文件。
  • 模型推理:在模型转换好后之后,在推理时,需要加plan文件进行反序列化加载模型,并通过TensorRT运行时进行模型推理,输出结果

模型转换

由于不同的深度学习框架的实现逻辑不同,TensorRT在转换模型时采用不同适配方法。以当前最流行深度学习框架TensorFlow和Pytorch为例为例。

由于pytorch采用动态的计算图,也就是没有图的概念,需要借助ONNX生成静态图。

Open Neural Network Exchange(ONNX,开放神经网络交换)格式,是一个用于表示深度学习模型的标准,可使模型在不同框架之间进行转移.最初的ONNX专注于推理(评估)所需的功能。 ONNX解释计算图的可移植,它使用graph的序列化格式

pth 转换为onnx

import onnx
import torch
def export_onnx(onnx_model_path, model, cuda, height, width, dummy_input=None):
    model.eval()
    if dummy_input is None:
        dummy_input = torch.randn(1, 3, height, width).float()
    dummy_input.requires_grad = True
    print("dummy_input shape: ", dummy_input.shape, dummy_input.requires_grad)

    if cuda:
        dummy_input = dummy_input.cuda()

    torch.onnx.export(
        model,  # model being run
        dummy_input,  # model input (or a tuple for multiple inputs)
        onnx_model_path,  # where to save the model (can be a file or file-like object)
        export_params=True,  # store the trained parameter weights inside the model file
        opset_version=10,  # the ONNX version to export the model to
        do_constant_folding=True,  # whether to execute constant folding for optimization
        verbose=True,
        input_names=['input'],  # the model's input names
        output_names=['output'],  # the model's output names
    )

从上可知,onnx通过pytorch模型完成一次模型输入和输出的过程来遍历整个网络的方式来构建完成的计算图的中间表示。

这里需要注意三个重要的参数:

  • opset_version: 这个是onnx支持的op算子的集合的版本,因为onnx目标是在不同深度学习框架之间做模型转换的中间格式,理论上onnx应该支持其他框架的所有算子,但是实际上onnx支持的算子总是滞后的,所以需要知道那个版本支持什么算子,如果转换存在问题,大部分当前的版本不支持需要转换的算子。
  • input_names:模型的输入,如果是多个输入,用列表的方式表示,如[“input”, “scale”]
  • output_names: 模型的输出, 多个输出,通input_names

onnx转换为plan engine模型

这里给出的通过TensorRT的python接口来完成onnx到plan engine模型的转换。

import tensorrt as trt
def build_engine(onnx_path):
          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_batch_size = 128
            builder.max_workspace_size = 1<<15
            builder.fp16_mode = True
            builder.strict_type_constraints = True
            with open(onnx_path, 'rb') as model:
                parser.parse(model.read())
            # Build and return an engine.
            return builder.build_cuda_engine(network)

从上面的转换过程可知,TensortRT的转换涉及到几个关键的概念:builder 、 network 、parser

  • builder:TensorRT构建器,在构建器中设置模型,解析器和推理的参数设置等 trt.Builder(TRT_LOGGER)
  • network: TensorRT能识别的模型结构(计算图)
  • parser:这里是指解析onnx模型结构(计算图)

从总体上看,TensorRT的转换模型是,将onnx的模型结构(以及参数)转换到TensorRT的network中,同时设置模型推理和优化的参数(如精度裁剪等)。 用一张图来总结下上述过程:

保存engine和读取engine

#解析模型,构建engine并保存
with build_engine(onnx_path) as engine:
    with open(engine_path, "wb") as f:
        f.write(engine.serialize())

#直接加载engine   
with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

TensorFlow / Keras

TensorFlow或者Keras(后台为TensorFlow)采用的是静态的计算图,本身就有图的完整结构,一般模型训练过程会保留ckpt格式,有很多冗余的信息,需要转换为pb格式。针对TensorFlow,TensorRT提供了两种转换方式,一种是pb直接转换,这种方式加速效果有限所以不推荐;另一种是转换uff格式,加速效果明显。

  • 转换为pb
from tensorflow.python.framework import graph_io
from tensorflow.python.framework import graph_util
from tensorflow.python.platform import gfile
# 设置输出节点为固定名称
OUTPUT_NODE_PREFIX = 'output_'
NUMBER_OF_OUTPUTS = 1
#输入和输出节点名称
output_names = ['output_']
input_names = ['input_']
input_tensor_name = input_names[0] + ":0"
output_tensor_name = output_names[0] + ":0"

def keras_to_pb(model_path, pb_path):
    K.clear_session()#可以保持输入输出节点的名称每次执行都一致
    K.set_learning_phase(0)
    sess = K.get_session()
    try:
        model = load_model(model_path)# h5 model file_path
    except ValueError as err:
        print('Please check the input saved model file')
        raise err

    output = [None]*NUMBER_OF_OUTPUTS
    output_node_names = [None]*NUMBER_OF_OUTPUTS
    for i in range(NUMBER_OF_OUTPUTS):
        output_node_names[i] = OUTPUT_NODE_PREFIX+str(i)
        output[i] = tf.identity(model.outputs[i], name=output_node_names[i])
    
    try:
        frozen_graph = graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), output_node_names)
        graph_io.write_graph(frozen_graph, os.path.dirname(pb_path), os.path.basename(pb_path), as_text=False)
        print('Frozen graph ready for inference/serving at {}'.format(pb_path))
    except:
        print("error !")
  • pb 到uff

采用TensorRT提供的uff模块的from_tensorflow_frozen_model()将pb格式模型转换成uff格式模型

import uff
def pb_to_uff(pb_path, uff_path, output_names):
        uff_model = uff.from_tensorflow_frozen_model(pb_path, output_names, output_filename=uff_path)

uff转换成plan engine模型

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.INFO)
img_size_tr = (3,224,224) #CHW
input_names = ['input_0']
output_names = ['output_0']

def build_engine(uff_path):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
        builder.max_batch_size = 128 #must bigger than batch_size
        builder.max_workspace_size =1<<15  #cuda buffer size
        builder.fp16_mode = True  #set dtype: fp32, fp16, int8
        builder.strict_type_constraints = True
        # Parse the Uff Network
        parser.register_input(input_names[0], img_size_tr)#NCHW
        parser.register_output(output_names[0])
        parser.parse(uff_path, network)
        # Build and return an engine.
        return builder.build_cuda_engine(network)

在绑定完输入输出节点之后,parser.parse()可以解析uff格式文件,并保存相应网络到network。而后通过builder.build_cuda_engine()得到可以直接在cuda执行的engine文件。该engine文件的构建需要一定时间,可以保存下来,下次直接加载该文件,而不需要解析模型后再构建。

TensorFlow的模型转换基本和onnx是一样的,主要是解析器不一样是UffParser。

#解析模型,构建engine并保存
with build_engine(uff_path) as engine:
    with open(engine_path, "wb") as f:
        f.write(engine.serialize())

#直接加载engine   
with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

模型推理

通过TensorRT的模型转换后,外部训练好的模型都被TensorRT统一成TensorRT可识别的engine文件(并优化过)。在推理时,只要通过TensorRT的推理SDK就可以完成推理。

具体的推理过程如下:

  • 通过TensorRT运行时,加载转换好的engine
  • 推理前准备:(1)在CPU中处理好输入(如读取数据和标准化等)(2)利用TensorRT的推理SDK中common模块进行输入和输出GPU显存分配
  • 执行推理:(1)将CPU的输入拷贝到GPU中 (2)在GPU中进行推理,并将模型输出放入GPU显存中
  • 推理后处理:(1)将输出从GPU显存中拷贝到CPU中 (2)在CPU中进行其他后处理
import common
import numpy as np
import cv2
import tensorrt as trt

def inference_test(engine_path, img_file):

    # process input
    input_image = cv2.imread(img_file)
    input_image = input_image[..., ::-1] / 255.0
    input_image = np.expand_dims(input_image, axis=0)
    
    input_image = input_image.transpose((0, 3, 1, 2))  # NCHW for pytorch
    input_image = input_image.reshape(1, -1)  # .ravel()
        
    # infer
    batch_size = 1
    TRT_LOGGER = trt.Logger(trt.Logger.INFO)
    with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())
        # Allocate buffers and create a CUDA stream
        inputs, outputs, bindings, stream = common.allocate_buffers(engine, batch_size)
        # Contexts are used to perform inference.
        with engine.create_execution_context() as context:
             np.copyto(inputs[0].host, input_image)
             [output] = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream, batch_size=batch_size)

TensorRT进阶和缺点

前面较全面了介绍了TensorRT的特点(优点)和工作流程;希望能感受到TensorRT的魅力所在。

在实际代码中主要是通过python的接口来讲解,TensorRT也提供了C++的转换和推理方式,但是主要的关键概念是一样

那TensorRT有什么局限性吗?

首先,TensorRT只支持NVIDIA自家的设备,并根据自家设备的特点,做了很多的优化,如果是其他设备,TensorRT就不适用了。这时候可以考虑其他的推理框架,比如以推理编译为基础的TVM, 针对移动平台推理NCNN,MACE、MNN以及TFLite等,以及针对Intel CPU的OPENVINO。

其次,算子的支持程度;这几乎是所有第三方推理框架都遇到的问题,TensorRT在某些不支持的算子的情况下,TensorRT提供了plugin的方式,plugin提供了标准接口,允许自己开发新的算子,并以插件的方式加入TensorRT(后面会专门介绍,欢迎关注)。

总结

  • 训练需要前向计算和反向梯度更新,推理只需要前向计算
  • 推理框架优化:低精度优化、算子编译优化、内存优化、计算图调度
  • TensorRT是针对NVIDIA设备的高性能推理框架
  • TensorRT工作流程包括模型转换和模型推理
  • 针对Pytorch, TensorRT模型转换链路为:pth->onnx->trt plan
  • 针对TensorFlow,TensorRT模型转换链路为:ckpt->pb->uff->trt plan
  • TensorRT模型转换关键点为build,network和parse
  • TensorRT模型推理关键点为:tensorrt runtime,engine context,显存操作和推理

ONNX—-模型部署教程(1)

转自:mmdeploy

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

官网1:https://pytorch.org/docs/stable/onnx.html#functions

官网2:https://onnxruntime.ai/docs/get-started/

前言

OpenMMLab 的算法如何部署?是很多社区用户的困惑。而模型部署工具箱 MMDeploy 的开源,强势打通了从算法模型到应用程序这 “最后一公里”!

今天我们将开启模型部署入门系列教程,在模型部署开源库 MMDeploy 的辅助下,介绍以下内容:

  • 中间表示 ONNX 的定义标准
  • PyTorch 模型转换到 ONNX 模型的方法
  • 推理引擎 ONNX Runtime、TensorRT 的使用方法
  • 部署流水线 PyTorch – ONNX – ONNX Runtime/TensorRT 的示例及常见部署问题的解决方法
  • MMDeploy C/C++ 推理 SDK

希望通过本系列教程,带领大家学会如何把自己的 PyTorch 模型部署到 ONNX Runtime/TensorRT 上,并学会如何把 OpenMMLab 开源体系中各个计算机视觉任务的模型用 MMDeploy 部署到各个推理引擎上。

初识模型部署

在软件工程中,部署指把开发完毕的软件投入使用的过程,包括环境配置、软件安装等步骤。类似地,对于深度学习模型来说,模型部署指让训练好的模型在特定环境中运行的过程。相比于软件部署,模型部署会面临更多的难题:

1)运行模型所需的环境难以配置。深度学习模型通常是由一些框架编写,比如 PyTorch、TensorFlow。由于框架规模、依赖环境的限制,这些框架不适合在手机、开发板等生产环境中安装。

2)深度学习模型的结构通常比较庞大,需要大量的算力才能满足实时运行的需求。模型的运行效率需要优化。

因为这些难题的存在,模型部署不能靠简单的环境配置与安装完成。经过工业界和学术界数年的探索,模型部署有了一条流行的流水线:

为了让模型最终能够部署到某一环境上,开发者们可以使用任意一种深度学习框架来定义网络结构,并通过训练确定网络中的参数。之后,模型的结构和参数会被转换成一种只描述网络结构的中间表示,一些针对网络结构的优化会在中间表示上进行。最后,用面向硬件的高性能编程框架(如 CUDA,OpenCL)编写,能高效执行深度学习网络中算子的推理引擎会把中间表示转换成特定的文件格式,并在对应硬件平台上高效运行模型。

这一条流水线解决了模型部署中的两大问题:使用对接深度学习框架和推理引擎的中间表示,开发者不必担心如何在新环境中运行各个复杂的框架;通过中间表示的网络结构优化和推理引擎对运算的底层优化,模型的运算效率大幅提升。

中间表示 – ONNX

在介绍 ONNX 之前,我们先从本质上来认识一下神经网络的结构。神经网络实际上只是描述了数据计算的过程,其结构可以用计算图表示。比如 a+b 可以用下面的计算图来表示:

为了加速计算,一些框架会使用对神经网络“先编译,后执行”的静态图来描述网络。静态图的缺点是难以描述控制流(比如 if-else 分支语句和 for 循环语句),直接对其引入控制语句会导致产生不同的计算图。比如循环执行 n 次 a=a+b,对于不同的 n,会生成不同的计算图:

ONNX (Open Neural Network Exchange)是 Facebook 和微软在2017年共同发布的,用于标准描述计算图的一种格式。目前,在数家机构的共同维护下,ONNX 已经对接了多种深度学习框架和多种推理引擎。因此,ONNX 被当成了深度学习框架到推理引擎的桥梁,就像编译器的中间语言一样。由于各框架兼容性不一,我们通常只用 ONNX 表示更容易部署的静态图。

创建 PyTorch 模型

让我们用 PyTorch 实现一个超分辨率模型,并把模型部署到 ONNX Runtime 这个推理引擎上。

# 安装 ONNX Runtime, ONNX, OpenCV 
pip install onnxruntime onnx opencv-python

在一切都配置完毕后,用下面的代码来创建一个经典的超分辨率模型 SRCNN。

import os 
 
import cv2 
import numpy as np 
import requests 
import torch 
import torch.onnx 
from torch import nn 
 
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) 
 
        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): 
        x = self.img_upsampler(x) 
        out = self.relu(self.conv1(x)) 
        out = self.relu(self.conv2(out)) 
        out = self.conv3(out) 
        return out 
 
# Download checkpoint and test image 
urls = ['https://download.openmmlab.com/mmediting/restorers/srcnn/srcnn_x4k915_1x16_1000k_div2k_20200608-4186f232.pth', 
    'https://raw.githubusercontent.com/open-mmlab/mmediting/master/tests/data/face/000001.png'] 
names = ['srcnn.pth', 'face.png'] 
for url, name in zip(urls, names): 
    if not os.path.exists(name): 
        open(name, 'wb').write(requests.get(url).content) 
 
def init_torch_model(): 
    torch_model = SuperResolutionNet(upscale_factor=3) 
 
    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)).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.png", torch_output)

SRCNN 先把图像上采样到对应分辨率,再用 3 个卷积层处理图像。为了方便起见,我们跳过训练网络的步骤,直接下载模型权重(由于 MMEditing 中 SRCNN 的权重结构和我们定义的模型不太一样,我们修改了权重字典的 key 来适配我们定义的模型),同时下载好输入图片。为了让模型输出成正确的图片格式,我们把模型的输出转换成 HWC 格式,并保证每一通道的颜色值都在 0~255 之间。如果脚本正常运行的话,一幅超分辨率的人脸照片会保存在 “face_torch.png” 中。

在 PyTorch 模型测试正确后,我们来正式开始部署这个模型。我们下一步的任务是把 PyTorch 模型转换成用中间表示 ONNX 描述的模型。

让我们用下面的代码来把 PyTorch 的模型转换成 ONNX 格式的模型:

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

其中,torch.onnx.export 是 PyTorch 自带的把模型转换成 ONNX 格式的函数。让我们先看一下前三个必选参数:前三个参数分别是要转换的模型、模型的任意一组输入、导出的 ONNX 文件的文件名。转换模型时,需要原模型和输出文件名是很容易理解的,但为什么需要为模型提供一组输入呢?这就涉及到 ONNX 转换的原理了。从 PyTorch 的模型到 ONNX 的模型,本质上是一种语言上的翻译。直觉上的想法是像编译器一样彻底解析原模型的代码,记录所有控制流。但前面也讲到,我们通常只用 ONNX 记录不考虑控制流的静态图。因此,PyTorch 提供了一种叫做追踪(trace)的模型转换方法:给定一组输入,再实际执行一遍模型,即把这组输入对应的计算图记录下来,保存为 ONNX 格式。export 函数用的就是追踪导出方法,需要给任意一组输入,让模型跑起来。我们的测试图片是三通道,256×256大小的,这里也构造一个同样形状的随机张量。

剩下的参数中,opset_version 表示 ONNX 算子集的版本。深度学习的发展会不断诞生新算子,为了支持这些新增的算子,ONNX会经常发布新的算子集,目前已经更新15个版本。我们令 opset_version = 11,即使用第11个 ONNX 算子集,是因为 SRCNN 中的 bicubic (双三次插值)在 opset11 中才得到支持。剩下的两个参数 input_names, output_names 是输入、输出 tensor 的名称,我们稍后会用到这些名称。

如果上述代码运行成功,目录下会新增一个”srcnn.onnx”的 ONNX 模型文件。我们可以用下面的脚本来验证一下模型文件是否正确。

import onnx 
 
onnx_model = onnx.load("srcnn.onnx") 
try: 
    onnx.checker.check_model(onnx_model) 
except Exception: 
    print("Model incorrect") 
else: 
    print("Model correct")

其中,onnx.load 函数用于读取一个 ONNX 模型。onnx.checker.check_model 用于检查模型格式是否正确,如果有错误的话该函数会直接报错。我们的模型是正确的,控制台中应该会打印出”Model correct”。

接下来,让我们来看一看 ONNX 模型具体的结构是怎么样的。我们可以使用 Netron (开源的模型可视化工具)来可视化 ONNX 模型。把 srcnn.onnx 文件从本地的文件系统拖入网站,即可看到如下的可视化结果:

点击 input 或者 output,可以查看 ONNX 模型的基本信息,包括模型的版本信息,以及模型输入、输出的名称和数据类型。

点击某一个算子节点,可以看到算子的具体信息。比如点击第一个 Conv 可以看到:

每个算子记录了算子属性、图结构、权重三类信息。

  • 算子属性信息即图中 attributes 里的信息,对于卷积来说,算子属性包括了卷积核大小(kernel_shape)、卷积步长(strides)等内容。这些算子属性最终会用来生成一个具体的算子。
  • 图结构信息指算子节点在计算图中的名称、邻边的信息。对于图中的卷积来说,该算子节点叫做 Conv_2,输入数据叫做 11,输出数据叫做 12。根据每个算子节点的图结构信息,就能完整地复原出网络的计算图。
  • 权重信息指的是网络经过训练后,算子存储的权重信息。对于卷积来说,权重信息包括卷积核的权重值和卷积后的偏差值。点击图中 conv1.weight, conv1.bias 后面的加号即可看到权重信息的具体内容。

现在,我们有了 SRCNN 的 ONNX 模型。让我们看看最后该如何把这个模型运行起来。

推理引擎 -ONNX Runtime

ONNX Runtime 是由微软维护的一个跨平台机器学习推理加速器,也就是我们前面提到的”推理引擎“。ONNX Runtime 是直接对接 ONNX 的,即 ONNX Runtime 可以直接读取并运行 .onnx 文件, 而不需要再把 .onnx 格式的文件转换成其他格式的文件。也就是说,对于 PyTorch – ONNX – ONNX Runtime 这条部署流水线,只要在目标设备中得到 .onnx 文件,并在 ONNX Runtime 上运行模型,模型部署就算大功告成了。

通过刚刚的操作,我们把 PyTorch 编写的模型转换成了 ONNX 模型,并通过可视化检查了模型的正确性。最后,让我们用 ONNX Runtime 运行一下模型,完成模型部署的最后一步。

ONNX Runtime 提供了 Python 接口。接着刚才的脚本,我们可以添加如下代码运行模型:

import onnxruntime 
 
ort_session = onnxruntime.InferenceSession("srcnn.onnx") 
ort_inputs = {'input': input_img} 
ort_output = ort_session.run(['output'], 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.png", ort_output)

这段代码中,除去后处理操作外,和 ONNX Runtime 相关的代码只有三行。让我们简单解析一下这三行代码。onnxruntime.InferenceSession用于获取一个 ONNX Runtime 推理器,其参数是用于推理的 ONNX 模型文件。推理器的 run 方法用于模型推理,其第一个参数为输出张量名的列表,第二个参数为输入值的字典。其中输入值字典的 key 为张量名,value 为 numpy 类型的张量值。输入输出张量的名称需要和torch.onnx.export 中设置的输入输出名对应。

如果代码正常运行的话,另一幅超分辨率照片会保存在”face_ort.png”中。这幅图片和刚刚得到的”face_torch.png”是一模一样的。这说明 ONNX Runtime 成功运行了 SRCNN 模型,模型部署完成了!以后有用户想实现超分辨率的操作,我们只需要提供一个 “srcnn.onnx” 文件,并帮助用户配置好 ONNX Runtime 的 Python 环境,用几行代码就可以运行模型了。或者还有更简便的方法,我们可以利用 ONNX Runtime 编译出一个可以直接执行模型的应用程序。我们只需要给用户提供 ONNX 模型文件,并让用户在应用程序选择要执行的 ONNX 模型文件名就可以运行模型了。

总结

  • 模型部署,指把训练好的模型在特定环境中运行的过程。模型部署要解决模型框架兼容性差和模型运行速度慢这两大问题。
  • 模型部署的常见流水线是“深度学习框架-中间表示-推理引擎”。其中比较常用的一个中间表示是 ONNX。
  • 深度学习模型实际上就是一个计算图。模型部署时通常把模型转换成静态的计算图,即没有控制流(分支语句、循环语句)的计算图。
  • PyTorch 框架自带对 ONNX 的支持,只需要构造一组随机的输入,并对模型调用 torch.onnx.export 即可完成 PyTorch 到 ONNX 的转换。
  • 推理引擎 ONNX Runtime 对 ONNX 模型有原生的支持。给定一个 .onnx 文件,只需要简单使用 ONNX Runtime 的 Python API 就可以完成模型推理。

关于模型部署

当我们千辛万苦完成了前面的数据获取、数据清洗、模型训练、模型评估等等步骤之后,终于等到“上线”啦。想到辛苦训练出来的模型要被调用还有点小激动呢,可是真当下手的时候就有点懵了:模型要怎么部署?部署在哪里?有什么限制或要求?

模型训练重点关注的是如何通过训练策略来得到一个性能更好的模型,其过程似乎包含着各种“玄学”,被戏称为“炼丹”。整个流程包含从训练样本的获取(包括数据采集与标注),模型结构的确定,损失函数和评价指标的确定,到模型参数的训练,这部分更多是业务方去承接相关工作。一旦“炼丹”完成(即训练得到了一个指标不错的模型),如何将这颗“丹药”赋能到实际业务中,充分发挥其能力,这就是部署方需要承接的工作。

目前来说,我还没有真正的接触工业界的模型应用,仅仅只是在学术界进行模型训练和探索。部署只是简单的将模型推理代码直接部署在服务器上,因此也没有考虑过模型在工业界的部署,对于工业应用来说,模型的表现和推理速度同等重要,因此,如何提高模型的推理速度成为模型能否落地的关键因素。

部署流程大致分为以下几个步骤:模型转换、模型量化压缩、模型打包封装 SDK。这里我们主要探讨模型转换。

模型转换主要用于模型在不同框架之间的流转,常用于训练和推理场景的连接。目前主流的框架都以 ONNX 或者 caffe 为模型的交换格式,另外,根据需要,还可以在中间插入计算图优化,对计算机进行推理加速(诸如常见的 CONV/BN 的算子融合)。

模型部署

在软件工程中,部署指把开发完毕的软件投入使用的过程,包括环境配置、软件安装等步骤。类似地,对于深度学习模型来说,模型部署指让训练好的模型在特定环境中运行的过程。相比于软件部署,模型部署会面临更多的难题:

1)运行模型所需的环境难以配置。深度学习模型通常是由一些框架编写,比如 PyTorch、TensorFlow。由于框架规模、依赖环境的限制,这些框架不适合在手机、开发板等生产环境中安装。

2)深度学习模型的结构通常比较庞大,需要大量的算力才能满足实时运行的需求。模型的运行效率需要优化。

因为这些难题的存在,模型部署不能靠简单的环境配置与安装完成。经过工业界和学术界数年的探索,模型部署有了一条流行的流水线:

为了让模型最终能够部署到某一环境上,开发者们可以使用任意一种深度学习框架来定义网络结构,并通过训练确定网络中的参数。之后,模型的结构和参数会被转换成一种只描述网络结构的中间表示,一些针对网络结构的优化会在中间表示上进行。最后,用面向硬件的高性能编程框架(如 CUDA,OpenCL)编写,能高效执行深度学习网络中算子的推理引擎会把中间表示转换成特定的文件格式,并在对应硬件平台上高效运行模型。

这一条流水线解决了模型部署中的两大问题:使用对接深度学习框架和推理引擎的中间表示,开发者不必担心如何在新环境中运行各个复杂的框架;通过中间表示的网络结构优化和推理引擎对运算的底层优化,模型的运算效率大幅提升。

MonoViT—基于ViT的自监督单目深度估计

Self-Supervised Monocular Depth Estimation witha Vision Transformer

paper: https://arxiv.org/pdf/2208.03543.pdf

基于深度学习单目深度估计任务简介

深度估计是计算机视觉领域的一个基础性问题,其可以应用在机器人导航、增强现实、三维重建、自动驾驶等领域。而目前大部分深度估计都是基于二维RGB图像到RBG-D图像的转化估计,主要包括从图像明暗、不同视角、光度、纹理信息等获取场景深度形状的Shape from X方法,还有结合SFM(Structure from motion)和SLAM(Simultaneous Localization And Mapping)等方式预测相机位姿的算法。其中虽然有很多设备可以直接获取深度,但是设备造价昂贵。也可以利用双目进行深度估计,但是由于双目图像需要利用立体匹配进行像素点对应和视差计算,所以计算复杂度也较高,尤其是对于低纹理场景的匹配效果不好。而单目深度估计则相对成本更低,更容易普及。

那么对于单目深度估计,顾名思义,就是利用一张或者唯一视角下的RGB图像,估计图像中每个像素相对拍摄源的距离。对于人眼来说,由于存在大量的先验知识,所以可以从一只眼睛所获取的图像信息中提取出大量深度信息。那么单目深度估计不仅需要从二维图像中学会客观的深度信息,而且需要提取一些经验信息,后者则对于数据集中相机和场景会比较敏感。

摘要:

自监督单眼深度估计是一种有吸引力的解决方案,它不需要难以获取的深度标签来进行训练。 卷积神经网络 (CNN) 最近在这项任务中取得了巨大成功。 然而,它们有限的接受域限制了现有的网络架构只能在局部进行推理,从而削弱了自我监督范式的有效性。 鉴于 Vision Transformers (ViTs) 最近取得的成功,我们提出了 MonoViT,这是一个全新的框架,结合了 ViT 模型支持的全局推理和自监督单目深度估计的灵活性。 通过将普通卷积与 Transformer 块相结合,我们的模型可以在局部和全局进行推理,以更高的细节和准确性产生深度预测,从而使 MonoViT 在已建立的 KITTI 数据集上实现sota的性能。 此外,MonoViT 在 Make3D 和 Driving Stereo 等其他数据集上证明了其卓越的泛化能力。

介绍:

Transformers (ViTs)最近表现出杰出的目标检测和 语义分割等任务的结果,这要归功于它们能够建立像素之间的长距离关系,因此是全局感受野。另外,有相关工作将VIT应用于深度估计,但不是采用自监督单目深度估计。 本文弥补了这个缺失的步骤,提出了MonoViT architecture。它在其骨干网中结合了卷积层和最先进的 (SoTA) MPViT块【1】进而对图片中的局部信息(objects)和全局信息(前景和背景之间的关系,以及物体之间)进行建模。 该策略使我们能够消除由 CNN 编码器的有限感知域引起的瓶颈,产生自然更细粒度的预测。

作者在KITTI dataset进行实验,表现优于其他sota模型,还分析了模型泛化能力跨不同的数据集,将 MonoViT 与它在 Make3D 和 Driving-Stereo datasets进行比较,也突出显示了 MonoViT 的卓越泛化能力

模型架构:

Deep Network:

Joint CNN & Transformer Layer used in depthencoder:

PoseNet:

PoseNet 倾向于简单而有效的实现。 具体来说, PoseNet 使用 ResNet18【2】的轻量级结构。 接收相邻图像 [I, I†] 作为输入,输出视频序列相邻帧之间的 6 DoF 相对位姿 T。这个网络用于最终辅助计算loss,提供监督信息。

Loss损失函数

View reconstruction loss:

Smoothness loss. As in previous works, the edge-aware smoothness loss is used to improve the inverse depth map d:

【1】MPViT: Multi-Path Vision Transformer for Dense Prediction

【2】Deep residual learning for image recognition

单张图片实现MIP:Single-view view synthesis with multiplane images

CVPR 2020:https://single-view-mpi.github.io/

最近在视图合成方面的一项工作是在已知的视点上通过给定两个或更多的输入图像,利用深度学习来生成多平面图像(以相机为中心的分层三维展示)。我们将这种表示方法应用于单视角的视图合成,这是一个更具挑战性的问题,但可能有更广泛的应用。我们的方法直接从单个图像输入中学习预测多平面图像,并引入了尺度不变量视图合成的监督,使我们能够在在线视频上进行训练。我们展示了这种方法适用于几个不同的数据集,它还能额外生成合理的深度图,并能学习实现对背景层中的前景对象边缘后面的内容进行填充。

备注:SFM(Structure From Motion),主要基于多视觉几何原理,用于从运动中实现3D重建,也就是从无时间序列的2D图像中推算三维信息,是计算机视觉学科的重要分支。

单目图像深度估计 – 自监督方法

近几年有关单目图像深度识别的算法以CNN为主流,更细的说是以无监督的同时对深度、计算机角度、光流等同时计算的端到端深度网络为主流。所谓无监督其实是指在训练过程中不需要输入真实的深度值,这样做有一个好处就是目前能够测量到深度信息的传感器还不够精确,因此由不够精确的label训练出的model得到的预测结果必然不会特别令人满意;
所谓同时计算呢,在我理解是指在训练过程中,用一个能够表征时间序列上有前后关系的帧之间的差别的loss同时训练多个网络,而在得到model后每个网络可以单独使用。
很聪明,不同作用的网络相当于人为的特征提取过程,最后的预测基于这个人为的特征提取结果,但这种方法也有其缺点,我能想到的就是参数的增加,网络结构的复杂化和人为特征对最终预测结果有没有起引导作用只能用实验去证明。

详细说呢,首先,所谓的“自监督”虽然不需要输入真实深度信息,但需要输入双目摄像头获取到的同一时刻不同角度的图像或者前后帧图像。

自监督的单目深度估计包括: 基于双目训练的无监督模型 和 基于视频序列的无监督模型

基于双目训练的无监督模型:

UnSupervised Monocular Depth Estimation with Left-Right Consistency, CVPR, 2017

学习方法在单目深度估计上面有比较好的结果了,但还把深度估计问题作为一个有监督的回归问题。所以需要大量相对应的ground truth用来训练。记录有质量的各种场景的深度信息是一个比较难的事情。作者做了一个新方法,替代了现在直接用深度图数据训练。这个方法是用容易获得的双目立体视觉的角度。

作者提出了全新的训练函数目标函数,可以让卷积神经网络学习到深度估计,虽然没有深度信息的ground truth。利用一些对极几何约束,作者产生了视差损失。还发现只用图像重建结果会产生低质量的深度图。为了克服这个困难,作者构建了一个新颖的训练损失,可以加强左右视差图的一致性,这样能够提升性能和鲁棒性。作者的方法在KITTI数据集的单目深度估计上达到了state-of-the-art,超过用ground-truth深度训练的有监督方法。

基于视频序列的无监督模型:

UnSupervised Learning of Depth and Ego-Motion from Video,CVPR,2017

这篇文章提出了一种非监督的多功能网络,主要思想就像之前提到过的用一个loss同时训练两个网络。网络的结果如图,其中第一个网络可接受一幅图片作为输入,输出其对应的深度图片;第二个网络为姿态网络,接受t,t+1和t-1三个时刻三幅图片作为输入,输出从t到t+1和从t到t-1的相机姿态变化矩阵。

输入为前中后三帧连续的图片,同时训练两个网络,一个得到深度预测结果,一个得到视差矩阵结果

[ICCV 2019] Digging into Self-Supervised Monocular Depth Prediction

目前自监督深度学习取得最好进展的地方,一般说来自监督不需要标注,使用内在几何(通常是多视图几何)关系监督学习,从另一个侧面说明3d视觉才是视觉的本质。

这篇文章中的作者开发了一种方法,该方法使用深度估计和姿态估计网络的组合来预测单帧图像中的深度。 它通过在一系列运动的图像(包括单目和双目)序列上训练一个建立在自监督损失函数上的架构来实现,这一架构包括两个网络,一个用来在单目图像上预测深度,另一个在运动图像之间预测姿态。 此方法不需要标注训练数据集。 相反,它使用图像序列中的连续时间帧和姿态的重投影关系来进行训练。 稍后将更详细地描述重建过程。 论文的主要贡献是:

1.一种自动mask技术,可消除loss对不重要像素的注意力,减少它们的影响

2.用深度图修正光度重建误差

3.多尺度深度估计

模型架构:本文的方法使用深度网络和姿态网络。 深度网络是经典的U-Net [2]编码器-解码器模型结构。 编码器是经过预训练的ResNet模型,当然也可以考虑其他模型。 深度解码器将输出转换为深度值。作者使用基于ResNet18的姿势网络,该姿态网络经过修改后,可以使用两个彩色图像作为输入来预测单个6自由度相对姿势或旋转和平移参数。 姿势网络使用前后两帧而不是典型的立体图像对作为图像对输入。 它可以从序列中的前一帧和后一帧通过角度预测目标图像的外观。

训练:下图说明了模型的训练过程。

可以同时使用运动自监督和双目自监督

光度重建误差:目标图像位于第0帧,并且用于姿态估计过程的图像可以是前一帧或后一帧,也就是帧+1或帧-1。该损失是基于目标图像和重建的目标图像之间的相似性。重建过程通过使用姿态网络从源帧(帧+1或帧-1)计算转换矩阵开始。然后使用旋转和平移的信息来计算从源帧到目标帧的映射。最后使用从深度网络预测的目标图像的深度图和从姿势网络转换的矩阵,将其投影到具有固有内参矩阵K的摄像机中,以获取重建的目标图像。此过程需要先将深度图转换为3D点云,通过姿态将点云转换到另一个坐标系后再使用相机内参将3D点转换为2D点。所得的点用作从目标图像进行双线性插值的采样网格。

光度重建误差,ssim是经常用在图像相似度里的度量

这种loss的目的是减少目标图像和重建的目标图像之间的差异,在目标图像和重建的目标图像中,姿态和深度估计的过程中都需要使用它。

自动mask:最终的光度重建误差要乘以一个mask,该mask解决与假设相机在静态场景(例如静态场景)中移动的变化有关的问题。尤其是 一个物体正在以与相机相似的速度移动,或者在其他物体正在移动时照相机已停止,也就是那些在相机坐标系里静止的物体。 这些相对静止的物体理论上应该有无穷大的深度。 作者使用一种自动mask方法解决了这一问题,该方法可以过滤不会将外观从一帧更改为下一帧的像素,也就是那些和相机同步运动的像素。 mask是二进制的,如果目标图像和重建的目标图像之间的最小光度误差小于目标图像和源图像的最小光度误差,则为1,否则为0。

当相机是静止的时,这种方法会图像中的所有像素都被掩盖(实际场景中概率很低)。 当物体以与照相机相同的速度移动时,会导致图像中静止物体的像素被掩盖。

多尺度估计:作者将各个尺度的损失合并在一起。 将较低分辨率的深度图上采样到较高的输入图像分辨率,然后在较高的输入分辨率下重新投影,重新采样并计算光度误差。 作者声称,这使得各个比例尺上的深度图以实现相同的目标,即对目标图像进行精确的高分辨率重建。

其他形式的loss:作者加入了平均归一化的逆深度图值和输入/目标图像之间的边缘敏感的平滑度损失。 这鼓励模型学习尖锐的边缘并消除噪声。

Depth Prediction Without the Sensors: Leveraging Structure for Unsupervised Learning from Monocular Videos

项目地址:https://sites.google.com/view/struct2depth

目标运动建模:来自Google大脑的作者发表了该文章,该文章进一步扩展了Monodepth2。 它们通过预测单个目标而不是整个图像的运动来改善姿态网络估计。 因此,现在重建的图像序列不再是单个投影,而是组合在一起的一系列目标的投影。 这通过两个模型一个目标运动模型和一个相机运动估计网络(类似于前面几节中描述的姿态网络)来做到。 步骤如下:

1.预训练的MASK-RCNN [2]模型用于捕获潜在移动目标的语义分割。

2.使用二进制掩码从静态图像(帧-1,帧0和帧+1)中删除这些可能移动的对象

3.被掩盖的图像被发送到ego-motion(相机自身运动)网络,并输出帧-1和0与帧0和+1之间的转换矩阵。

4.mask过程可提取静态背景,然后提取ego-motion转换矩阵,而无需移动对象。 使用来自[3]的方程。

遮罩过程可提取静态背景,然后提取自我运动转换矩阵,而无需移动对象。

具体操作时需要注意以下几点

1.使用之前步骤3中产生的相机运动转换矩阵,并将其应用于帧-1和帧+1,用这个矩阵来变换帧0。

2.使用从步骤3得到的相机运动变换矩阵,并将其应用于可能移动的对象的分割mask到帧-1和帧+1,以获取每个目标相当于帧0的形变后的分割mask。

3.二进制掩码用于保持与变形分割掩码关联的像素。

4.mask图像与变形图像组合在一起,并传递到目标运动模型,该模型输出预测的对象运动。

结果显示必须知道相机必如何运动才能“解释”对象外观的变化。 然后,要根据目标运动建模过程的步骤4中生成的运动模型来移动对象。 最后,将形变后的对目标运动与形变后的静态背景结合起来,以获得最终的形变图像:

整体架构

学习目标尺度:虽然Monodepth2通过其自动mask技术解决了静态物体或以与照相机相同速度移动的物体的问题,但struct2 depth作者还是建议对模型进行约束,以识别物体的比例,从而改善物体运动的建模。

基于对象的类别(例如楼房)定义每个对象的比例loss,旨在基于对象比例的知识来限制深度。 loss是图像中对象的输出深度图与通过使用相机的焦距,基于对象类别的先验高度和图像中分割后的对象的实际高度计算出的近似深度图之间的差 ,两者均按目标图片的平均深度进行缩放

基于深度学习的单目深度估计综述

Monocular Depth Estimation

Monocular Depth Estimation is the task of estimating the depth value (distance relative to the camera) of each pixel given a single (monocular) RGB image. This challenging task is a key prerequisite for determining scene understanding for applications such as 3D scene reconstruction, autonomous driving, and AR. State-of-the-art methods usually fall into one of two categories: designing a complex network that is powerful enough to directly regress the depth map, or splitting the input into bins or windows to reduce computational complexity. The most popular benchmarks are the KITTI and NYUv2 datasets. Models are typically evaluated using RMSE or absolute relative error. 这项具有挑战性的任务是确定 3D 场景重建、自动驾驶和 AR 等应用场景理解的关键先决条件。

任务介绍

深度估计是计算机视觉领域的一个基础性问题,其可以应用在机器人导航、增强现实、三维重建、自动驾驶等领域。而目前大部分深度估计都是基于二维RGB图像到RBG-D图像的转化估计,主要包括从图像明暗、不同视角、光度、纹理信息等获取场景深度形状的Shape from X方法,还有结合SFM(Structure from motion)和SLAM(Simultaneous Localization And Mapping)等方式预测相机位姿的算法。其中虽然有很多设备可以直接获取深度,但是设备造价昂贵。也可以利用双目进行深度估计,但是由于双目图像需要利用立体匹配进行像素点对应和视差计算,所以计算复杂度也较高,尤其是对于低纹理场景的匹配效果不好。而单目深度估计则相对成本更低,更容易普及。

那么对于单目深度估计,顾名思义,就是利用一张或者唯一视角下的RGB图像,估计图像中每个像素相对拍摄源的距离。对于人眼来说,由于存在大量的先验知识,所以可以从一只眼睛所获取的图像信息中提取出大量深度信息。那么单目深度估计不仅需要从二维图像中学会客观的深度信息,而且需要提取一些经验信息,后者则对于数据集中相机和场景会比较敏感。

通过阅读文献,可以将基于深度学习的单目深度估计算法大致分为以下几类:

  • 监督算法

顾名思义,直接以2维图像作为输入,以深度图为输出进行训练::监督方法的监督信号基于深度图的地面真值,因此单目深度估计可以看作是一个回归问题。从单个深度图像设计神经网络来预测深度。利用预测深度图和实际深度图之间的差异来监督网络的训练 L2损失

上面给的例子是KITTI数据集中的一组例子,不过深度图可能看的不是很明显,我重新将深度图涂色之后:

深度网络通过近似真值的方法来学习场景的深度。基于不同结构和损失函数的方法:据我们所知,Eigen等人首先用CNNs解决单目深度估计问题。该体系结构由两个组成部分组成(全局粗尺度网络和局部精细尺度网络),在文献中用于从单个图像进行端到端的深度图预测。

基于条件随机场的方法:Li等人提出了一种基于多层的条件随机场(CRFs)的细化方法,该方法也被广泛应用于语义分割。在深度的估计中,考虑到深度的连续特征,可以广泛地使用CRF的深度信息,因此可以广泛地应用于深度的估计中。

基于对抗性学习的方法:由于提出的对抗性学习在数据生成方面的突出表现,近年来成为一个研究热点。各种算法、理论和应用已得到广泛发展。对抗式学习深度估计的框架如图所示。

file
  • 无监督算法

首先,所谓的“无监督”虽然不需要输入真实深度信息,但需要输入双目摄像头获取到的同一时刻不同角度的图像或者前后帧图像,只是这样就叫做无监督在我看来略显牵强。

有监督学习方法要求每幅RGB图像都有其对应的深度标签,而深度标签采集通常需要深度相机或激光雷达,前者范围受限后者成本昂贵。再者,采集的原始深度标签通常是一些稀疏的点,不能与原图很好的匹配。因此不用深度标签的无监督估计方法是近年的研究趋势,其基本思路是利用左右视图,结合对极几何与自动编码机的思想求解深度。 

由于深度数据的获取难度较高,所以目前有大量算法都是基于无监督模型的。即仅仅使用两个摄像机采集的双目图像数据进行联合训练。其中双目数据可彼此预测对方,从而获得相应的视差数据,再根据视差与深度的关系进行演化。亦或是将双目图像中各个像素点的对应问题看作是立体匹配问题进行训练。左视图-右视图示例:

视差,以我们人眼为例,两只眼睛看到的图像分别位于不同的坐标系。将手指从较远地方慢慢移动到眼前,会发现,手指在左眼的坐标系中越来越靠右,而在右眼坐标系中越来越靠左,这种差异性就是视差。与此同时,可以说明,视差与深度成反比。除此之外,由于摄像机参数也比较容易获取,所以也可以以相机位姿作为标签进行训练。

同时同一水平线上的两个照相机拍摄到的照片是服从以下物理规律的:

在图中, Z 为场景所距离我们的深度, X为三维场景映射到的二维图像平面,也就是最终我们得到的二维图像所在的平面。 f为相机的焦距。 b为两个相机之间的距离,Xl和 Xr 分别为相同物体在左右两个不同相机中成像的坐标。根据以上信息,和简单的三角形相似规律我们可以得到:

这种思路最先应用于使用单张图片生成新视角问题:DeepStereo 和 Deep3d之中, 在传统的视角生成问题之中,首先会利用两张图(或多张)求取图片之间的视差d,其次通过得到的视差(相当于三维场景)来生成新视角。

基于可解释性掩模的方法:基于投影函数的视图重建算法依赖于静态场景假设,即动态目标在相邻帧上的位置不满足投影函数,从而影响测光度误差和训练过程。

基于传统视觉里程计的方法:用传统的直接视觉里程计回归的位姿来辅助深度估计,而不是使用位姿网络估计的位姿。直接视觉里程计利用深度网络生成的深度图和一个三帧图像,通过最小化光度误差来估计帧间的姿态,然后将计算出的姿态发送回训练框架。因此,由于深度网络由更精确的姿态来监督,因此深度估计的精度显着提高。

基于多任务框架的方法:最近的方法在基本框架中引入了额外的多任务网络,如光流、物体运动和相机内参矩阵,作为一个附加的训练框架,加强了整个训练任务之间的关系

基于对抗学习的方法:将对抗学习框架引入到无监督的单目深度估计中。由于在无监督训练中没有真正的深度图。因此,将视图重建算法合成的图像和真实图像作为鉴别器的输入,而不是使用鉴别器来区分真实深度图和预测深度图。

  • Structure from motion/基于视频的深度估计(无监督学习)

这一部分中既包含了单帧视频的单目深度估计,也包含了多帧间视频帧的像素的立体匹配,从而近似获取多视角图像,对相机位姿进行估计。

评估指标:

在单目深度估计问题中,常用的精度评估指标有相对误差(REL)、均方根误差(RMS)、对数误差(LG)及阈值误差(% correct)

深度估计相关数据集 

在深度估计的研究中,由于室内外场景类型与深度范围具有较大的差异,对应不同的场景分别会构造不同的数据集

  • 真实场景数据集
    • NYU depth v2(来自纽约大学)是常用的室内数据集之一,
      • 选取了464个不同的场景,
      • 利用RGB相机和微软的Kinect深度相机同时采集室内场景的RGB信息深度信息,收集了407 024帧RGBD图像对构建数据集。
      • 由于红外相机和摄像机之间的位置偏差,深度相机采集的原始深度图存在缺失部分或是噪点,
      • 作者从中选取了1 449幅图像,利用着色算法对深度图进行填充得到稠密深度图,同时人工标注语义信息
    • Make3D(斯坦福大学)是常用的室外场景数据集之一,
      • 使用激光扫描仪采集室外场景的深度信息,
      • 选取的场景类型为白天的城市和自然风光,深度范围是5~81 m,大于该范围统一映射为81 m。
      • 数据集共包含534幅RGBD图像对,其中400幅用于训练,134幅用于测试。
    • KITTI(德国卡尔斯鲁厄理工学院和美国丰田技术研究院)自动驾驶领域常用的数据集之一,链接:http://www.cvlibs.net/datasets/kit
      • 包含深度数据标签
      • 通过一辆装配有2台高分辨率彩色摄像机、2台灰度摄像机、激光扫描仪和GPS定位系统的汽车采集数据,其中激光扫描仪的最大测量距离为120 m。
      • 图像场景包括卡尔斯鲁厄市、野外地区以及高速公路。
      • 数据集共包含93 000个RGBD训练样本。
    • Depth in the Wild(DIW)(密歇根大学)以相对深度作为标签的数据集
      • 从词典中随机选取单词作为搜索关键字,然后从互联网中收集得到原始的RGB图像。
      • 标注工作外包给专门的工作人员,为了更加高效,每一幅图像选取了两个高亮的点,工作人员只需判定两个点的远近关系即可。
      • 对于采样点对之间的位置关系,采用50%随机采样,另外50%对称采样的方法以使构建的数据集尽可能平衡。最终获得的有效标注图像约5×E11张。
    • Cityscapes. Cityscapes的数据取自德国的50多个城市的户外场景,其中数据包含有左右视角图像、视差深度图、相机校准、车辆测距、行人标定、目标分割等,同时也包含有类似于vKITTI的虚拟渲染场景图像。其中简单的左视角图像、相机标定、目标分割等数据需要利用学生账号注册获取,其他数据需要联系管理员获取。链接:https://www.cityscapes-dataset.com/
  • 虚拟场景数据集
    • SceneNet RGB-D数据集
    • SYNTHIA数据集
    • 由于是通过虚拟场景生成,数据集中包括更多天气、环境及光照,场景类型多样。各数据集有各自的优缺点,在实际研究中,应根据具体研究问题来选择合适的数据集。

 综上,可以看到基于深度学习的单目深度估计是本领域的发展方向。目前,该领域的发展主要集中在数据集和深度学习模型两方面。首先,数据集的质量在很大程度上决定了模型的鲁棒性与泛化能力,深度学习要求训练数据必须有更多的数量、更多的场景类型,如何构建满足深度学习的数据集成为一个重要的研究方向。目前,基于虚拟场景生成深度数据具有不需要昂贵的深度采集设备、场景类型多样、节省人力成本等优势,结合真实场景和虚拟场景的数据共同训练也是未来深度学习方法的趋势。其次,为了提高深度学习估计单幅图像深度的精度,要求更新的更复杂的深度框架。除了神经网络模型本身结构的优化,更新颖的算法设计也能有效地提升预测精度。研究工作大多采用有监督回归模型对连续的绝对深度值进行回归拟合。考虑到场景由远及近的特性,也有用分类模型进行绝对深度估计的方法。由深度信息和其他信息之间的互补性,部分工作结合表面法线等信息提升深度预测的精度。深度学习发展迅速,新的模型层出不穷,如何将这些模型应用于单幅图像深度估计问题中需要更加深入地研究。另外,探索神经网络在单目深度估计问题中学到的是何种特征也是一个重要的研究方向。

对于单目深度估计模型,目前主要分为基于回归/分类的监督模型,基于双目训练/视频序列的无监督模型,以及基于生成学习的图像风格迁移模型。大概从2017年起,即CVPR2018开始,单目深度估计的效果就已经达到了双目深度估计的效果,主要是监督模型。但是由于现有的数据集主要为KITTI、Cityscapes、NYU DepthV2等,其场景和相机都是固定的,从而导致监督学习下的模型无法适用于其他场景,尤其是多目标跟踪这类细节丰富的场景,可以从论文中看到,基本上每个数据集都会有一个单独的预训练模型。

对于GAN,其对于图像风格的迁移本身是一个很好的泛化点,既可以用于将场景变为晴天、雾天等情况,也可以用于图像分割场景。但是深度估计问题中,像素点存在相对大小,因此必定涉及到回归,因此其必定是监督学习模型,所以泛化性能也不好。对于无监督的算法,可能场景适应性会更好,但依旧不适用于对行人深度的估计。

参考文献

[1] Long J, Shelhamer E, Darrell T. Fully convolutional networks for semantic segmentation[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2015: 3431-3440.

[2] Ronneberger O, Fischer P, Brox T. U-net: Convolutional networks for biomedical image segmentation[C]//International Conference on Medical image computing and computer-assisted intervention. Springer, Cham, 2015: 234-241.

[3] Laina I, Rupprecht C, Belagiannis V, et al. Deeper depth prediction with fully convolutional residual networks[C]//2016 Fourth international conference on 3D vision (3DV). IEEE, 2016: 239-248.

[4] Fu H, Gong M, Wang C, et al. Deep ordinal regression network for monocular depth estimation[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018: 2002-2011.

[5] Godard C, Mac Aodha O, Brostow G J. Unsupervised monocular depth estimation with left-right consistency[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2017: 270-279.

[6] Dosovitskiy A, Fischer P, Ilg E, et al. Flownet: Learning optical flow with convolutional networks[C]//Proceedings of the IEEE international conference on computer vision. 2015: 2758-2766.

[7] Ilg E, Mayer N, Saikia T, et al. Flownet 2.0: Evolution of optical flow estimation with deep networks[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 2462-2470.

[8] Mayer N, Ilg E, Hausser P, et al. A large dataset to train convolutional networks for disparity, optical flow, and scene flow estimation[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2016: 4040-4048.

[9] Xie J, Girshick R, Farhadi A. Deep3d: Fully automatic 2d-to-3d video conversion with deep convolutional neural networks[C]//European Conference on Computer Vision. Springer, Cham, 2016: 842-857.

[10] Luo Y, Ren J, Lin M, et al. Single View Stereo Matching[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018.

[11] Zhou T, Brown M, Snavely N, et al. Unsupervised learning of depth and ego-motion from video[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2017: 1851-1858.

[12] Yin Z, Shi J. Geonet: Unsupervised learning of dense depth, optical flow and camera pose[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018: 1983-1992.

[13] Zhan H, Garg R, Saroj Weerasekera C, et al. Unsupervised learning of monocular depth estimation and visual odometry with deep feature reconstruction[C]// Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018: 340-349.

[14] Goodfellow I, Pouget-Abadie J, Mirza M, et al. Generative adversarial nets[C]//Advances in neural information processing systems. 2014: 2672-2680.

[15] Radford A , Metz L , Chintala S . Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks[J]. Computer Science, 2015.

[16] Arjovsky M, Chintala S, Bottou L. Wasserstein gan[J]. arXiv preprint arXiv:1701.07875, 2017.

[17] Gulrajani I, Ahmed F, Arjovsky M, et al. Improved training of wasserstein gans[C]//Advances in Neural Information Processing Systems. 2017: 5767-5777.

[18] Mao X, Li Q, Xie H, et al. Least squares generative adversarial networks[C]//Proceedings of the IEEE International Conference on Computer Vision. 2017: 2794-2802.

[19] Mirza M, Osindero S. Conditional generative adversarial nets[J]. arXiv preprint arXiv:1411.1784, 2014.

[20] Isola P, Zhu J Y, Zhou T, et al. Image-to-image translation with conditional adversarial networks[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 1125-1134.

[21] Wang T C, Liu M Y, Zhu J Y, et al. High-resolution image synthesis and semantic manipulation with conditional gans[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018: 8798-8807.

[22] Zhu J Y, Park T, Isola P, et al. Unpaired image-to-image translation using cycle-consistent adversarial networks[C]//Proceedings of the IEEE international conference on computer vision. 2017: 2223-2232.

[23] Wang T C , Liu M Y , Zhu J Y , et al. Video-to-Video Synthesis[J]. arXiv preprint arXiv:1808.06601,2018.

[24] Zheng C, Cham T J, Cai J. T2net: Synthetic-to-realistic translation for solving single-image depth estimation tasks[C]//Proceedings of the European Conference on Computer Vision (ECCV). 2018: 767-783.

[25] Atapour-Abarghouei A, Breckon T P. Real-time monocular depth estimation using synthetic data with domain adaptation via image style transfer[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018: 2800-2810.

[26] Nekrasov V , Dharmasiri T , Spek A , et al. Real-Time Joint Semantic Segmentation and Depth Estimation Using Asymmetric Annotations[J]. arXiv preprint arXiv:1809.04766,2018.

[27] Nekrasov V , Shen C , Reid I . Light-Weight RefineNet for Real-Time Semantic Segmentation[J]. arXiv preprint arXiv:1810.03272, 2018.

[28] Lin G , Milan A , Shen C , et al. RefineNet: Multi-Path Refinement Networks for High-Resolution Semantic Segmentation[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition.,2017:1925-1934

[29] Zou Y , Luo Z , Huang J B . DF-Net: Unsupervised Joint Learning of Depth and Flow using Cross-Task Consistency[C]//Proceedings of the European Conference on Computer Vision (ECCV). 2018:36-53.

[30] Ranjan A, Jampani V, Balles L, et al. Competitive collaboration: Joint unsupervised learning of depth, camera motion, optical flow and motion segmentation[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2019: 12240-12249.