AIE-ML 上的 MNIST ConvNet

版本：Vitis 2024.2

简介

本教程在 AMD Versal^TM 自适应 SoC AIE-ML 上实现了一个卷积神经网络分类器，用于识别来自 MNIST 数据库 的手写数字。目标是说明如何将一个简单的机器学习示例分区和向量化到 Versal AI 引擎。MNIST ConvNet 分类器是一个很好的学习示例，因为它只包含约 100,000 个参数和少量的层。本教程示例说明了使用 AIE API 自定义编码机器学习设计的许多关键主题，包括：

• 使用多通道矩阵乘法内在函数来向量化 ConvNet 层计算工作负载
• 使用内存块的 3D 寻址模式来访问计算消耗所需的层 I/O
• 使用本地块内存来捕获存储的网络权重
• 结构化 AIE-ML 内核代码以高效地实现卷积和池化层

虚拟 Python 环境设置

本教程使用 Jupyter Notebooks 中的 Python 实现 MNIST ConvNet 的计算机模型。最好使用 Python 虚拟环境来管理它。本教程的第一步是设置这样一个虚拟环境，其中包含所有相关开源 Python 包的所需版本。本教程的顶级 Makefile 基于脚本 create_venv.sh 构建 Python 虚拟环境，该脚本创建环境，然后加载一组特定版本的所需 Python 包。要创建此 Python 虚拟环境，请运行以下代码：

% make venv
% source my-venv/bin/activate
% python --version

这将在顶级文件夹中创建一个文件夹 my-venv，其中包含 Jupyter Notebooks、TensorFlow、matplotlib、pydot 和 bfloat16 所需的所有软件包（包括任何和所有依赖项）。第二个命令激活 Python 环境。第三个命令显示用于创建虚拟环境的 Python 版本。本教程已使用 Python 3.10.4 构建和测试。

Jupyter Notebook 模型

本教程的第一步是构建 MNIST ConvNet 的计算机模型，并训练该模型以获得可用于推理的一组权重。本教程为此目的使用 Keras 框架。Keras 提供了一个简单、灵活且强大的框架，构建在 TensorFlow 之上，用于构建机器学习网络模型。此模型的完整 Jupyter Notebook 由 MNIST-Convnet-demo.ipynb 提供。要运行 notebook，请执行以下命令：

% jupyter-notebook MNIST-Convnet-demo.ipynb

MNIST ConvNet 分类器的 Keras 模型源自 Keras 网站上的一个 示例，并由以下 Python 代码给出：

inputs = keras.Input(shape=(28,28,1),name="input")x1 = layers.Conv2D(filters=16,kernel_size=3,activation="relu",name="conv2d_w1")(inputs)x2 = layers.MaxPooling2D(pool_size=2,name="max_pooling2d_w2")(x1)x3 = layers.Conv2D(filters=64,kernel_size=3,activation="relu",name= "conv2d_w3")(x2)x4 = layers.MaxPooling2D(pool_size=2,name="max_pooling2d_w4")(x3)x5 = layers.Conv2D(filters=128,kernel_size=3,activation="relu",name="conv2d_w5")(x4)x6 = layers.Flatten(name="flatten_w6")(x5)outputs = layers.Dense(10,activation="softmax",name="dense_w7")(x6)model = keras.Model(inputs=inputs,outputs=outputs)model.compile(optimizer="rmsprop",loss="sparse_categorical_crossentropy",metrics=["accuracy"])

该网络包含七个层：

• 第一层是一个 2D 卷积层，将十六个 3x3 卷积核应用于输入层，并在其输出处应用 ReLU 激活函数。每个内核都有自己的一组 3 x 3 = 9 个乘法权重以及一个单一的加法偏置权重。因此，该层的总权重为 16 x (9+1) = 160。
• 第二层实现一个“最大池化”层，该层在两个图像维度中执行二分抽取。此层没有关联的权重。
• 第三层实现另一个 2D 卷积层，这次有 64 个 3x3 卷积核，并在其输出处应用 ReLU 激活函数。该层总共涉及 16 x 64 x 9 = 9,216 个乘法权重和 64 个加法偏置权重，总共 9,280 个权重。
• 第四层实现另一个二维“最大池化”层，没有额外的权重。
• 第五层实现第三个 2D 卷积层，具有 128 个 3x3 卷积核，并在其输出处应用 ReLU 激活。该层使用总共 64 x 128 x 9 = 73,728 个乘法权重和 128 个加法偏置权重，总共 73,856 个权重。
• 第六层执行展平功能，将总共 3 x 3 x 128 = 1,152 个连接折叠到单个 1D 总线中。
• 第七层由一个完全连接的密集网络组成，具有 1152 x 10 = 11,520 个乘法权重和 10 个加法偏置权重，总共 11,530 个权重。

此网络的总权重数为 94,826 个，或大约 100,000 个权重。这是一个小的 ConvNet。下图总结了 MNIST ConvNet 的层。

导入 MNIST 图像数据库

MNIST ConvNet 必须在 MNIST 图像数据库上进行训练。此数据库包含 60,000 个手写数字图像，并作为 Keras/TensorFlow 包的一部分分发。MNIST-Convnet-demo.ipynb 选择 3750 x 16 个图像用于以 16 个批次训练网络，并选择另外 512 x 16 个图像用于以 16 个批次测试经过训练的网络。每个图像由 28 x 28 个单色像素（范围为 0 到 255）表示。提取这些训练和测试图像集的代码如下所示，以及其中八个手写图像的示例。

BS=16
NB_trn = 3750
NB_tst = 512
# 加载 MNIST 数据库：
(trn_images,trn_labels), (tst_images,tst_labels) = mnist.load_data()
trn_images = trn_images.reshape((60000,28,28,1))
trn_images = trn_images[:NB_trn*BS,:,:,:]
tst_images = tst_images.reshape((10000,28,28,1))
# 提取可用数据：
tst_images = tst_images[:NB_tst*BS,:,:,:]
trn_labels = trn_labels[:NB_trn*BS]
tst_labels = tst_labels[:NB_tst*BS]
fig,ax = plt.subplots(nrows=2,ncols=4)
for rr in range(2):for cc in range(4):ax[rr,cc].imshow(trn_images[4*rr+cc],cmap=plt.cm.binary)

训练和测试 MNIST ConvNet 模型

Keras 框架提供内置函数用于训练和测试模型。这些如下所示。训练在训练图像集上执行。标记的数据集用于驱动反向传播算法，以调整权重以最小化所选的成本函数（在本例中为“稀疏分类交叉熵”）。选择“准确度”指标来评估结果质量。经过五个训练周期后，在测试数据集上实现了 99.1% 的网络准确度。

使用 MNIST ConvNet 进行推理

另一个内置的 Keras 例程可用于从测试输入图像中预测新的模型输出。下面的代码显示了四个示例图像，这些图像已标记有 ConvNet 生成的（正确的）数字识别。

提取权重和偏置以用于 AIE-ML 推理解决方案

现在已经获得了 MNIST ConvNet 分类器的训练模型，构建 AIE-ML 上的推理解决方案之前的最后一步是获得一组量化的权重以供实现使用。为了简单起见，本教程选择了一个 bfloat16 实现，因为量化很简单。每个权重和偏置使用相同的指数，但量化的尾数只有 8 位，而不是 Keras 网络模型生成的全精度浮点值使用的 24 位。下面的代码从 Keras 模型中提取权重和偏置，将它们量化为 bfloat16，然后将它们保存在文件中，用于验证下面将在 AIE-ML 中设计的网络的每一层。

AIE-ML 推理解决方案

本节概述了 MNIST ConvNet 分类器的最终 AIE-ML 设计，包括对设计中所有层利用的关键原则的审查。然后，在 各个层设计 中给出了每个单独层的详细信息。

设计方法

本教程采用了一种简单的设计方法，用于在 AIE-ML 中构建 MNIST ConvNet 分类器，旨在建立可行的数据流，并确定用于常见卷积和池化层的高效 AIE API 编码策略。重点不是吞吐量性能或资源利用率，而是确定一套可以应用于更高级设计的设计技术。该设计方法采用以下概念：

• 选择 bfloat16 数据类型用于层 I/O 数据以及权重和偏置。这简化了训练网络参数的量化。不需要特殊的工具或量化策略。
• 没有选择特定的吞吐量目标。该设计是一个玩具示例，因此其性能没有实际意义。
• 该设计通常将每个网络层分区到其自己的 AIE-ML 块（如果可行）。这简化了系统分区，并允许为要构建的每个 AIE-ML 内核定义明确的范围。
• 利用内存块零填充功能将输入张量形状从 (28,28,1) 扩展到 (28,32,1)，以满足 AI 引擎内存对齐和 PLIO 位宽要求。
• 利用内存块多维寻址功能，以高效地传输 I/O 数据以进行计算消耗，而无需额外的内核周期用于内核内的数据洗牌或通道调整。
• 用于 2D 卷积层的计算工作负载利用高效的 mac_4x8_8x4() 本征函数用于 bfloat16 数据类型，以在特定层可行的情况下实现每个周期 128 个 MAC 操作的最大效率。
• 计算工作负载利用效率较低的 mac_elem_16_2() 本征函数用于 bfloat16 数据类型，在 mac_4x8_8x4() 不可行的情况下，每个周期最多 64 个 MAC 操作（例如，仅从单个输入通道接收数据的 conv2d_w1() 层）。
• 该设计将 flatten_w6() 和 dense_w7() 层组合到同一个 AIE-ML 块中。
• 权重和偏置存储在本地块内存中，而不是存储在内存块中，因为前者允许使用异步缓冲区建立只读访问方案，该方案允许仅在启动时读取一次权重。具有数百万个权重的更大的 ML 网络需要基于内存块的流式解决方案；对于此处考虑的小型 MNIST 问题，这种复杂的解决方案是过度的，因为所有权重都可以轻松地存储在阵列中。扩展编程模型以支持内存块中的只读操作正在开发中。

Vitis 功能仿真

本教程使用一种名为 Vitis 功能仿真 (VFS) 的新工具功能，用于针对其 Python 行为模型验证 MNIST ConvNet 分类器 AI 引擎实现。VFS 功能生成异构 Versal 设计的 AI 引擎和 PL 部分的可执行“共享对象”，从而允许将其引入到熟悉的仿真框架中，即 MATLAB 和 Python。这允许在不离开您首选的仿真框架且不为此目的创建 I/O 文件的情况下，对 Versal AI 引擎和 PL 设计进行功能验证。

Vitis 功能仿真是 2024.2 中的 EA 工具功能。要获取说明和设计示例，请请求访问 Vitis 功能仿真抢先体验安全站点。

Vitis 功能仿真的 Python 版本用于验证 MNIST ConvNet 设计中的每个网络层。下面的屏幕截图显示了如何从用于验证该层的 Jupyter notebook gen_vectors.ipynb 调用 conv2d_w1() 层。为设计中的每一层构建了一个类似的 notebook。当任何源代码文件发生更改时，VFS 基础结构会自动在后台编译 AI 引擎图，然后在 Jupyter notebook 和 x86 功能模拟器之间传递 I/O 向量。

MNIST ConvNet：AI 引擎图视图

MNIST ConvNet 分类器的整体 AI 引擎图如下所示。

• 每个层通常分配到其自己的 AIE-ML 块，如上所述。
• 带有权重和偏置的层包含两个 AIE-ML 块。一个块对输入图像执行计算工作负载。权重从权重传递块传入。设计启动时将来自 PLIO 的权重加载到缓冲区，并传递到计算块。异步缓存机制读取 PLIO 中的权重，并在设计启动时将它们传送到权重输入缓冲区。计算块便可以在设计运行时持续地访问这些权重。
• conv2d_w5() 层在四个块中进行分区，以管理权重存储，该存储太大而无法放入可用的 32 KB 本地块存储中（ping/pong 缓冲区限制为 32 KB，因为它们必须分配给同一 64 KB 本地块内存）。根据其大约 74,000 个参数，存储必须在最少四个 AIE-ML 块上进行分割。其计算工作负载也在四个块上进行分区，每个块计算四分之一的输出层样本。下面概述了更多详细信息。
• 最后一个 AIE-ML 块包含 flatten_w6() 层、dense_w7() 层和 softmax() 计算工作负载以生成最终分类器输出。
  最大池化层 max_pooling2d_w2() 和 max_pooling2d_w4() 没有任何权重，因此使用每个块一个 AIE-ML 来实现。

MNIST ConvNet：AI 引擎布局视图

MNIST ConvNet 分类器的布局视图如下所示。放置约束将计算块放置在块的顶行中，并将权重传递块放置在块的下行中。该设计使用内存块进行层 I/O 排序和零填充，如下面更详细的概述。

MNIST ConvNet：AI 引擎资源利用率

设计的资源利用率如下图所示。该设计适合 2 x 9 的块网格，并利用五个内存块用于共享缓冲区。

3x3 Conv2D 层处理的向量化

MNIST ConvNet 中的 2D 卷积层使用 3x3 补丁处理，如下图所示。每一层都有许多大小为 $H_I\times W_I$ 像素的图像的输入通道 $C_I$。每个输入图像都由 3x3 补丁处理，以生成许多大小为 $H_O\times W_O$ 像素的输出通道 $C_O$，其中 $H_O=H_I-2$ 和 $W_O=W_I-2$。一个像素宽的外边框在输出图像中丢失，以允许 3x3 补丁完全跨越输入图像而不超过其边界。3x3 补丁处理涉及计算输入图像像素与 3x3 补丁的 9 个权重的 9 个点积，并将结果相加。将偏置项添加到总和中，这给出了 3x3 补丁中心处输出像素的输出值。下图显示了一个 $5\times 5$ 输入图像及其对应的 $3\times 3$ 输出图像。每个输出像素都是通过将补丁向右移动一个像素来计算的，然后在输出图像的每一行中重复此过程。每个输入/输出层对都有一组独特的 9 个权重 + 1 个偏置，用于计算从特定输入层到特定输出层的 3x3 补丁。

AIE-ML 数据路径针对矩阵乘法进行了优化，并且事实证明，上面概述的 3x3 卷积处理可以转换为这种形式。下图显示了 AIE-ML 中 bfloat16 数据类型的 mac_4x8_8x4() 本征函数执行的计算。此本征函数执行 $M=X\times Y$ 形式的矩阵乘法，其中矩阵 $X$ 的大小为 $[A\times B]$，矩阵 $Y$ 的大小为 $[B\times C]$，矩阵 $M$ 的大小为 $[A\times C]$。在这种情况下，$A\times B\times C=4\times 8\times 4$。输入矩阵 $X$ 为 $4\times 8$，输入矩阵 $Y$ 为 $8\times 4$，输出矩阵 $M$ 为 $4\times 4$。

下图显示了如何加载此本征函数以执行 3x3 卷积层处理。可以使用 8 个通道的每个通道 4 个像素来加载输入矩阵 $X$，其中像素存储在每个通道的列中。AI 引擎以行优先顺序将此输入矩阵映射到其向量通道中。因此，我们可以考虑将矩阵 $X$ 映射到 32 通道寄存器中，其中前 8 个通道包含来自所有 8 个通道的像素 0，第二个 8 个通道包含来自所有 8 个通道的像素 1，依此类推。

权重可以加载为矩阵 $Y$ 的列，其中每个权重 $w(C_i,C_o)$ 从输入通道映射到输出层。由于有 8 行和 4 列，因此每个权重从 8 个输入通道映射到 4 个输出通道。权重不是像素的函数；相同的权重用于特定的图像。

基于此向量化，我们可以在一个周期内从 32 个输入样本（来自八个输入通道的四个像素）计算 16 个输出样本（四个输出通道的每个通道四个像素）。然后，对于每个 AIE-ML 内核的技巧是高效地加载这 $4\times 8$ 个输入样本和 $8\times 4$ 个权重，以便持续地保持计算繁忙。

MNIST ConvNet：分析和向量加载

下表总结了 AIE-ML 上 MNIST ConvNet 分类器设计的分析特征。每个函数和内在函数都针对每一层列出，以及实现所需的 AI 引擎周期数。#MACs 列列出了理论上基于其输入和输出张量形状的每一层所需的乘法累加运算总数。可以通过硅能够执行的 MAC 操作数来缩放该值，以评估每一层的平均“向量加载”的估计值。根据具体层，实现了 13% 到 30% 到 50% 的负载。这些值略低于在“真实世界”网络中通常实现的值。MNIST ConvNet 使用非常小的图像，从仅 28 x 28 像素开始。因此，开销成本更高，因为我们无法在这些小图像上保持连续的计算周期，而不会快速超出内部计算循环。

MNIST ConvNet：吞吐量

基于 AI 引擎仿真，MNIST ConvNet 分类器的吞吐量约为每秒 70,000 帧。从这个小例子的实际角度来看，这并没有什么意义，但为了完整起见，将其包含在内。

各个层设计

层设计详情：conv2d_w1()

下图总结了 conv2d_w1() 层设计的关键方面。用于验证的 Jupyter Notebook 是 gen_vectors.ipynb。

• 使用输入内存块将输入图像从 (28,28,1) 的张量零填充到 (28,32,1) 的张量。此零填充在图像的右侧引入 4 列零，因此总列数是 32 的倍数。只有列维度需要填充，因为它构成了内部循环。由于该设计使用 bfloat16 数据，并且内存块需要 32 位对齐，因此输入内存块设计为采用四个图像（即，具有 (4,28,28,1) 张量），并且 conv2d_w1_graph 设置为多速率解决方案，其中内存块的 repetition_count=1 和计算内核的 repetition_count=4。这是贯穿整个设计的关键原则。
• 由于该层只有一个输入通道，因此 mac_elem_16_2() 本征函数以 50% 的容量使用；它默认处理两个通道，这里一个通道被归零。这会影响其整体向量效率。
• 实现了内部循环 II=17 以传递 9 个 MAC 操作，这只有 26% 的效率。由于此处图像的性质较小，因此难以用更多的 MAC 操作来填充管道。在更大的设计中可以轻松地完成此操作。
• 整体循环结构采用输出图像行的外部循环和输出图像列的内部循环。这非常适合所选的本征函数。
• 请注意内存块的平铺参数如何用于使用四个附加的零值像素来填充输入图像的列维度。

层设计详情：max_pooling2d_w2()

下图总结了 max_pooling2d_w2() 层设计的关键方面。用于验证的 Jupyter Notebook 是 gen_vectors.ipynb。

• 最大池化通过在 2x2 补丁中的所有四个像素上应用 max() 操作来将每个维度中的输入图像抽取 2 倍。连续补丁以 2 为步长，因此它们是非重叠的。可以使用 AIE API 的 aie::max() 函数有效地向量化此计算工作负载。
• 该层被编码为输出图像行的外部循环和图像列的内部循环。向量化为每个通道 4 个像素创建 16 个输出通道。
• 内部循环的软件流水线处理非常完美，II=16。

层设计详情：conv2d_w3()

下图总结了 conv2d_w3() 层设计的关键方面。用于验证的 Jupyter Notebook 是 gen_vectors.ipynb。

• 输入内存块以线性顺序写入样本，首先按最右边的维度（如在 Numpy 约定中）按顺序使用形状为 (13,16,16) 的输入张量。以平铺方式从输入内存块中提取样本，以便样本可以立即被计算消耗，而无需任何额外的洗牌。8 x 16 个样本的平铺从每个 8 个输入层提取 16 个像素。这对应于四个单独的 4 x 8 样本块，这些样本块可以被 mac_4x8_8x4() 本征函数使用。
• 计算工作负载计划在输出图像行的外部循环和输出图像列的内部循环中进行。每次计算在一个周期内生成四个像素（用于四个输出通道）。该层生成 (11,16,64) 的输出张量形状。
• 实现了循环 II=12 个周期，用于 12 个 MAC 操作（100% 内部循环效率）。
• 输出内存块以 4x4 个平铺存储样本，这样内核本身就不需要输出洗牌；它由内存块 DMA 引擎处理。以这种平铺方式将样本写入输出内存块中，以便可以在输出张量顺序 (11,16,64) 中轻松提取它们，以供后续层遵循。

层设计详情：max_pooling2d_w4()

下图总结了 max_pooling2d_w4() 层设计的关键方面。其设计与第二层非常相似，图像尺寸略小，但要处理的 I/O 通道更多。代码结构相似，并且编译器实现了高效的软件流水线调度。该层不需要任何内存块进行样本重新排序。用于验证的 Jupyter Notebook 是 gen_vectors.ipynb。

层设计详情：conv2d_w5()

下图总结了 conv2d_w5() 层设计的关键方面。用于验证的 Jupyter Notebook 是 gen_vectors.ipynb。该第五层有许多独特的方面：

• 由于其权重的沉重存储需求，该层在四个块上进行分割。每个块处理总输出通道的四分之一，并且每个本地块内存存储总权重的四分之一。
• 该层使用输入和输出内存块，使用与 conv2d_w3() 层类似的方法，以按计算消耗所需的顺序提取和替换样本。这可以有效地利用使用 max_4x8_8x4() 的高性能计算，并可以实现完美的内部循环软件流水线处理。
• 由于计算在四个计算块上进行分区，因此必须在处理完成后执行“收集”函数。输出通道结果必须从每个计算块传递到第四个块以进行收集和重新排序。该设计为此目的使用 I/O 流。这需要一个额外的样本重新排序过程，该过程将在下面进一步详细概述。

下图说明了 conv2d_w5() 层所需的输出样本收集过程，以便收集并恢复其四个块计算出的输出通道集合中的样本顺序。输入内存块以 8 x 8 平铺模式提取输入样本，并且这些样本广播到所有四个计算块。每个块计算其分配的输出通道部分。一旦被第四个块收集，这些四个数据集必须通过第四个计算块重新洗牌为正确的顺序，然后才能通过输出内存块使用的输出 4 x 8 平铺模式进行提取。下图中的彩色块有助于识别每个块计算的各个输出图像行。必须按行大小交错这些行以恢复层输出处所需的 (3,8,128) 张量形状。这是通过第四个块的额外计算周期来执行的。收集时，每个数据集都复制到本地块中的草稿内存中。然后，以正确的顺序读取这些四个草稿区域，以产生所需的洗牌。

层设计详情：dense_w7()

下图总结了 dense_w7() 层设计的关键方面。用于验证的 Jupyter Notebook 是 gen_vectors.ipynb。该块包括三个函数，flatten_w6() 层、dense_w7() 层和最终的 softmax() 计算。该层使用 mac_elem_16_2() 本征函数每次计算两个输出。每个本征函数消耗 32 个输入通道。softmax() 激活函数使用 Softmax 函数 Vitis 教程 中概述的方法进行计算

总结

本教程介绍了 AIE-ML 中 MNIST ConvNet 分类器的设计。该解决方案具有约 100,000 个参数，需要约 18 个块。它实现了约每秒 70K 帧的吞吐量。该设计说明了使用为信号处理开发的 Vitis AI 引擎工具流程的“数据流”模式为机器学习网络构建“白盒”设计的许多方面。此简单设计的各个层之间代码结构的相似性表明，实际上，机器学习层库应该是可行的。