今天我们先看一下通用的模型部署场景,然后讨论一下Libtorch模型转换。


深度学习模型部署

深度学习模型训练完成后,需要将模型部署应用,该部门一般无需再考虑如何修改训练方式或者修改网络结构以提高模型精度,更多的是需要明确部署的场景、部署方式(中心服务化还是本地终端部署)、模型的优化指标,以及如何提高吞吐率和减少延迟等。

根据部署方式的不同,通常分为梁总:云端部署和设备端部署。云端部署常见的模式是,模型部署在云端服务器,用户通过网页访问或者API接口调用等形式向云端服务器发出请求,云端收到请求后处理并返回结果。设备端部署则主要用于嵌入式设备,主要通过将模型打包封装到SDK,集成到嵌入式设备,数据的处理和模型推理都在终端设备上执行。

在实际的部署中,除去Demo性质的场景可以用Python进行简单的展示,其余场景考虑到性能问题,通常都会用C++来实现推理部门,对于某些极地资源的场景,还会考虑C语言,本文后续会在X86和ARM平台下展示一个典型的部署,所以会用C++来实现。


Libtorch模型转换

为了将Pytorch模型用C++进行部署,通常会采用第三方推理框架,目前最常见的主要有官方提供的Libtorch库和微软推出的ONNXRuntime,不管采用哪一个方式,都需要先将Python训练好的模型转换到对应支持的格式,然后在C++里调用相应的模型做推理得到最终结果。

本文先介绍Libtorch,它是由Pytorch官方提供推理库,可以方便的把转换的模型在C++层面进行调用。需要将模型转换为script model才能由Libtorch加载并进行推理。目前官方提供了两种方式进行转换:


1. Trace

Trace的方法模拟模型推理的过程,给模型一组可能的输入,进行一遍网络推理,然后由torch.ji.trace记录一下路径上的信息并保存即可。示例如下:


使用该方法时,如果模型中存在控制流比如if-else语句,一组输入只能遍历一个分支,这种情况下就没办法完整的把模型信息记录下来。


2. Script

另一种方法时直接在训练脚本中通过torch.jit.script编译模块,将其转换为ScriptModule。示例如下:


• forward方法会被默认编译,forward中被调用的方法也会按照被调用的顺序被编译

• 如果想要编译一个forward以外且未被forward调用的方法,可以添加 @torch.jit.export.

• 如果想要方法不被编译,可使用@torch.jit.ignore 或者 @torch.jit.unused

在使用Script进行模型转换时,会遇到一些限制,可以参考文末的资料,主要是try-catch等分支不被支持,我们在SpeechBrain的ECAPA-TDNN模型中也遇到了这个问题,解决办法就是重新修改模型结构,尽量的避免或融合这个分支。


SpeechBrain模型转换

回到我们的声纹模型转换,我们需要将模型保存成script的形式供C++调用,具体的代码已经实现,在这里:convert_to_torchscript

整个代码也非常简单:


在上面的代码中,首先根据输入的模型类型判断是'TDNN'还是'ECAPA'模型,然后调用SpeechBrain的构造方法构造一个初始化的模型,得到相应的模型后,调用model.load_state_dict(torch.load(args.raw_ckpt))去加载已经训练好的模型参数,至此已经可以做推理或转换了,我们在这里用的是上面提到的Script方法,直接转换并保存scripted模型,此外,还实现了一个量化版本的quantized_model,后续介绍模型优化的时候会用到,这里先写上了。

可以用如下命令执行该转换:


执行完后就得到了最终的模型embedding.ckpt.pt ,这就是我们C++推理需要用到的模型文件。


文章来源知乎,文章作者:蘑菇炖提莫