[模型部署]：深度学习模型部署（已更Pytorch篇）

模型部署一、Pytorch 模型部署1. 模型保存1.1. 使用torch.save()保存1.2. 使用[torch.jit.save()](/docs/stable/generated/torch.jit.save.html)保存 2. 模型部署 or 模型编译2.1. TorchScript部署2.2. TensorRT部署2.3. TVM部署2.4. Laser部署

模型部署

一、Pytorch 模型部署

1. 模型保存

1.1. 使用torch.save()保存

保存和加载state_dict，即只保存模型参数

保存：

torch.save(model.state_dict(), SAVE_PATH)

加载：

model = ModelClass(*args, **kwargs)model.load_state_dict(torch.load(SAVE_PATH))model.eval()

保存完整的模型文件

保存：

torch.save(model, SAVE_PATH)

加载：

model = torch.load(SAVE_PATH)model.eval()

保存模型并恢复训练状态，断点重训

# 保存state = {'epoch': epoch,'state_dict': model.state_dict(),'optimizer': optimizer.state_dict(),...}torch.save(state, filepath)# 模型与参数加载model.load_state_dict(state['state_dict'])optimizer.load_state_dict(state['optimizer'])

1.2. 使用torch.jit.save()保存

用于保存编译过后的模型，跨平台使用，模型中只能使用pytorch的函数，可以使用被加载到C ++ APItorch::jit::load(filename)或与Python的APItorch.jit.load

torch.jit.save支持保存script类型和trace类型的模为TorchScript，其中script为全量模型，trace为仅保存运行过的路径。

什么是JIT？TorchScript？

JIT（just-in-time compilation，即时编译）模式选择了构建well-defined IR的方式。这样带来的好处有：

将python runtime和计算图（model graph）解耦获得C++相比带GIL的Python的收益获得整个程序，拿到全局信息，来进行优化将易于调试的模式和易于部署/优化的模式进行切分（算法调参和算法部署各司其职）

TorchScript是JIT模式的具体形式，是一种从PyTorch代码创建可序列化和可优化模型的方法。任何TorchScript程序都可以从Python进程中保存，并加载到没有Python依赖的进程中。

trace和script的区别？

1、trace只记录走过的tensor和对tensor的操作，不会记录任何控制流信息，如if条件句和循环。因为没有记录控制流的另外的路，也没办法对其进行优化。好处是trace深度嵌入python语言，复用了所有python的语法，在计算流中记录数据流。

2、script会去理解所有的code，真正像一个编译器一样去进行lexer、parser、Semantic analusis的分析「也就是词法分析语法分析句法分析，形成AST树，最后再将AST树线性化」。script相当于一个嵌入在Python/Pytorch的DSL，其语法只是pytorch语法的子集，这意味着存在一些op和语法script不支持，这样在编译的时候就会遇到问题。此外，script的编译优化方式更像是CPU上的传统编译优化，重点对于图进行硬件无关优化，并对IF、loop这样的statement进行优化。

【Pytorch部署】TorchScript

PyTorch系列「一」PyTorch JIT —— trace/ script的代码组织和优化方法

2. 模型部署 or 模型编译

什么是模型编译？深度学习编译器？

传统的编译器是以高层语言作为输入，避免直接去写汇编，机器码；而深度学习编译器作用相仿，其输入是高度抽象的计算图，输出包括CPU或者GPU等硬件平台是哪个的底层代码和执行引擎，即深度学习编译器将不同框架描述的深度学习模型为某个硬件平台生成优化的代码。AOT(ahead-of-time compliation)

深度学习编译器普遍采用的设计架构：

这类编译器的通用设计体系结构主要包含两部分：编译器前端和编译器后端。 中间表示（IR）横贯前端和后端。 通常IR是程序的抽象，用于程序优化。 具体而言，深度学习模型在编译器中转换为多级IR，其中高级IR驻留在前端，而低级IR驻留在后端。 基于高级IR，编译器前端负责独立于硬件的转换和优化。 基于低级IR，编译器后端负责特定于硬件的优化、代码生成和编译。

2.1. TorchScript部署

参见上文

2.2. TensorRT部署

tensorrt 部署 pth：

使用pytorch训练得到pt文件；将pt文件转换为onnx中间件；使用onnxsim.simplify对转换后的onnx进行简化；解析onnx文件构建trt推理引擎；cd tensorrt_path/bin && ./trtexec --onnx=*.onnx --saveEngine=*.trt导出网络定义以及相关权重；解析网络定义以及相关权重；根据显卡算子构造出最优执行计划；将执行计划序列化存储；反序列化执行计划；进行推理

pytorch中使用TensorRT

2.3. TVM部署

模型在使用 TVM 编译器框架进行转换时所采取的步骤。

从Tensorflow、PyTorch或Onnx等框架导入模型。

TVM 导入层是从其他框架（如 Tensorflow、PyTorch 或 ONNX）导入模型，如果在将模型导入 TVM 时遇到问题，想尝试将其转换为ONNX（更适配）。

将导入的模型翻译成TVM 的高级模型语言Relay。

已导入 TVM 的模型在 Relay 中表示。Relay 是模型的中间表示 (Intermediate representation, IR)。支持：

传统数据流表示Functional-style scoping, let-binding which makes it a fully featured differentiable language允许混合两种编程风格的能力

Relay 应用图级优化传递来优化模型。

Lower到Tensor expression(TE) 表示。

Lower是指将高级表示转换为低级表示。在应用高级优化后，Relay 运行 FuseOps pass 将模型划分为许多小子图，并将子图降低为 TE 表示。张量表达式 (TE) 是一种用于描述张量计算的domain-specific语言。TE 还提供了几个schedule来指定low-level loop优化，例如tiling、vectorization、parallelization、unrolling和fusion。为了帮助将Relay representation转换为 TE represention的过程，TVM 包括一个张量运算符清单 (Tensor Operator Inventory, TOPI)，它具有预定义的常用张量运算符模板（例如，conv2d、转置）。

使用auto-tuning模块AutoTVM或AutoScheduler搜索最佳schedule。

schedule为 TE 中定义的运算符或子图指定的低级循环进行优化。auto-tuning模块搜索最佳schedule并将其与cost model和on-device measurement进行比较。TVM 中有两个auto-tuning模块。

AutoTVM：基于模板的自动调整模块。它运行搜索算法来找到最佳值。对于常用的算子，TOPI 中已经提供了模板。AutoScheduler（又名 Ansor）：一个无模板的自动调整模块。它不需要预定义的schedule模板。相反，它通过分析计算定义自动生成搜索空间。然后它在生成的搜索空间中搜索最佳schedule。

选择模型编译的最佳配置。

调优后，自动调优模块会生成JSON格式的调优记录。此步骤为每个子图选择最佳schedule。

Lower to Tensor Intermediate Representation (TIR)，TVM 的低级中间表示。

在根据调整步骤选择最佳配置后，每个 TE 子图都降低到 TIR 并通过低级优化通道进行优化。接下来，将优化后的 TIR 降低到硬件平台的目标编译器。这是生成可部署到生产中的优化模型的最终代码生成阶段。TVM 支持多种不同的编译器后端，包括：

LLVM，它可以针对任意微处理器架构，包括标准 x86 和 ARM 处理器、AMDGPU 和 NVPTX 代码生成，以及 LLVM 支持的任何其他平台。专用编译器，例如 NVCC，NVIDIA 的编译器。嵌入式和专用目标，通过 TVM 的Bring Your Own Codegen (BYOC) 框架实现。

编译成机器代码。

在此过程结束时，特定于编译器的生成代码可以lower为机器代码。

TVM 可以将模型编译为可链接的对象模块，然后使用轻量级 TVM runtime 提供的 C API 来动态加载模型然后运行，以及使用其他语言（如 Python 和 Rust）的接口。TVM 还可以构建捆绑部署，其中runtime与模型结合在一个包中。

使用tvmc的python API

TVM 实战练习：

[模型部署]：TVM模型部署实战

2.4. Laser部署