在为 SOFIE 的 RModel 实现 Keras 解析器之后,我开始开发 PyTorch 解析器。 PyTorch 支持原生导出到.onnx格式,因此具有跟踪 PyTorch 模型中使用的 ONNX 运算符的内部机制。 PyTorch 是一个快速发展的深度学习框架,通过 Python 绑定的 PyRoot 接口在 ROOT TMVA 中得到支持。为了保存和导入已保存的模型,ROOT 使用 TorchScript。

在我的 GSoC'21 与 CERN-HSF 系列的这篇博客中,我讨论了 PyTorch 解析器的实现,用于将.pt文件转换为 SOFIE 的 RModel 以用于后续生成推理代码。

目录

1.PyTorch & TorchScript

2、解析PyTorch.pt文件

3.测试

1. PyTorch 和 TorchScript

正如 PyTorch 的官方存储库将其描述为 Tensors 和 Python 中具有强大 GPU 加速的动态神经网络的框架,PyTorch 提供了一个强大的基于 Python 的平台和一个基于 Torch 的 C++ 后端,用于实现神经网络和进行基于张量的高级计算。开源框架由 Facebook 开发,后来被公开。 PyTorch 采用命令式方法,维护一个动态计算图,用于定义计算顺序和计算神经网络梯度下降优化的梯度。计算图是在声明变量时动态构建的。因此,在每次连续的训练迭代之后,再次构建该图。它的动态特性使其更加灵活和易于调试。在 TensorFlow2 Alpha 出现之前,TensorFlow 维护了一个静态计算图。这些范式的主要区别可以通过各种类比来理解。遵循_“定义并运行”_哲学的 TensorFlow 的工作方式与任何编译器一样,其中所有变量首先在图中定义、连接、初始化,然后被重用,但内部结构无法在运行中再次重建。 PyTorch 遵循“定义并运行”的概念,因此 DCG 就像解释器一样工作,在流程的每个步骤中都构建了图形,并且可以在运行时修改其内部结构。

TorchScript

与 Keras & TensorFlow 的model.save(filename)保存完整的深度学习模型及其配置和权重不同,使用 PyTorch 的torch.save()保存模型将序列化数据绑定到特定的类和目录结构,因此需要在加载之前定义模型类模型。

TorchScript 提供了优化 PyTorch 模型并从 Python 进程序列化它们的功能,该进程稍后可以加载到不一定依赖 Python 的环境中,例如 C++ 独立程序。将 TorchScript 添加到 PyTorch 实现了一个统一的框架,用于将模型从研究过渡到生产。将 PyTorch 模块序列化为 TorchScript 可以让我们保存整个模型以及配置和权重,以便以后用于推理。 TorchScript 提供了两种执行方法,即Trace和Script。跟踪模型需要一个输入张量,该输入张量和计算在张量上运行,直到记录输出,因此可以稍后用于类似形状的输入张量。为模型编写脚本更加通用,其中整个模型的所有配置都是序列化的,以后可以与任何适用形状的张量一起使用。

Root TMVA 的 PyTorch 接口允许在 TMVA 环境中训练 PyTorch 模型,并使用 TorchScript 对它们进行序列化,稍后可用于推理。

2.解析 PyTorch.pt文件

由于 PyTorch 在构建神经网络模型中提供了在图中动态添加模型计算的灵活性,因此很难提取所包含层的信息,从而解析配置并构建 RModel 对象。 PyTorch 模型可以构建在nn.Sequential、nn.Module和nn.ModuleList容器中,并且还支持将层封装到单元中以进行子类化。由于 ROOT TMVA 的 PyTorch 接口使用 TorchScript 保存和加载模型的脚本模块,该模块只包含用于推理的模型例程和权重,因此与 Keras Sequential API 和函数式不同,提取构成 PyTorch 模型的底层特征很困难API模型,图可以迭代遍历。

为了解决这个问题并找到提取包含层及其属性的方法,我们可以借助 PyTorch 提供的 Native ONNX 支持。

https://github.com/pytorch/pytorch/tree/master/torch/onnx

PyTorch 具有将模型转换为其 ONNX 等效图的内部实现。为此,需要一个通过模型传递的虚拟输入张量,并记录其过程中应用的运算符以开发 ONNX 图。在为此使用保存的脚本模块时,我们还需要形状推断的预期示例输出。为了处理具有动态轴的张量,我们再次使用 PyTorch 内部功能为它们设置形状推断。

#Converting the model to ONNX Graph representation
import torch
from torch.onnx.utils import _model_to_graph
from torch.onnx.symbolic_helper import _set_onnx_shape_inference

model= torch.jit.load('model.pt')
model.eval()
model.float() 

#Building the dummy input tensor
s1=[120,1]                       
dummy=torch.rand(*s1) 
input=[dummy]
output=model(*input)

#Setting inference for dynamic axes
_set_onnx_shape_inference(True)      

#Generating the equivalent ONNX graph
graph=_model_to_graph(model,input,example_outputs=output)

图的输出看起来像这样

>>> graph
(graph(%x.1 : Float(2, 6, strides=[6, 1], requires_grad=0, device=cpu),
      %fc1.weight : Float(12, 6, strides=[6, 1], requires_grad=0, device=cpu),
      %fc1.bias : Float(12, strides=[1], requires_grad=0, device=cpu)):
  %3 : Float(2, 12, strides=[12, 1], device=cpu) = onnx::Gemm[alpha=1., beta=1., transB=1](%x.1, %fc1.weight, %fc1.bias) 
  %4 : Float(2, 12, strides=[12, 1], device=cpu) = onnx::Relu(%3)
  %5 : Float(12, 2, strides=[1, 12], requires_grad=1, device=cpu) = onnx::Transpose[perm=[1, 0]](%4) 
  return (%5)
, {'fc1.bias': tensor([ 0.1858, -0.4581, -0.3743, -0.0609, -0.1052,  0.0538,  0.0467, -0.1544,
        -0.3622, -0.3611, -0.2269,  0.1716]), 'fc1.weight': tensor([[ 1.9536e-01,  3.3694e-02, -9.5294e-02, -3.2092e-01, -2.6642e-01,
          3.2872e-01],
        [ 2.5622e-01, -2.5419e-01,  3.2304e-01, -5.7926e-02,  1.8545e-01,
         -3.3965e-01],
        [ 2.3855e-01, -3.0029e-01, -3.2645e-01, -3.6467e-01, -1.1971e-01,
         -1.5546e-01],
        [-3.9777e-01, -2.4680e-01,  4.6634e-02,  7.5872e-05,  5.0814e-02,
         -3.0577e-01],
        [ 9.8371e-03, -3.5033e-01, -1.8954e-01,  2.4746e-01, -1.9828e-01,
         -3.3955e-01],
        [-2.3662e-02, -3.1937e-02, -8.0010e-02,  1.6615e-01,  2.1353e-01,
         -3.5624e-01],
        [-1.0537e-01, -1.0636e-01,  3.1500e-01, -2.0943e-01,  6.0947e-01,
          2.3454e-01],
        [-1.7886e-01, -1.6449e-01, -7.1962e-02, -1.4553e-01,  2.7605e-01,
          2.8320e-02],
        [-4.0435e-01,  3.2550e-01, -2.1856e-01, -3.2801e-02, -3.6782e-01,
         -2.4415e-01],
        [ 1.9953e-01,  2.1276e-01, -6.8311e-02,  7.6783e-02,  1.1836e-01,
          1.2552e-01],
        [ 1.5363e-01, -2.6658e-01, -4.5881e-01, -2.5509e-01,  1.9992e-01,
          6.6146e-02],
        [-3.3991e-01, -1.3903e-01, -1.1859e-01, -8.4289e-02,  1.4896e-02,
          1.6103e-01]])}, None)

这里的图表是一个元组,其中graph[0]保存模型配置,graph[1]保存模型权重。由于图是ONNX图格式,所以我们可以直接遍历它,提取节点数据。

modelData = []

for i in graph[0].nodes():
 globals().update(locals())

 #For type of node
 nodeData = []
 nodeData.append(i.kind())

 #For attrbutes of node
 nodeAttributeNames = [x for x in i.attributeNames()]
 nodeAttributes     = {j:i[j] for j in nodeAttributeNames}
 nodeData.append(nodeAttributes)

 #For input tensor info of node
 nodeInputs     = [x for x in i.inputs()]
 nodeInputNames = [x.debugName() for x in nodeInputs]
 nodeData.append(nodeInputNames)

 #For output tensor info of node
 nodeOutputs     = [x for x in i.outputs()]
 nodeOutputNames = [x.debugName() for x in nodeOutputs]
 nodeData.append(nodeOutputNames)

 modelData.append(nodeData)

然后以与 Keras 解析器类似的方式使用这些提取的数据,将运算符、输入和输出张量信息以及初始化的张量添加到 RModel 对象中,以实现成功的翻译。 Root 是一个主要用 C++ 开发的框架,因此 C-Python API 用于所有这些 Python 变量的提取和操作。官方的C-Python API可用于理解用法,使用IBM 的 DeveloperWorks提供的更简单指南可在此处找到。

可以使用PR#8684跟踪 SOFIE 的 PyTorch 解析器的开发。

https://github.com/root-project/root/pull/8684

3.测试

为验证已实现的解析器而开发了测试,这与 Keras 解析器的测试类似。开发 Python 脚本,从 C-Python API 运行以生成并保存 PyTorch 模型,然后将该模型解析为 RModel 对象,随后生成推理代码。然后通过比较 PyTorch 模型的输出和在相同输入张量上调用时生成的推理代码,在 Google 的 GTest 框架上找到正确性。

这标志着为 SOFIE 实现 PyTorch 解析器的完成,至此,我的 Google Summer of Code 项目的主要任务也完成了。随着官方编码阶段的剩余时间,我也开始为 BDT 模型开发 Root-Storage,用于它们的序列化和推理代码生成。在 GSoC 之后,我将描述 BDT 根存储的开发,届时将实现大部分功能,我的以下博客将在 GSoC 的最终报告中总结和总结这 10 周的惊人旅程.

暂时告别,

San级板

图片来源

infoq.com/presentations/pytorch-torchscript..

miro.medium.com/max/700/1*5PLIVNA5fIqEC8-kZ..