PyTorch深度学习实战54——使用TorchScript和ONNX导出通用PyTorch模型0. 前言1. TorchScript 简介2. 使用 TorchScript 进行模型追踪3. 使用 TorchScript 进行模型脚本化4. 在 C 中运行 PyTorch 模型5. 使用 ONNX 导出 PyTorch 模型小结系列链接0. 前言我们已经深入探讨了 PyTorch 模型部署这可能是将PyTorch模型投入生产系统中最关键的一环。在本节中我们将聚焦另一个重要维度模型导出。我们已经学习了如何在经典的 Python 脚本环境中保存和加载 PyTorch 模型。但是我们还需要更多的方式来导出PyTorch模型主要是出于以下考虑首先Python解释器通过全局解释器锁 (Global Interpreter Lock,GIL) 限制单线程运行 (目前Python正在逐步移除GIL)这阻碍了操作并行化其次目标运行环境(如某些系统或设备)可能不支持Python为解决这些问题PyTorch提供了高效且与平台/语言无关的模型导出方案使模型能脱离训练环境运行。我们首先探讨TorchScript它能将序列化优化后的PyTorch模型转换为中间表示 (Intermediate Representation,IR)从而在非Python环境(如C程序)中运行。接着探讨ONNX标准该技术支持以通用格式保存模型以便导入其他深度学习框架或跨编程语言使用。1. TorchScript 简介TorchScript是将PyTorch模型投入生产环境中的关键工具原因有两个PyTorch默认采用即时执行模式 (eager execution)。虽然这种模式便于调试但逐操作执行时频繁的内存读写会导致推理延迟增高且难以进行全局优化。为此PyTorch提供了基于Python子集的即时 (just-in-time,JIT) 编译器解决方案。JIT编译器通过整体分析模型所有操作将其编译为单一复合计算图(而非逐行解释执行)生成的TorchScript代码实质上是静态类型的Python子集。这种编译方式带来多重性能提升消除全局解释器锁 (Global Interpreter Lock,GIL) 实现多线程支持同时可应用多种图结构优化生产环境往往需要更高性能(如更快速、内存效率更高)的C等语言支持或需在不支持Python的硬件设备上部署。此时TorchScript展现出独特优势——当PyTorch模型被编译为中间表示 (Intermediate Representation,IR) 后可通过TorchScript编译器序列化为C兼容格式最终借助PyTorch C API(LibTorch) 在C推理程序中加载运行本节已多次提及PyTorch模型的JIT编译机制现在我们将具体探讨两种将PyTorch模型编译为TorchScript格式的方法。2. 使用 TorchScript 进行模型追踪一种方法是通过模型追踪 (tracing) 实现PyTorch代码到TorchScript的转换。该方法需要提供PyTorch模型对象和一个虚拟输入样例 (dummy input)。顾名思义追踪机制会记录这个虚拟输入在模型(神经网络)中的流动过程捕捉所有运算操作最终生成可视图化的TorchScript中间表示 (Intermediate Representation,IR)——既支持图形化展示也可转换为TorchScript代码。接下来同样以手写数字分类模型为例演示追踪PyTorch模型的过程。(1)首先导入所需库importtorchimporttorch.nnasnnimporttorch.nn.functionalasFimporttorch.optimasoptimfromtorch.utils.dataimportDataLoaderfromtorchvisionimportdatasets,transformsimportnumpyasnpfromPILimportImage(2)接下来定义并实例化模型对象classConvNet(nn.Module):def__init__(self):super(ConvNet,self).__init__()self.cn1nn.Conv2d(1,16,3,1)self.cn2nn.Conv2d(16,32,3,1)self.dp1nn.Dropout(0.10)self.dp2nn.Dropout(0.25)self.fc1nn.Linear(4608,64)# 4608 is basically 12 X 12 X 32self.fc2nn.Linear(64,10)defforward(self,x):xself.cn1(x)xF.relu(x)xself.cn2(x)xF.relu(x)xF.max_pool2d(x,2)xself.dp1(x)xtorch.flatten(x,1)xself.fc1(x)xF.relu(x)xself.dp2(x)xself.fc2(x)opF.log_softmax(x,dim1)returnop modelConvNet()(3)然后加载模型权重PATH_TO_MODEL./convnet.pthmodel.load_state_dict(torch.load(PATH_TO_MODEL,weights_onlyFalse,map_locationcpu))(4)然后加载一张示例图像imageImage.open(./digit_image.jpg)(5)接下来定义数据预处理函数并将预处理函数应用于示例图像defimage_to_tensor(image):gray_imagetransforms.functional.to_grayscale(image)resized_imagetransforms.functional.resize(gray_image,(28,28))input_image_tensortransforms.functional.to_tensor(resized_image)input_image_tensor_normtransforms.functional.normalize(input_image_tensor,(0.1302,),(0.3069,))returninput_image_tensor_norm input_tensorimage_to_tensor(image)(6)同时还需要执行以下代码否则所有被追踪模型的参数都需要梯度计算因此必须在torch.no_grad()上下文中加载模型model.eval()forpinmodel.parameters():p.requires_grad_(False)(7)加载具有预训练权重的PyTorch模型对象后使用虚拟输入来追踪模型demo_inputtorch.ones(1,1,28,28)traced_modeltorch.jit.trace(model,demo_input)本节所用的虚拟输入是一个所有像素值都设置为1的图像。(8)查看追踪生成的模型计算图traced_model.graph输出结果如下所示直观来看计算图起始部分显示模型各层初始化过程(如fc2、dp2等层)末端则呈现最后的softmax层。可以看到该计算图采用静态类型变量的低级语言描述其语法形式与TorchScript语言非常相似。(9)除了查看计算图我们还可以通过以下命令获取追踪模型对应的TorchScript代码print(traced_model.code)执行后将输出如下类Python代码这些代码定义了模型的forward传播方法这段代码正是我们在步骤 2中使用PyTorch编写的等效TorchScript实现。(10)接下来导出(保存)追踪后的模型torch.jit.save(traced_model,traced_convnet.pt)(11)加载保存的模型loaded_traced_modeltorch.jit.load(traced_convnet.pt)需要注意的是这里我们不需要分别加载模型的架构和参数。(12)最后使用该模型进行推理loaded_traced_model(input_tensor.unsqueeze(0))输出结果如下tensor([[-9.9644e00, -1.1098e01, -2.3822e-03, -8.3610e00, -8.2592e00, -1.1988e01, -8.2979e00, -1.0842e01, -6.4784e00, -1.1229e01]])(13)可以通过在原模型上重新运行推理来验证结果model(input_tensor.unsqueeze(0))输出结果与第 12 步相同验证了我们的追踪模型运行正常。得益于TorchScript无需GIL的特性可以使用追踪模型替代原始PyTorch模型对象来构建更高效的 Flask 模型服务器。虽然追踪是JIT编译PyTorch模型的有效方法但它也存在一些局限性。例如若模型的前向传播包含条件分支(如if语句)或循环结构(如for语句)追踪机制仅能捕获其中一条执行路径的运算流程。为确保此类控制流模型的准确转换我们需要采用另一种编译机制——脚本化 (scripting)。3. 使用 TorchScript 进行模型脚本化复用上一小节代码然后继续完成本节内容。(1)脚本化转换无需提供虚拟输入直接通过以下代码将PyTorch模型转为TorchScriptscripted_modeltorch.jit.script(model)(2)查看脚本化后的模型图scripted_model.graph输出格式与追踪模型类似如下图所示可见计算图仍以逐行描述的底层脚本形式呈现但注意此处计算图结构与上一小节的追踪结果存在差异这表明使用追踪与脚本化在代码编译策略上有所不同。(3)通过运行以下代码查看等效的TorchScript代码print(scripted_model.code)输出结果如下所示本质上流程与上一小节类似然而由于编译策略的不同代码签名上存在细微差异。(4)同样地脚本化模型可通过以下方式导出并重新加载torch.jit.save(scripted_model,scripted_convnet.pt)loaded_scripted_modeltorch.jit.load(scripted_convnet.pt)(5)最后使用脚本化模型进行推理loaded_scripted_model(input_tensor.unsqueeze(0))输出结果应与上一小节步骤 12完全一致由此验证脚本化模型运行符合预期。与追踪模型类似脚本化模型同样不受GIL限制因此在Flask服务部署时可显著提升性能。下表展示了两种编译策略的对比特征TracingScripting需要虚拟输入不需要虚拟输入通过将虚拟输入传递给模型记录固定的数学操作序列通过检查 PyTorch 代码中的 nn.Module 内容生成 TorchScript 代码/图无法处理模型前向传播中的多个控制流(如 if-else)对处理各种控制流(如 if-else、循环等)非常有用即使模型包含不被 TorchScript 支持的 PyTorch 功能也能正常工作只有在 PyTorch 模型不包含 TorchScript 不支持的功能时才有效我们已经展示了如何将PyTorch模型转换并序列化为TorchScript模型。在下一节中我们将暂时脱离Python展示如何使用C加载TorchScript序列化模型。4. 在 C 中运行 PyTorch 模型在某些场景下Python可能成为性能瓶颈或者我们无法在目标环境中运行基于PyTorch和Python训练的机器学习模型。为此本节将利用导出的TorchScript序列化模型(包括追踪和脚本化两种方式)演示如何在C代码中执行模型推理。在开始之前我们需要安装CMake以支持C代码编译。完成安装后在当前工作目录下创建名为cpp_convnet的文件夹后续操作都将在此目录中进行。(1)编写用于运行模型推理流程的C文件#includetorch/script.h#includeopencv2/core.hpp#includeopencv2/imgcodecs.hpp#includeopencv2/highgui.hpp#includeopencv2/imgproc.hpp#includeiostreamusingnamespacecv;usingnamespacestd;intmain(intargc,char**argv){Mat imgimread(argv[2],IMREAD_GRAYSCALE);首先使用OpenCV库将.jpg图像文件读取为灰度图像。(2)然后将灰度图像调整为28x28像素即本节模型所需输入规格resize(img,img,Size(28,28));(3)接着将图像数组转换为PyTorch张量autoinput_torch::from_blob(img.data,{img.rows,img.cols,img.channels()},at::kByte);本节所有与torch相关的操作都需使用libtorch库——这是PyTorch C API的核心组件。若已安装PyTorch则无需单独安装LibTorch。(4)由于OpenCV读取的灰度图像维度为(28, 28, 1)我们需要将其转换为(1, 28, 28)格式以满足PyTorch要求。接着将张量重塑为(1,1,28,28)形状其中第一个维度表示推理时的batch_size第二个维度为通道数对于灰度图像而言为1autoinputinput_.permute({2,0,1}).unsqueeze_(0).reshape({1,1,img.rows,img.cols}).toType(c10::kFloat).div(255);input(input-0.1302)/0.3069;由于OpenCV读取的图像像素值范围为0-255我们首先将其归一化到[0,1]区间(通过除以255实现)。接着按照预处理标准使用均值0.1302和标准差0.3069对图像进行标准化处理。(5)加载已经导出的JIT编译后的TorchScript模型对象automoduletorch::jit::load(argv[1]);std::vectortorch::jit::IValueinputs;inputs.push_back(input);(6)最后进入模型预测阶段使用加载的模型对象利用提供的输入数据执行前向推理autooutput_module.forward(inputs).toTensor();(7)输出变量output_存储着每个类别的预测概率分布。提取概率最高的类别标签并打印autooutputoutput_.argmax(1);coutoutput\n;(8)最终退出C例程return0;}(9)此外还需要在同一工作目录中编写一个CMakeLists.txt文件cmake_minimum_required(VERSION3.12FATAL_ERROR)set(CMAKE_CXX_STANDARD17)set(CMAKE_CXX_STANDARD_REQUIRED ON)set(CMAKE_CXX_EXTENSIONS OFF)add_compile_options(-stdc17)project(cpp_convnet)find_package(Torch REQUIRED)find_package(OpenCV REQUIRED)set(CMAKE_CXX_FLAGS${CMAKE_CXX_FLAGS}${TORCH_CXX_FLAGS})add_executable(cpp_convnet cpp_convnet.cpp)target_link_libraries(cpp_convnet pthread)target_link_libraries(cpp_convnet${TORCH_LIBRARIES}${OpenCV_LIBS})set_property(TARGET cpp_convnet PROPERTY CXX_STANDARD14)该文件本质上是库安装和构建脚本类似于Python项目中的setup.py文件。此外还需要设置OpenCV_DIR环境变量$exportOpenCV_DIR/usr/lib/x86_64-linux-gnu/cmake/opencv4/(10)接下来需要实际运行CMakeLists文件完成构建。我们在当前工作目录中创建一个新目录并从该目录中运行构建过程来。在命令行中需执行以下命令$mkdirbuild $cdbuild $ cmake-DCMAKE_PREFIX_PATH/home/brainiac/Documents/myenv_torch/lib/python3.12/site-packages/torch..$ cmake--build.--configRelease第三行需要提供LibTorch的安装路径。要获取系统中的路径可在Python中执行以下代码importtorch torch.__path__输出类似以下结果# /lib/python3.12/site-packages/torch执行第三行命令将输出类似以下结果执行第四行命令将输出以下内容成功执行以上步骤后将生成名为cpp_convnet的C编译二进制文件。运行该可执行程序即通过C模型对样本图像进行推理。可以选择以下两种方式输入模型使用脚本化模型作为输入$ ./cpp_convnet../../scripted_convnet.pt../../digit_image.jpg或者使用追踪模型作为输入$ ./cpp_convnet../../traced_convnet.pt../../digit_image.jpg两种方式均应输出如下结果2[CPULongType{1}]可以看到C模型正常运行。需注意的是由于C使用OpenCV而Python使用PIL进行图像处理像素编码方式存在细微差异这将导致预测概率略有不同。但只要正确实施归一化处理最终预测结果并不会出现显著差异。至此我们完成了PyTorch模型在C环境中的推理实践。在本节中我们学习了如何将训练好的PyTorch深度学习模型移植到C环境以提高预测效率同时为在没有Python环境的系统(例如某些嵌入式系统、无人机等)中部署模型开辟了可能性。接下来我们将讨论一个通用的神经网络建模格式——ONNX(Open Neural Network Exchange)它能实现跨深度学习框架、编程语言和操作系统的模型互通。具体而言我们将介绍如何将PyTorch训练的模型导入TensorFlow进行推理。5. 使用 ONNX 导出 PyTorch 模型在生产系统中有些已部署的机器学习模型大多是使用某个深度学习库(例如TensorFlow)编写的并配有成熟的模型服务基础设施。如果某个模型是使用PyTorch编写的我们希望能够使用TensorFlow来运行它可以通过ONNX等标准化框架可使其兼容TensorFlow服务策略。ONNX(Open Neural Network Exchange) 是一种通用格式它将深度学习模型的基本操作(如矩阵乘法和激活函数)进行标准化这些操作在不同的深度学习库中有不同的实现方式。ONNX使我们能够灵活地使用不同的深度学习库、编程语言甚至不同的运行环境来执行同一个深度学习模型。在本节我们将介绍如何在TensorFlow中运行一个使用PyTorch训练的模型。我们将首先将PyTorch模型导出为ONNX格式然后将该ONNX模型加载到TensorFlow代码中。除了TensorFlow外还需要安装onnx和onnx2tf这两个库。首先复用《使用 TorchScript 进行模型追踪》小节中的步骤1到11然后继续执行以下步骤。(1)首先安装所需库$ pipinstallonnx onnx2tf tf_keras onnx_graphsurgeon ai_edge_litert sng4onnx(2)与模型追踪类似我们将一个虚拟输入传入已加载的模型demo_inputtorch.ones(1,1,28,28)torch.onnx.export(model,demo_input,convnet.onnx)这将保存一个ONNX格式的模型文件。其底层使用的模型序列化机制与模型追踪中的机制相同。(3)接下来加载保存的ONNX模型并将其转换为TensorFlow模型onnx2tf.convert(input_onnx_file_pathconvnet.onnx,output_folder_pathconvnet_tf,non_verboseTrue,)(4)然后加载序列化的TensorFlow模型以解析模型的计算图。以验证我们是否正确加载了模型结构并识别计算图的输入和输出节点modeltf.saved_model.load(./convnet_tf/)model输出结果如下所示tensorflow.python.saved_model.load.Loader._recreate_base_user_object.locals._UserObject at0x7e0fc83c6ae0(5)最后我们在TensorFlow模型上运行推理为示例图像生成预测结果# Perform inferenceoutputmodel(input_tensor.unsqueeze(-1))# Print the outputprint(Model Output:,output)输出结果如下所示可以看到TensorFlow和PyTorch版本的模型预测结果完全相同。这验证了ONNX框架的成功运行。我们可以进一步分析TensorFlow模型理解ONNX如何通过利用模型计算图中的底层数学运算在不同的深度学习库中重新生成完全相同的模型。小结在本节中我们将深入探讨使用TorchScript导出PyTorch模型。通过序列化TorchScript使模型与Python生态系统独立从而使得模型可以在其他环境中加载例如基于C的环境。我们还跨越Torch框架与Python生态的边界研究机器学习通用开放格式ONNX该技术能帮助我们将PyTorch训练的模型导出至非PyTorch甚至非Python环境。系列链接PyTorch深度学习实战1——神经网络与模型训练过程详解PyTorch深度学习实战2——PyTorch基础PyTorch深度学习实战3——使用PyTorch构建神经网络PyTorch深度学习实战4——常用激活函数和损失函数详解PyTorch深度学习实战5——计算机视觉基础PyTorch深度学习实战6——神经网络性能优化技术PyTorch深度学习实战7——批大小对神经网络训练的影响PyTorch深度学习实战8——批归一化PyTorch深度学习实战9——学习率优化PyTorch深度学习实战10——过拟合及其解决方法PyTorch深度学习实战11——卷积神经网络PyTorch深度学习实战12——数据增强PyTorch深度学习实战13——可视化神经网络中间层输出PyTorch深度学习实战14——类激活图PyTorch深度学习实战15——迁移学习PyTorch深度学习实战16——面部关键点检测PyTorch深度学习实战17——多任务学习PyTorch深度学习实战18——目标检测基础PyTorch深度学习实战19——从零开始实现R-CNN目标检测PyTorch深度学习实战20——从零开始实现Fast R-CNN目标检测PyTorch深度学习实战21——从零开始实现Faster R-CNN目标检测PyTorch深度学习实战22——从零开始实现YOLO目标检测PyTorch深度学习实战23——从零开始实现SSD目标检测PyTorch深度学习实战24——使用U-Net架构进行图像分割PyTorch深度学习实战25——从零开始实现Mask R-CNN实例分割PyTorch深度学习实战26——多对象实例分割PyTorch深度学习实战27——自编码器(Autoencoder)PyTorch深度学习实战28——卷积自编码器(Convolutional Autoencoder)PyTorch深度学习实战29——变分自编码器(Variational Autoencoder, VAE)PyTorch深度学习实战30——对抗攻击(Adversarial Attack)PyTorch深度学习实战31——神经风格迁移PyTorch深度学习实战32——DeepfakesPyTorch深度学习实战33——生成对抗网络(Generative Adversarial Network, GAN)PyTorch深度学习实战34——DCGAN详解与实现PyTorch深度学习实战35——条件生成对抗网络(Conditional Generative Adversarial Network, CGAN)PyTorch深度学习实战36——Pix2Pix详解与实现PyTorch深度学习实战37——CycleGAN详解与实现PyTorch深度学习实战38——StyleGAN详解与实现PyTorch深度学习实战39——少样本学习(Few-shot Learning)PyTorch深度学习实战40——零样本学习(Zero-Shot Learning)PyTorch深度学习实战41——循环神经网络与长短期记忆网络PyTorch深度学习实战42——图像字幕生成PyTorch深度学习实战43——手写文本识别PyTorch深度学习实战44——基于 DETR 实现目标检测PyTorch深度学习实战45——强化学习PyTorch深度学习实战46——深度Q学习PyTorch深度学习实战47——使用PyTorch构建Transformer模型PyTorch深度学习实战48——基于Transformer实现机器翻译PyTorch深度学习实战49——扩散模型Diffusion Model详解与实现PyTorch深度学习实战50——PyTorch分布式训练PyTorch深度学习实战51——自动混合精度训练PyTorch深度学习实战52——PyTorch深度学习模型部署PyTorch深度学习实战53——使用TorchServe部署PyTorch模型
PyTorch深度学习实战（54）——使用TorchScript和ONNX导出通用PyTorch模型
发布时间：2026/5/30 6:31:55
PyTorch深度学习实战54——使用TorchScript和ONNX导出通用PyTorch模型0. 前言1. TorchScript 简介2. 使用 TorchScript 进行模型追踪3. 使用 TorchScript 进行模型脚本化4. 在 C 中运行 PyTorch 模型5. 使用 ONNX 导出 PyTorch 模型小结系列链接0. 前言我们已经深入探讨了 PyTorch 模型部署这可能是将PyTorch模型投入生产系统中最关键的一环。在本节中我们将聚焦另一个重要维度模型导出。我们已经学习了如何在经典的 Python 脚本环境中保存和加载 PyTorch 模型。但是我们还需要更多的方式来导出PyTorch模型主要是出于以下考虑首先Python解释器通过全局解释器锁 (Global Interpreter Lock,GIL) 限制单线程运行 (目前Python正在逐步移除GIL)这阻碍了操作并行化其次目标运行环境(如某些系统或设备)可能不支持Python为解决这些问题PyTorch提供了高效且与平台/语言无关的模型导出方案使模型能脱离训练环境运行。我们首先探讨TorchScript它能将序列化优化后的PyTorch模型转换为中间表示 (Intermediate Representation,IR)从而在非Python环境(如C程序)中运行。接着探讨ONNX标准该技术支持以通用格式保存模型以便导入其他深度学习框架或跨编程语言使用。1. TorchScript 简介TorchScript是将PyTorch模型投入生产环境中的关键工具原因有两个PyTorch默认采用即时执行模式 (eager execution)。虽然这种模式便于调试但逐操作执行时频繁的内存读写会导致推理延迟增高且难以进行全局优化。为此PyTorch提供了基于Python子集的即时 (just-in-time,JIT) 编译器解决方案。JIT编译器通过整体分析模型所有操作将其编译为单一复合计算图(而非逐行解释执行)生成的TorchScript代码实质上是静态类型的Python子集。这种编译方式带来多重性能提升消除全局解释器锁 (Global Interpreter Lock,GIL) 实现多线程支持同时可应用多种图结构优化生产环境往往需要更高性能(如更快速、内存效率更高)的C等语言支持或需在不支持Python的硬件设备上部署。此时TorchScript展现出独特优势——当PyTorch模型被编译为中间表示 (Intermediate Representation,IR) 后可通过TorchScript编译器序列化为C兼容格式最终借助PyTorch C API(LibTorch) 在C推理程序中加载运行本节已多次提及PyTorch模型的JIT编译机制现在我们将具体探讨两种将PyTorch模型编译为TorchScript格式的方法。2. 使用 TorchScript 进行模型追踪一种方法是通过模型追踪 (tracing) 实现PyTorch代码到TorchScript的转换。该方法需要提供PyTorch模型对象和一个虚拟输入样例 (dummy input)。顾名思义追踪机制会记录这个虚拟输入在模型(神经网络)中的流动过程捕捉所有运算操作最终生成可视图化的TorchScript中间表示 (Intermediate Representation,IR)——既支持图形化展示也可转换为TorchScript代码。接下来同样以手写数字分类模型为例演示追踪PyTorch模型的过程。(1)首先导入所需库importtorchimporttorch.nnasnnimporttorch.nn.functionalasFimporttorch.optimasoptimfromtorch.utils.dataimportDataLoaderfromtorchvisionimportdatasets,transformsimportnumpyasnpfromPILimportImage(2)接下来定义并实例化模型对象classConvNet(nn.Module):def__init__(self):super(ConvNet,self).__init__()self.cn1nn.Conv2d(1,16,3,1)self.cn2nn.Conv2d(16,32,3,1)self.dp1nn.Dropout(0.10)self.dp2nn.Dropout(0.25)self.fc1nn.Linear(4608,64)# 4608 is basically 12 X 12 X 32self.fc2nn.Linear(64,10)defforward(self,x):xself.cn1(x)xF.relu(x)xself.cn2(x)xF.relu(x)xF.max_pool2d(x,2)xself.dp1(x)xtorch.flatten(x,1)xself.fc1(x)xF.relu(x)xself.dp2(x)xself.fc2(x)opF.log_softmax(x,dim1)returnop modelConvNet()(3)然后加载模型权重PATH_TO_MODEL./convnet.pthmodel.load_state_dict(torch.load(PATH_TO_MODEL,weights_onlyFalse,map_locationcpu))(4)然后加载一张示例图像imageImage.open(./digit_image.jpg)(5)接下来定义数据预处理函数并将预处理函数应用于示例图像defimage_to_tensor(image):gray_imagetransforms.functional.to_grayscale(image)resized_imagetransforms.functional.resize(gray_image,(28,28))input_image_tensortransforms.functional.to_tensor(resized_image)input_image_tensor_normtransforms.functional.normalize(input_image_tensor,(0.1302,),(0.3069,))returninput_image_tensor_norm input_tensorimage_to_tensor(image)(6)同时还需要执行以下代码否则所有被追踪模型的参数都需要梯度计算因此必须在torch.no_grad()上下文中加载模型model.eval()forpinmodel.parameters():p.requires_grad_(False)(7)加载具有预训练权重的PyTorch模型对象后使用虚拟输入来追踪模型demo_inputtorch.ones(1,1,28,28)traced_modeltorch.jit.trace(model,demo_input)本节所用的虚拟输入是一个所有像素值都设置为1的图像。(8)查看追踪生成的模型计算图traced_model.graph输出结果如下所示直观来看计算图起始部分显示模型各层初始化过程(如fc2、dp2等层)末端则呈现最后的softmax层。可以看到该计算图采用静态类型变量的低级语言描述其语法形式与TorchScript语言非常相似。(9)除了查看计算图我们还可以通过以下命令获取追踪模型对应的TorchScript代码print(traced_model.code)执行后将输出如下类Python代码这些代码定义了模型的forward传播方法这段代码正是我们在步骤 2中使用PyTorch编写的等效TorchScript实现。(10)接下来导出(保存)追踪后的模型torch.jit.save(traced_model,traced_convnet.pt)(11)加载保存的模型loaded_traced_modeltorch.jit.load(traced_convnet.pt)需要注意的是这里我们不需要分别加载模型的架构和参数。(12)最后使用该模型进行推理loaded_traced_model(input_tensor.unsqueeze(0))输出结果如下tensor([[-9.9644e00, -1.1098e01, -2.3822e-03, -8.3610e00, -8.2592e00, -1.1988e01, -8.2979e00, -1.0842e01, -6.4784e00, -1.1229e01]])(13)可以通过在原模型上重新运行推理来验证结果model(input_tensor.unsqueeze(0))输出结果与第 12 步相同验证了我们的追踪模型运行正常。得益于TorchScript无需GIL的特性可以使用追踪模型替代原始PyTorch模型对象来构建更高效的 Flask 模型服务器。虽然追踪是JIT编译PyTorch模型的有效方法但它也存在一些局限性。例如若模型的前向传播包含条件分支(如if语句)或循环结构(如for语句)追踪机制仅能捕获其中一条执行路径的运算流程。为确保此类控制流模型的准确转换我们需要采用另一种编译机制——脚本化 (scripting)。3. 使用 TorchScript 进行模型脚本化复用上一小节代码然后继续完成本节内容。(1)脚本化转换无需提供虚拟输入直接通过以下代码将PyTorch模型转为TorchScriptscripted_modeltorch.jit.script(model)(2)查看脚本化后的模型图scripted_model.graph输出格式与追踪模型类似如下图所示可见计算图仍以逐行描述的底层脚本形式呈现但注意此处计算图结构与上一小节的追踪结果存在差异这表明使用追踪与脚本化在代码编译策略上有所不同。(3)通过运行以下代码查看等效的TorchScript代码print(scripted_model.code)输出结果如下所示本质上流程与上一小节类似然而由于编译策略的不同代码签名上存在细微差异。(4)同样地脚本化模型可通过以下方式导出并重新加载torch.jit.save(scripted_model,scripted_convnet.pt)loaded_scripted_modeltorch.jit.load(scripted_convnet.pt)(5)最后使用脚本化模型进行推理loaded_scripted_model(input_tensor.unsqueeze(0))输出结果应与上一小节步骤 12完全一致由此验证脚本化模型运行符合预期。与追踪模型类似脚本化模型同样不受GIL限制因此在Flask服务部署时可显著提升性能。下表展示了两种编译策略的对比特征TracingScripting需要虚拟输入不需要虚拟输入通过将虚拟输入传递给模型记录固定的数学操作序列通过检查 PyTorch 代码中的 nn.Module 内容生成 TorchScript 代码/图无法处理模型前向传播中的多个控制流(如 if-else)对处理各种控制流(如 if-else、循环等)非常有用即使模型包含不被 TorchScript 支持的 PyTorch 功能也能正常工作只有在 PyTorch 模型不包含 TorchScript 不支持的功能时才有效我们已经展示了如何将PyTorch模型转换并序列化为TorchScript模型。在下一节中我们将暂时脱离Python展示如何使用C加载TorchScript序列化模型。4. 在 C 中运行 PyTorch 模型在某些场景下Python可能成为性能瓶颈或者我们无法在目标环境中运行基于PyTorch和Python训练的机器学习模型。为此本节将利用导出的TorchScript序列化模型(包括追踪和脚本化两种方式)演示如何在C代码中执行模型推理。在开始之前我们需要安装CMake以支持C代码编译。完成安装后在当前工作目录下创建名为cpp_convnet的文件夹后续操作都将在此目录中进行。(1)编写用于运行模型推理流程的C文件#includetorch/script.h#includeopencv2/core.hpp#includeopencv2/imgcodecs.hpp#includeopencv2/highgui.hpp#includeopencv2/imgproc.hpp#includeiostreamusingnamespacecv;usingnamespacestd;intmain(intargc,char**argv){Mat imgimread(argv[2],IMREAD_GRAYSCALE);首先使用OpenCV库将.jpg图像文件读取为灰度图像。(2)然后将灰度图像调整为28x28像素即本节模型所需输入规格resize(img,img,Size(28,28));(3)接着将图像数组转换为PyTorch张量autoinput_torch::from_blob(img.data,{img.rows,img.cols,img.channels()},at::kByte);本节所有与torch相关的操作都需使用libtorch库——这是PyTorch C API的核心组件。若已安装PyTorch则无需单独安装LibTorch。(4)由于OpenCV读取的灰度图像维度为(28, 28, 1)我们需要将其转换为(1, 28, 28)格式以满足PyTorch要求。接着将张量重塑为(1,1,28,28)形状其中第一个维度表示推理时的batch_size第二个维度为通道数对于灰度图像而言为1autoinputinput_.permute({2,0,1}).unsqueeze_(0).reshape({1,1,img.rows,img.cols}).toType(c10::kFloat).div(255);input(input-0.1302)/0.3069;由于OpenCV读取的图像像素值范围为0-255我们首先将其归一化到[0,1]区间(通过除以255实现)。接着按照预处理标准使用均值0.1302和标准差0.3069对图像进行标准化处理。(5)加载已经导出的JIT编译后的TorchScript模型对象automoduletorch::jit::load(argv[1]);std::vectortorch::jit::IValueinputs;inputs.push_back(input);(6)最后进入模型预测阶段使用加载的模型对象利用提供的输入数据执行前向推理autooutput_module.forward(inputs).toTensor();(7)输出变量output_存储着每个类别的预测概率分布。提取概率最高的类别标签并打印autooutputoutput_.argmax(1);coutoutput\n;(8)最终退出C例程return0;}(9)此外还需要在同一工作目录中编写一个CMakeLists.txt文件cmake_minimum_required(VERSION3.12FATAL_ERROR)set(CMAKE_CXX_STANDARD17)set(CMAKE_CXX_STANDARD_REQUIRED ON)set(CMAKE_CXX_EXTENSIONS OFF)add_compile_options(-stdc17)project(cpp_convnet)find_package(Torch REQUIRED)find_package(OpenCV REQUIRED)set(CMAKE_CXX_FLAGS${CMAKE_CXX_FLAGS}${TORCH_CXX_FLAGS})add_executable(cpp_convnet cpp_convnet.cpp)target_link_libraries(cpp_convnet pthread)target_link_libraries(cpp_convnet${TORCH_LIBRARIES}${OpenCV_LIBS})set_property(TARGET cpp_convnet PROPERTY CXX_STANDARD14)该文件本质上是库安装和构建脚本类似于Python项目中的setup.py文件。此外还需要设置OpenCV_DIR环境变量$exportOpenCV_DIR/usr/lib/x86_64-linux-gnu/cmake/opencv4/(10)接下来需要实际运行CMakeLists文件完成构建。我们在当前工作目录中创建一个新目录并从该目录中运行构建过程来。在命令行中需执行以下命令$mkdirbuild $cdbuild $ cmake-DCMAKE_PREFIX_PATH/home/brainiac/Documents/myenv_torch/lib/python3.12/site-packages/torch..$ cmake--build.--configRelease第三行需要提供LibTorch的安装路径。要获取系统中的路径可在Python中执行以下代码importtorch torch.__path__输出类似以下结果# /lib/python3.12/site-packages/torch执行第三行命令将输出类似以下结果执行第四行命令将输出以下内容成功执行以上步骤后将生成名为cpp_convnet的C编译二进制文件。运行该可执行程序即通过C模型对样本图像进行推理。可以选择以下两种方式输入模型使用脚本化模型作为输入$ ./cpp_convnet../../scripted_convnet.pt../../digit_image.jpg或者使用追踪模型作为输入$ ./cpp_convnet../../traced_convnet.pt../../digit_image.jpg两种方式均应输出如下结果2[CPULongType{1}]可以看到C模型正常运行。需注意的是由于C使用OpenCV而Python使用PIL进行图像处理像素编码方式存在细微差异这将导致预测概率略有不同。但只要正确实施归一化处理最终预测结果并不会出现显著差异。至此我们完成了PyTorch模型在C环境中的推理实践。在本节中我们学习了如何将训练好的PyTorch深度学习模型移植到C环境以提高预测效率同时为在没有Python环境的系统(例如某些嵌入式系统、无人机等)中部署模型开辟了可能性。接下来我们将讨论一个通用的神经网络建模格式——ONNX(Open Neural Network Exchange)它能实现跨深度学习框架、编程语言和操作系统的模型互通。具体而言我们将介绍如何将PyTorch训练的模型导入TensorFlow进行推理。5. 使用 ONNX 导出 PyTorch 模型在生产系统中有些已部署的机器学习模型大多是使用某个深度学习库(例如TensorFlow)编写的并配有成熟的模型服务基础设施。如果某个模型是使用PyTorch编写的我们希望能够使用TensorFlow来运行它可以通过ONNX等标准化框架可使其兼容TensorFlow服务策略。ONNX(Open Neural Network Exchange) 是一种通用格式它将深度学习模型的基本操作(如矩阵乘法和激活函数)进行标准化这些操作在不同的深度学习库中有不同的实现方式。ONNX使我们能够灵活地使用不同的深度学习库、编程语言甚至不同的运行环境来执行同一个深度学习模型。在本节我们将介绍如何在TensorFlow中运行一个使用PyTorch训练的模型。我们将首先将PyTorch模型导出为ONNX格式然后将该ONNX模型加载到TensorFlow代码中。除了TensorFlow外还需要安装onnx和onnx2tf这两个库。首先复用《使用 TorchScript 进行模型追踪》小节中的步骤1到11然后继续执行以下步骤。(1)首先安装所需库$ pipinstallonnx onnx2tf tf_keras onnx_graphsurgeon ai_edge_litert sng4onnx(2)与模型追踪类似我们将一个虚拟输入传入已加载的模型demo_inputtorch.ones(1,1,28,28)torch.onnx.export(model,demo_input,convnet.onnx)这将保存一个ONNX格式的模型文件。其底层使用的模型序列化机制与模型追踪中的机制相同。(3)接下来加载保存的ONNX模型并将其转换为TensorFlow模型onnx2tf.convert(input_onnx_file_pathconvnet.onnx,output_folder_pathconvnet_tf,non_verboseTrue,)(4)然后加载序列化的TensorFlow模型以解析模型的计算图。以验证我们是否正确加载了模型结构并识别计算图的输入和输出节点modeltf.saved_model.load(./convnet_tf/)model输出结果如下所示tensorflow.python.saved_model.load.Loader._recreate_base_user_object.locals._UserObject at0x7e0fc83c6ae0(5)最后我们在TensorFlow模型上运行推理为示例图像生成预测结果# Perform inferenceoutputmodel(input_tensor.unsqueeze(-1))# Print the outputprint(Model Output:,output)输出结果如下所示可以看到TensorFlow和PyTorch版本的模型预测结果完全相同。这验证了ONNX框架的成功运行。我们可以进一步分析TensorFlow模型理解ONNX如何通过利用模型计算图中的底层数学运算在不同的深度学习库中重新生成完全相同的模型。小结在本节中我们将深入探讨使用TorchScript导出PyTorch模型。通过序列化TorchScript使模型与Python生态系统独立从而使得模型可以在其他环境中加载例如基于C的环境。我们还跨越Torch框架与Python生态的边界研究机器学习通用开放格式ONNX该技术能帮助我们将PyTorch训练的模型导出至非PyTorch甚至非Python环境。系列链接PyTorch深度学习实战1——神经网络与模型训练过程详解PyTorch深度学习实战2——PyTorch基础PyTorch深度学习实战3——使用PyTorch构建神经网络PyTorch深度学习实战4——常用激活函数和损失函数详解PyTorch深度学习实战5——计算机视觉基础PyTorch深度学习实战6——神经网络性能优化技术PyTorch深度学习实战7——批大小对神经网络训练的影响PyTorch深度学习实战8——批归一化PyTorch深度学习实战9——学习率优化PyTorch深度学习实战10——过拟合及其解决方法PyTorch深度学习实战11——卷积神经网络PyTorch深度学习实战12——数据增强PyTorch深度学习实战13——可视化神经网络中间层输出PyTorch深度学习实战14——类激活图PyTorch深度学习实战15——迁移学习PyTorch深度学习实战16——面部关键点检测PyTorch深度学习实战17——多任务学习PyTorch深度学习实战18——目标检测基础PyTorch深度学习实战19——从零开始实现R-CNN目标检测PyTorch深度学习实战20——从零开始实现Fast R-CNN目标检测PyTorch深度学习实战21——从零开始实现Faster R-CNN目标检测PyTorch深度学习实战22——从零开始实现YOLO目标检测PyTorch深度学习实战23——从零开始实现SSD目标检测PyTorch深度学习实战24——使用U-Net架构进行图像分割PyTorch深度学习实战25——从零开始实现Mask R-CNN实例分割PyTorch深度学习实战26——多对象实例分割PyTorch深度学习实战27——自编码器(Autoencoder)PyTorch深度学习实战28——卷积自编码器(Convolutional Autoencoder)PyTorch深度学习实战29——变分自编码器(Variational Autoencoder, VAE)PyTorch深度学习实战30——对抗攻击(Adversarial Attack)PyTorch深度学习实战31——神经风格迁移PyTorch深度学习实战32——DeepfakesPyTorch深度学习实战33——生成对抗网络(Generative Adversarial Network, GAN)PyTorch深度学习实战34——DCGAN详解与实现PyTorch深度学习实战35——条件生成对抗网络(Conditional Generative Adversarial Network, CGAN)PyTorch深度学习实战36——Pix2Pix详解与实现PyTorch深度学习实战37——CycleGAN详解与实现PyTorch深度学习实战38——StyleGAN详解与实现PyTorch深度学习实战39——少样本学习(Few-shot Learning)PyTorch深度学习实战40——零样本学习(Zero-Shot Learning)PyTorch深度学习实战41——循环神经网络与长短期记忆网络PyTorch深度学习实战42——图像字幕生成PyTorch深度学习实战43——手写文本识别PyTorch深度学习实战44——基于 DETR 实现目标检测PyTorch深度学习实战45——强化学习PyTorch深度学习实战46——深度Q学习PyTorch深度学习实战47——使用PyTorch构建Transformer模型PyTorch深度学习实战48——基于Transformer实现机器翻译PyTorch深度学习实战49——扩散模型Diffusion Model详解与实现PyTorch深度学习实战50——PyTorch分布式训练PyTorch深度学习实战51——自动混合精度训练PyTorch深度学习实战52——PyTorch深度学习模型部署PyTorch深度学习实战53——使用TorchServe部署PyTorch模型
