Unrolling Arm Ethos-U55和U65的RNN模型

https://community.arm.com/arm-community-blogs/b/ai-and-ml-blog/posts/rnn-models-ethos-u

Richard Burton 2021年11月5日

概述

递归神经网络（RNN）是一种非常适合处理序列数据或可表示为序列的数据的神经网络。它们的工作方式是按时间顺序处理多个输入，并具有一个内部状态，允许它们对以前看到的输入信息进行编码。重要的是，RNN的所有权重在每个步骤之间共享。递归神经网络的一些流行用例包括：机器翻译、语音合成和时间序列预测。其中许多用例适用于小型嵌入式平台，但也可以通过在专业硬件（如神经处理单元（NPU））上运行而受益。

2262.png-459x598.png

图1:t时间步的简单递归神经网络布局。

在本指南中，我们将向您展示如何使用TensorFlow来训练和量化与Arm嵌入式NPU（如Arm Ethos-U55和Arm Ethos-U65）兼容的简单递归神经网络。

在你开始之前需准备

可从以下位置下载可从命令行运行的完整代码示例：https://github.com/ARM-software/ML-examples/tree/master/rnn-unrolling-tflite

您应该确保您的计算机上安装了Python3。

运行以下命令可创建Python环境并安装运行代码示例所需的库。

$ python3 -m venv env 
$ source env/bin/activate 
$ pip3 install –-upgrade pip3 
$ pip3 install tensorflow == 2.5.0
$ pip3 install numpy == 1.19.5
$ pip3 install ethos-u-vela == 3.0.0

训练

首先，我们将训练一种RNN类型，简称为选通循环单元（GRU）。这是一种稍微复杂一点的RNN类型，它的特性使它能够比标准RNN更长久地保留过去的重要信息。

我们使用GRU的任务是使用MNIST数据集对手写数字进行分类。通常情况下，RNN不用于图像分类任务，但是为了简单起见，我们选择在这里这样做。为了使用GRU对图像进行分类，我们将28x28个输入图像中的每一行视为GRU的一次性输入。通过这种方式，我们有28个图像行的序列要输入到GRU。在执行了28个时间步长的计算后，GRU已经“看到”了整个图像，我们可以将最终输出传递到一个完全连接的层，该层将为我们进行分类。

0246.png-438x700.png

图2：我们训练的GRU模型的Layout。

1.在TensorFlow中定义一个模型，该模型由一个GRU层和一个完全连接的层组成。以下代码显示了如何执行此操作：


def rnn_model(time_steps):
    model_input = tf.keras.Input(shape=(time_steps, 28), name='input')

    gru_out = tf.keras.layers.GRU(units=256)(model_input)

    model_output = tf.keras.layers.Dense(10, activation='softmax')(gru_out)

    model = tf.keras.Model([model_input], [model_output])

    return model

2.接下来，加载用于训练和测试RNN模型的数据，并使用以下代码对其进行规范化：


(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
# Normalize between [0, 1] to help training.
x_train = x_train / 255.0  
x_test = x_test / 255.0

3.使用以下代码创建先前定义的模型的实例：

model = rnn_model(time_steps=28)

时间步长设置为28，即每个图像中的行数。

4.编译模型，然后使用标准的“fit”方法调用对其进行训练。为了节省时间，我们只训练一个时代，训练更长的时间当然会得到更精确的模型。如以下代码

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])
model.fit(x=x_train, y=y_train, epochs=1, validation_data=(x_test, y_test))

5.经过一段时期的训练，我们在验证数据集上的准确率达到了几乎98%，考虑到这不是RNN的理想任务，这是一个不错的结果。训练完成后，我们将模型的权重保存到文件中，以便稍后再次使用：

model.save_weights('gru_mnist')

部署到TFLite

6.现在，我们已经训练了模型权重，我们将定义一个稍微不同的graph ，用于部署。此graph 对于推断和转换为TFLite格式更为优化，最终允许模型量化并在Arm Ethos-U55上运行。或Arm Ethos-U65以下代码示例显示了TensorFlow中的新graph定义：


def rnn_model_tflite(time_steps):
    # We need to specify batch_size=1 to get an optimal graph for inference.
    input_node = tf.keras.Input(shape=(time_steps, 28), batch_size=1, name='input')

    gru_out = tf.keras.layers.GRU(units=256, unroll=True)(input_node)

    prediction = tf.keras.layers.Dense(10, activation='softmax')(gru_out)

    model = tf.keras.Model([input_node], [prediction])

    return model

它看起来与前面的定义几乎相同，但有两个重要的变化。第一个是将输入节点上的batch_size参数设置为1。由于我们只计划使用经过训练的模型一次进行单一推断，因此我们可以安全地进行这项工作。这样做还有一个好处，那就是删除可能必须自动添加的任何额外操作，以读取批量大小或处理大于1的批量大小。

第二个变化是定义GRU层时使用的参数。这是最重要的更改，它涉及将“unroll ”设置为True。

全部展开（Unroll）

通常，RNN没有定义它可以循环的时间步数的结束。因此，它可以处理任何长度的输入。然而，对于许多现实生活中的用例，我们提前知道需要RNN循环的时间步数。在我们的示例中，我们知道需要有28个时间步来说明每个输入图像中的28行。将unroll设置为True将从RNN中删除循环，如图3所示，并将其Unroll到层的输入设置的时间步数“t”。结果将是一个严格的前馈神经网络，然后可以量化并轻松转换为TensorFlow Lite。

0412.png-1000x371.png

图3：RNN的Unroll。

7.我们创建了一个新的“部署就绪”模型的实例，并使用以下代码将以前保存的权重加载到此模型中：

model_for_tflite = rnn_model_tflite(time_steps=28)
model_for_tflite.load_weights('gru_mnist')

量化（Quantize）与TFLite

现在，我们已经创建了最佳模型图，并且加载了来自先前训练的权重，我们希望继续并量化它。在小型边缘设备上部署模型时，量化是必不可少的一步，因为它可以减少模型大小并加快推理时间。此外，像Arm Ethos-U55这样的NPU只支持对量化到8位的模型进行推理。

以下量化模型的步骤来自培训后量化（https://www.tensorflow.org/lite/performance/post_training_quantization#full_integer_quantization）的公共TensorFlow教程#tensorflow#tensorflowlite#量子化

8.根据训练后量化的需要，我们定义了一个生成函数，该函数将从训练数据集中生成固件。以下代码显示了如何执行此操作：

def rep_dataset():
    for i in range(50):  # Only need a few examples.
        img = x_train[i].astype(np.float32)
        img = np.expand_dims(img, 0)
        yield [img]

然后我们将新的keras模型加载到TFLiteConverter中，设置训练后量化的所有属性，并将模型转换为TensorFlow Lite格式。最后，我们将模型保存到文件中，以便使用Vela进行优化，并部署到Arm Ethos-U55或Arm Ethos-U65上。以下代码显示了这些步骤：


converter_quant = tf.lite.TFLiteConverter.from_keras_model(model_for_tflite)
converter_quant.optimizations = [tf.lite.Optimize.DEFAULT]
converter_quant.representative_dataset = rep_dataset

converter_quant.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter_quant.inference_input_type = tf.int8
converter_quant.inference_output_type = tf.int8
tflite_model = converter_quant.convert()
open('gru_mnist.tflite', 'wb').write(tflite_model)

测试量化精度

以往当量化你的神经网络模型时，可能会经历一些精度损失。因此，您应该在量化后检查模型的准确性。在我们的GRU模型中这样做，我们看到精度几乎保持不变，这是一个好消息。

$ Testing quantized TFLite model...
$ Accuracy of quantized TFLite model on test set: 97.57%

输入和获取隐藏状态

对于您的特定模型，您可能希望手动设置初始隐藏状态或从RNN层检索最终隐藏状态。幸运的是，这非常简单，只需要修改导出到TensorFlow Lite时定义模型的方式。例如，要获得一个能够访问这些输入和输出状态的TensorFlow Lite模型，我们只需定义我们的模型，如以下代码片段所示：


def rnn_model_tflite_with_state_transfer(time_steps):
    # We need to specify batch_size=1 to get an optimal graph for inference.
    input_node = tf.keras.Input(shape=(time_steps, 28), batch_size=1, name='input')

    # Input tensor for the initial GRU hidden state.
    initial_gru_state = tf.keras.Input(shape=256, batch_size=1, name='initial_gru_state')

    # Return the final GRU hidden state.
    gru_out, final_gru_state = tf.keras.layers.GRU(units=256, unroll=True, return_state=True)(input_node, initial_state=initial_gru_state)

    prediction = tf.keras.layers.Dense(10, activation='softmax')(gru_out)

    model = tf.keras.Model([input_node, initial_gru_state], [prediction, final_gru_state])

    return model

在这个模型定义中，我们现在定义了一个新的 Input tensor，即GRU隐藏层的大小，这就是我们将如何输入初始GRU状态。在GRU层中，我们将return_state设置为True，以便捕获最终的隐藏状态。当我们调用层时，我们提供了我们在前一行中定义的初始状态。最后，在创建keras模型时，我们将这个额外的输入和输出提供给相应的输入和输出节点列表。

如果您使用的是LSTM层而不是基本的RNN或GRU层，那么代码只需要从上面进行一点小的修改。您需要为单元状态添加一个额外的 Input tensor，还需要捕获额外的最终单元状态。在制作TensorFlow keras模型时，还需要将其添加到输入和输出列表中。

注意：在启用此状态传输的模型上进行训练后量化时，您需要在代表性数据生成器函数中提供这些输入隐藏状态的示例数据。如果您的用例涉及将最终状态作为新的初始状态提供给您的RNN，那么建议您收集这些最终状态的示例，并迭代地将它们添加到具有代表性的输入状态数据集中。

Vela

准备在Arm Ethos-U55或Arm Ethos-U65上部署的模型的最后一步是通过Vela编译器（https://pypi.org/project/ethos-u-vela/）运行量化的TensorFlow Lite文件 #compile NN #optimize #Arm Ethos NPU

请参阅Vela页面上关于如何通过Vela优化器运行模型的说明。按照本指南中的步骤并通过Vela运行您的模型后，您会发现自己拥有一个完全可以在Arm Ethos-U55或Arm Ethos-U65上运行的模型。

结论

通过遵循本指南，您成功地在TensorFlow中培训了一个简单的基于RNN的模型。然后，您已经了解了如何Unroll RNN模型，以便可以将其量化并轻松转换为TensorFlow Lite格式，以便在Arm Ethos-U55上部署。您还了解了如何创建一个TFLite模型，其中可以提供初始RNN状态作为模型的输入，以及如何从模型中捕获最终RNN状态。

概述