https://community.arm.com/arm-community-blogs/b/ai-and-ml-blog/posts/pytorch-to-tensorflow-lite-for-deploying-on-arm-ethos-u55-and-u65
Richard Burton 2021年9月28日
概述
目前，如果您希望在Arm Ethos-U55（https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u55?_ga=2.132077903.1534040332.1634568224-418006553.1606625914）或Arm Ethos-U65（https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u65?_ga=2.132077903.1534040332.1634568224-418006553.1606625914）上部署您的机器学习模型，您需要为微控制器使用TensorFlow Lite（https://www.tensorflow.org/li...）。这样做需要首先将模型设置为TensorFlow Lite格式。如果您已经使用TensorFlow训练了您的模型，那么转换为TensorFlow Lite的过程已经有很好的文档记录。
但是，您可能已经使用PyTorch来训练模型，而PyTorch本身不允许您导出为TensorFlow Lite格式。本指南将展示如何使用ONNX将经过培训的PyTorch模型转换为TensorFlow Lite，然后对其进行量化，以便在Arm Ethos-U55或Arm Ethos-U65上部署。
注意：在继续之前，请注意，尽管本指南介绍了如何将PyTorch培训的模型转换为TensorFlow Lite，但如果可能，我们建议您使用TensorFlow进行模型培训。这避免了从PyTorch转换时可能出现的任何转换错误，并且您还应该看到更好的性能和层支持。
在本指南中，我们将首先在CIFAR10数据集上的PyTorch中训练一个小型卷积神经网络，然后在将其导出为ONNX格式之前对图形执行优化。接下来，ONNX格式的模型将转换为TensorFlow保存的模型格式，最后加载到TensorFlow中进行量化并转换为TensorFlow Lite。
图1。从Pytork获取您的模型并为Arm Ethos-U55或Arm Ethos-U65做好准备的操作流程。
代码
您可以从以下位置下载完整的代码示例（可从命令行运行）：https://github.com/ARM-software/ML-examples/tree/master/pytorch-to-tflite.
您应该确保您的计算机上安装了Python3。
运行以下命令将创建Python环境并安装运行代码示例所需的库。
$ python3 -m venv env
$ source env/bin/activate
$ pip3 install –-upgrade pip3
$ pip3 install torch==1.8.1+cpu torchvision==0.9.1+cpu torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
$ pip3 install onnx==1.9.0
$ pip3 install tensorflow==2.4.1
$ git clone https://github.com/onnx/onnx-tensorflow.git
$ cd onnx-tensorflow
$ pip3 install .
$ pip3 install ethos-u-vela==3.0.0从PyTorch转换到TensorFlow Lite
1.使用以下代码在PyTorch中定义一个简单的卷积神经网络，该网络将在CIFAR10数据集上进行训练：
# A small convolutional network to test PyTorch to TFLite conversion.
class SimpleNetwork(nn.Module):
def __init__(self):
super(SimpleNetwork, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=(3, 3))
self.bn1 = nn.BatchNorm2d(16)
self.conv2 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=(3, 3))
self.bn2 = nn.BatchNorm2d(16)
self.conv3 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3))
self.bn3 = nn.BatchNorm2d(32)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv4 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3, 3), padding=(1, 1)) # Test padding conversion.
self.bn4 = nn.BatchNorm2d(64)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
# An affine operation: y = Wx + b
self.conv5 = nn.Conv2d(in_channels=64,out_channels=10, kernel_size=(6, 6)) # Feature size is 6*6 here.
self.softmax = nn.Softmax()
def forward(self, x):
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = self.pool1(x)
x = F.relu(self.bn3(self.conv3(x)))
x = F.relu(self.bn4(self.conv4(x)))
x = self.pool2(x)
x = self.conv5(x)
x = x.view(-1, 10)
x = self.softmax(x)
return x在本例中，模型在CIFAR10数据集上进行训练，因此加载该数据集，然后使用基本训练循环来训练模型。如果您想了解这是如何实现的，请参考完整的代码示例。
2.模型训练完成后，将其置于评估模式，以便可以使用以下代码进行导出：
model.eval() 我们还在这一点上测试模型的准确性，以便在转换后进行比较。仅经过两次训练，我们就获得了约65%的准确率——比随机机会好得多，正如我们在以下输出中所看到的：
$ Training finished...
$ Accuracy of PyTorch model on test set: 65.97%3.当我们在模型中使用批处理规范化层时，我们可以做的一个优化就是将这些层折叠或融合到前面的卷积操作中。可通过调用模型中必须熔合在一起的层名称列表上的torch.quantization.fuse_模块来折叠熔合，如以下代码所示：
torch.quantization.fuse_modules(model, [[‘conv1’, ‘bn1’], [‘conv2’, ‘bn2’], [‘conv3’, ‘bn3’], [‘conv4’, ‘bn4’]], inplace=True) 批处理规范化层现在融合到它们前面的卷积层中。这样做有助于减少网络执行的不必要计算量，减少权重数量，并有助于使转换更平滑。
4.接下来，我们将经过培训的PyTorch模型导出为ONNX格式。ONNX是一种侧重于推理的模型交换格式，作为不同神经网络框架之间的中间格式。PyTorch本机支持导出为ONNX格式。
我们可以通过调用以下函数将模型导出到ONNX：
torch.onnx.export(model, test_input, “model.onnx”, input_names=[‘input’], output_names=[‘output’])为了导出，我们需要提供一些示例输入数据，以便PyTorch能够准确地识别需要导出的图形部分。export函数还允许我们根据需要更改模型中输入和输出节点的名称。
5.导出为ONNX格式后，我们可以使用Netron检查模型文件，并查看它与原始PyTorch模型几乎相同。请注意，PyTorch模型中的标识操作是熔合批处理规范化层的遗留操作。在ONNX模型中，输入和输出节点现在具有我们在导出时指定的新名称。
图2：左侧为PyTorch模型，右侧为其导出的ONNX版本
一旦模型采用ONNX格式，我们就可以利用ONNX和可用的ONNX转换器，将模型加载并转换为TensorFlow格式。
6.加载ONNX模型，准备将其转换为TensorFlow，然后使用以下代码以TensorFlow保存的模型格式保存到文件中：
onnx_model = onnx.load(“model.onnx”)
tf_rep = onnx_tf.backend.prepare(onnx_model, device=’cpu’)
tf_rep.export_graph(“model_tf”)我们选择将设备设置为“cpu”，以强制操作采用TensorFlow Lite所需的NHWC格式。
7.使用TFLite转换器将我们的模型加载到TensorFlow中，现在该模型是TensorFlow保存模型格式，使用以下代码：
converter = tf.lite.TFLiteConverter.from_saved_model(“model_tf”)注：Arm Ethos-U55和Arm Ethos-U65旨在加速神经网络推理，仅支持8位或16位激活的8位权重。因此，要利用Arm Ethos-U55或Arm Ethos-U65，必须将模型从32位浮点量化为8位定点格式。
8.在使用TFLiteConverter加载转换后的模型后，我们执行训练后量化。该过程的结果将是一个权重和激活都完全量化的模型，以便可以部署在Arm Ethos-U55或Arm Ethos-U65上。以下代码显示了如何执行此量化：
def rep_dataset():
"""Generator function to produce representative dataset for post-training quantization."""
# Use a few samples from the training set.
for _ in range(100):
img = iter(train_loader).next()[0].numpy()
img = [(img.astype(np.float32))]
yield img
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = rep_dataset
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
tflite_model = converter.convert()9使用以下代码将量化模型保存到文件：
open(‘model.tflite’, ‘wb’).write(tflite_model) 10再次测试模型的准确性，因为量化后模型可能会略有变化。我们可以看到，在这种情况下，量化我们的模型对改变我们模型的准确性几乎没有什么作用。该测试还帮助我们确认转换过程也没有影响模型的质量。以下显示了测试的示例输出：
$ Testing quantized TFLite model...
$ Accuracy of quantized TFLite model on test set: 66.02%NHWC与NCHW
如果我们用Netron检查生成的TensorFlow Lite文件，我们可以看到在图的开头添加了转置操作符。PyTorch使用通道优先（NCHW）数据布局进行操作，而TensorFlow和TensorFlow Lite主要使用通道最后（NHWC）。转换器保持与原始PyTorch模型相同的输入形状，以便相同的输入数据可以在转换后的模型中重复使用，而无需更改。因此，转换器添加了转置或重塑操作，以便转换后的卷积操作能够正确处理该数据。
图3。当输入通道大于1时添加转置操作，否则添加重塑操作。
根据模型输出的形状，转换器还可以在模型末尾添加转置操作。此操作将输出形状返回到原始PyTorch模型输出形状。
此外，根据您的模型，您可能会注意到转换器在TensorFlow Lite模型的某些层之前和之后添加了转置操作。这些添加也是转换过程的结果，因为某些操作员即使在转换后仍希望NCHW输入。
例如，当使用某些激活函数（如ReLU6）时，可能会发生这些添加。生成的TensorFlow Lite图如下所示：
图4。ReLU6操作在转换为TensorFlow Lite后被转置运算符包装。
TensorFlow Lite图的一个显著添加是在卷积层之前添加填充操作。TensorFlow有“有效”和“相同”填充的概念，而PyTorch只允许解释要使用的填充。当前，转换为TensorFlow时，转换过程会保持这种显式填充，即使卷积层中的“相同”填充可能是等效的。
图5。转换为TensorFlow Lite后，在卷积层之前添加显式填充操作。
幸运的是，正如我们稍后看到的，如果我们决定在Arm Ethos-U55或Arm Ethos-U65上部署模型并通过Vela运行，这些添加的pad操作可以再次融合。
注意：根据转换为TensorFlow Lite时使用的TensorFlow版本，您可以看到转换器还围绕这些Pad操作添加了转置操作。在这种情况下，无法进行优化。
进一步限制
由于转换过程的当前限制，转换后的TFLite模型在推理速度方面可能不如在TensorFlow中本机训练然后转换为TensorFlow Lite的模型最优。造成这种情况的主要原因是转换器在ONNX和TensorFlow之间转换时可以添加到图形中的额外转置操作。
通过将PyTorch转换的模型与TensorFlow中本机培训的相同模型体系结构进行基准测试，我们可以更具体地看到这一点。目前实现这一点的一种方法是使用一块装载了Arm Ethos-U55和Arm Cortex-M55的MPS3 FPGA板，然后使用Arm ML嵌入式评估工具包的通用推理运行程序应用程序分析推理速度。
我们可以从下表中看到，从PyTorch转换的模型的CPU周期计数比本机TensorFlow模型高一个数量级，而它们都具有相似的NPU周期计数。添加的转置运算符将解释CPU周期计数的增加。这个额外的CPU操作显示在转换模型的额外挂钟推断时间中。
由于本指南中使用的模型非常小，因此用于转置操作的时间占总挂钟时间的相当大比例。对于更大和更深的模型，我们希望看到花费在这个操作上的时间比例大大减少。
另一个限制是，ONNX到TensorFlow转换器并不专注于使您的模型与TensorFlow Lite完全兼容。因此，您可能会遇到模型可以转换为TensorFlow但无法转换为TensorFlow Lite的情况。
克服限制
不幸的是，对于转换过程中的许多限制，似乎没有简单的解决方案，并且尝试将PyTorch模型转换为TensorFlow Lite是不可避免的事实。克服这些问题的最佳建议是在TensorFlow中对您的模型进行本机培训，并从那里转换为TensorFlow Lite。
Vela优化
现在我们有了量化的TensorFlow Lite文件，我们可以通过Arm Vela编译器运行它来优化模型，以便在Arm Ethos-U55或Arm Ethos-U65上部署。
通过运行以下命令，可以从命令行完成此优化：
$ vela model.tflite使用了与优化有关的默认参数，以及Arm Ethos-U55的目标参数，但在Vela对其进行优化后，只需查看结果模型就可以了。看见https://pypi.org/project/etho... 有关可与Vela一起使用的不同命令行选项的详细信息。
图6。通过Vela编译器运行后转换的TensorFlow Lite模型。
在创建的输出文件夹中，我们可以找到TensorFlow Lite格式的Vela优化模型。检查优化后的模型，我们可以看到，除了转置层之外，每个节点都被编译成一个ethos-u操作符。这意味着所有这些层都可以在Arm Ethos-U55或Arm Ethos-U65 NPU上成功运行和加速。转置层在伴随的Arm Cortex-M CPU上运行。
结论
通过遵循本指南，您应该已经成功地将PyTorch模型训练并转换为TensorFlow格式。然后，该模型在TensorFlow中量化并转换为TensorFlow Lite格式，最后通过Vela运行，准备部署在Arm Ethos-U55或Arm Ethos-U65上。
我们还讨论了此转换过程的一些限制，主要是引入的额外转置操作，以及在Arm Ethos-U55或Arm Ethos-U65上运行时这会如何影响最终性能。
最终，如果您有选择，我们强烈建议您使用TensorFlow框架以本机方式训练您的模型，并从那里转换为TensorFlow Lite。这将在您部署它时产生最佳的模型，并在过程中为您带来最少的问题。然而，如果这是不可能的，那么正如我们在本指南中所示，从PyTorch转换可能会产生一个可用的模型。
您可以在这里找到Jupyter笔记本，亲自尝试一下：
