第6章:实战演练——把MNIST CNN映射到AIE-ML阵列
恭喜你来到最后一章!前面我们学会了数据搬运、高速公路管理、实时调音、防止堵车,还指挥过FFT交响乐团。现在,我们要把这些技能全部用上,搭建一个真正能识别手写数字的AI加速器——就像在Versal芯片上开一家专门的“数字识别工厂”。
6.1 首先,我们要搭什么工厂?
想象一下:你开了一家工厂,专门识别手写的0-9数字。工厂的入口是一张28x28像素的黑白照片,出口是一个标签,写着“这是5!”。
这家工厂有三条核心生产线:
- 卷积车间:用“放大镜”(卷积核)扫描照片,找边缘、圆圈这些特征
- 池化车间:把扫描结果缩小,保留最明显的特征
- 分类车间:把所有特征汇总,判断到底是哪个数字
我们的任务,就是把这三条生产线搬到AIE-ML的空间阵列上。
6.2 工厂规划:整体架构
先看看我们的数字识别工厂长什么样:
graph TD
subgraph "PS 主控室"
PS[ARM CPU 主机]
PS_MM2S[DMA 读控制器]
PS_S2MM[DMA 写控制器]
end
subgraph "PL 运输部"
PL_IN[HLS 数据输入打包器]
PL_W[HLS 权重加载器]
PL_OUT[HLS 结果输出解包器]
end
subgraph "AIE-ML 生产车间"
subgraph "Layer 1"
CONV1[Conv2D 5x5 车间]
POOL1[MaxPool 2x2 车间]
end
subgraph "Layer 2"
CONV2[Conv2D 5x5 车间]
POOL2[MaxPool 2x2 车间]
end
subgraph "Layer 3"
DENSE[Dense 分类车间]
end
W_FIFO[权重暂存区]
end
PS --> PS_MM2S
PS_MM2S --> PL_IN
PL_IN --> CONV1
PS --> PS_MM2S
PS_MM2S --> PL_W
PL_W --> W_FIFO
W_FIFO --> CONV1
W_FIFO --> CONV2
W_FIFO --> DENSE
CONV1 --> POOL1
POOL1 --> CONV2
CONV2 --> POOL2
POOL2 --> DENSE
DENSE --> PL_OUT
PL_OUT --> PS_S2MM
PS_S2MM --> PS
图6.1 MNIST识别系统的整体架构
架构解读:
- PS主控室:负责把照片和“放大镜说明书”(权重)送进工厂,再把结果取出来
- PL运输部:负责把数据打包成AIE-ML喜欢的“流”格式,就像把散装零件装进传送带
- AIE-ML生产车间:核心计算区,每层都是一条小流水线
- 权重暂存区:让每个车间随时能拿到“放大镜说明书”,不用等PS慢慢送
6.3 第一条生产线:卷积车间
卷积是找特征的核心步骤——就像用不同形状的放大镜扫描照片,每次找一种特征。
6.3.1 卷积的空间映射
我们把一个卷积核拆分到多个AIE-ML tile上,就像把一个大任务拆给多个工人:
graph TD
subgraph "Input Feature Map"
I1[像素块 1]
I2[像素块 2]
I3[像素块 3]
I4[像素块 4]
end
subgraph "AIE-ML Tile 阵列"
T1[Tile 0,0
负责卷积核 0-3] T2[Tile 1,0
负责卷积核 4-7] T3[Tile 0,1
负责乘加累加] T4[Tile 1,1
负责乘加累加] end subgraph "Weights Stream" W1[权重 0-3] W2[权重 4-7] end I1 --> T1 I2 --> T1 I3 --> T2 I4 --> T2 W1 --> T1 W2 --> T2 T1 --> T3 T2 --> T3 T1 --> T4 T2 --> T4 T3 --> O1[输出特征图 0-3] T4 --> O2[输出特征图 4-7]
负责卷积核 0-3] T2[Tile 1,0
负责卷积核 4-7] T3[Tile 0,1
负责乘加累加] T4[Tile 1,1
负责乘加累加] end subgraph "Weights Stream" W1[权重 0-3] W2[权重 4-7] end I1 --> T1 I2 --> T1 I3 --> T2 I4 --> T2 W1 --> T1 W2 --> T2 T1 --> T3 T2 --> T3 T1 --> T4 T2 --> T4 T3 --> O1[输出特征图 0-3] T4 --> O2[输出特征图 4-7]
图6.2 卷积层的空间拆分示意
卷积车间的工作原理:
- 输入拆分:把28x28的照片拆成小像素块,每个tile负责处理一部分
- 权重拆分:把多个“放大镜”(卷积核)分给不同的tile
- 并行计算:每个tile用自己的像素块和放大镜做乘加运算
- 结果汇总:把相邻tile的结果拼起来,得到完整的输出特征图
这就像让4个工人一起拼拼图——每人拼一块,最后合在一起。
6.4 第二条生产线:池化车间
池化的工作很简单:把扫描结果缩小,只保留最亮的那个点(最大池化)。这就像拍集体照时,只拍最高的那个人。
6.4.1 池化的数据流
池化车间不需要太多计算,但需要巧妙地安排数据流动:
flowchart LR
A[输入特征图
24x24] --> B[滑动窗口
2x2] B --> C{找最大值?} C -->|是| D[保存最大值] C -->|否| E[丢弃] D --> F[移动窗口] F --> B D --> G[输出特征图
12x12]
24x24] --> B[滑动窗口
2x2] B --> C{找最大值?} C -->|是| D[保存最大值] C -->|否| E[丢弃] D --> F[移动窗口] F --> B D --> G[输出特征图
12x12]
图6.3 最大池化的工作流程
池化车间的优化:
- 滑动窗口:一次看4个像素,只留最大的
- 流水线并行:和卷积车间连在一起,卷积做完一个窗口,池化马上处理
- 内存复用:尽量在tile的本地内存里完成,不用去远处搬数据
6.5 第三条生产线:分类车间
分类车间是最后一步——把所有特征汇总,判断是哪个数字。这就像侦探把所有线索拼起来,找出凶手。
6.5.1 全连接层的矩阵乘法
全连接层本质上是一个大矩阵乘法——我们把它拆成小矩阵,分到多个tile上:
classDiagram
class InputVector {
+float[128] features
}
class WeightMatrix {
+float[10][128] weights
}
class TileGroup1 {
+float[5][64] local_weights
+multiply_add()
}
class TileGroup2 {
+float[5][64] local_weights
+multiply_add()
}
class OutputVector {
+float[10] scores
}
InputVector --> TileGroup1
InputVector --> TileGroup2
WeightMatrix --> TileGroup1
WeightMatrix --> TileGroup2
TileGroup1 --> OutputVector
TileGroup2 --> OutputVector
图6.4 全连接层的并行拆分
分类车间的秘密武器:
- 矩阵拆分:把10x128的大矩阵拆成4个小矩阵,每个tile算一个
- 乘加累加:用AIE-ML的SIMD指令,一次算好多个乘法
- 结果聚合:把所有tile的分数加起来,最高分对应的就是识别结果
6.6 把所有车间连起来:完整流水线
现在我们把三条生产线连起来,形成一条完整的“数字识别流水线”:
sequenceDiagram
participant PS as ARM 主控室
participant PL as PL 运输部
participant C1 as Conv1 车间
participant P1 as Pool1 车间
participant C2 as Conv2 车间
participant P2 as Pool2 车间
participant D as Dense 分类车间
PS->>PL: 发送图片和权重
PL->>C1: 图片流
PL->>C1: 权重流
C1->>P1: 特征图1
P1->>C2: 缩小特征图1
PL->>C2: 权重流
C2->>P2: 特征图2
P2->>D: 缩小特征图2
PL->>D: 权重流
D->>PL: 识别结果
PL->>PS: 返回结果
图6.5 完整流水线的工作序列
流水线的关键技巧:
- 乒乓缓冲:每个车间都有两个“工作台”——一个在计算,一个在装数据
- 权重预加载:在计算开始前,先把权重送到每个车间的暂存区
- 流控平衡:调整每个车间的FIFO深度,不让任何一个车间堵车
6.7 动手试试:运行MNIST示例
现在我们来真正运行这个数字识别工厂!
6.7.1 步骤概览
flowchart TD
A[克隆 Vitis-Tutorials 仓库] --> B[进入 MNIST 教程目录]
B --> C[设置 Vitis 环境变量]
C --> D[运行软件仿真 x86sim]
D --> E[检查仿真结果是否正确]
E --> F[运行 AIE 周期精确仿真 aiesim]
F --> G[查看性能报告]
G --> H[生成硬件比特流]
H --> I[在开发板上运行]
图6.6 运行MNIST示例的步骤
6.7.2 关键命令(简化版)
# 1. 进入教程目录
cd Vitis-Tutorials/AI_Engine_Development/AIE-ML/Design_Tutorials/02-MNIST_ConvNet
# 2. 设置环境
source /tools/Xilinx/Vitis/2024.1/settings64.sh
# 3. 运行软件仿真(验证功能)
make x86sim
# 4. 运行AIE仿真(验证性能)
make aiesim
# 5. 查看结果
cat aiesimulator_output/output.txt
6.8 本章总结与全书回顾
恭喜你!你已经完成了整个AIE-ML初学者指南的学习!
本章收获:
- 知道了如何把CNN的卷积、池化、全连接层映射到AIE-ML阵列
- 理解了空间并行和流水线并行的结合
- 学会了运行完整的MNIST识别示例
全书回顾:
- 第1章:学会了数据如何在PS、PL、AIE之间流动
- 第2章:掌握了用包交换管理多条数据流
- 第3章:学会了用RTP实时调整参数
- 第4章:掌握了用FIFO防止堵车和死锁
- 第5章:学会了分解和映射复杂的FFT算法
- 第6章:(本章)完成了完整的CNN映射实战
下一步:
- 试着修改MNIST示例,用更大的卷积核
- 尝试把其他网络(比如小的YOLO)映射到AIE-ML
- 探索Vitis AI Library,里面有更多预优化的AI算子
祝你在AIE-ML的开发之旅中一路顺风!