3. KPU核心概念

KPU(Knowledge Processing Unit)是K210芯片内置的神经网络硬件加速器,其设计目标是在低功耗、低成本的边缘场景下完成轻量级CNN模型的实时推理。与CPU上运行TensorFlow或PyTorch不同,KPU并非通用的可编程处理器,而是一条高度定制化的卷积运算流水线。理解KPU的运算原理、支持的网络层类型、输入输出格式以及内存限制,是进行模型转换、算子适配和性能调优的前提。本章将从硬件架构视角切入,帮助你在部署前建立正确的“KPU心智模型”。

KPU 核心概念总览:从输入到输出的数据通路 输入特征图 尺寸 ≤ 320×320 通道数 ≤ 1024 KPU 卷积运算核心 3×3 / 1×1 卷积 深度可分离卷积 最大池化 / 平均池化 激活函数(ReLU等) 输出特征图 尺寸逐层递减 通道数逐层递增 全连接层(可选) 仅支持1×1卷积替代 KPU 内存架构 内部SRAM:2MB(AI专用) 输入/输出缓冲区:独立寻址 权重存储:压缩后加载 行缓冲器:用于卷积滑动窗口 外部DMA通路 6MB 通用SRAM(DMA访问) Flash 存储 kmodel 数据搬运由CPU/DMA控制 推理时数据分块加载 关键限制与适配要求 • 输入通道数必须为 1/3/4 的倍数(通常为4的倍数) • 输出通道数必须为 8 的倍数(KPU硬件对齐要求) • 不支持Softmax、Sigmoid、Tanh等非线性激活(需在CPU上执行) • 不支持LSTM、GRU、Transformer等时序/注意力结构

上图展示了KPU从输入特征图到输出特征图的完整数据通路。可以看到,KPU内部集成了卷积运算核心、内存管理单元和DMA控制器,外部通过6MB通用SRAM与Flash进行数据交换。理解这条通路上的每一个环节,是后续进行模型转换和算子适配的基础。

3.1 KPU运算原理:硬件加速的本质

KPU的运算原理可以概括为“数据流驱动的卷积引擎”。与CPU的冯·诺依曼架构不同,KPU内部采用类似脉动阵列(Systolic Array)的设计,将卷积运算中的乘加操作(MAC)并行化。具体而言,KPU在单个时钟周期内可以完成多次乘加运算,其峰值算力可达0.8 TOPS(每秒万亿次操作),在400MHz主频下,处理224×224×3的输入图像时,典型推理延迟在100ms以内。

KPU的运算流程分为三个关键阶段:

  1. 数据加载阶段:DMA控制器从外部SRAM或Flash中将输入特征图、权重参数和偏置项搬运到KPU内部的SRAM缓冲区。这一阶段是性能瓶颈的主要来源,因为外部存储的带宽远低于KPU的运算速度。
  2. 卷积计算阶段:KPU内部的卷积引擎按照预设的步长(Stride)和填充(Padding)规则,对输入特征图进行滑动窗口计算。每个卷积核的权重被预先加载到寄存器中,与输入数据逐元素相乘并累加,最终加上偏置项。
  3. 激活与池化阶段:卷积结果经过激活函数(如ReLU)处理后,可选择性地进行池化操作(最大池化或平均池化),然后写入输出缓冲区。

值得注意的是,KPU的卷积引擎对通道数有严格的对齐要求。输入通道数必须是4的倍数,输出通道数必须是8的倍数。这是因为KPU的硬件乘法器阵列按4×8的维度组织,如果通道数不满足对齐要求,硬件会自动填充零值,导致计算浪费。

核心原理小结:KPU不是通用的可编程处理器,而是一条针对CNN卷积运算优化的专用流水线。它的优势在于高能效比(每瓦特算力远高于CPU),但代价是灵活性受限——所有运算必须符合硬件预定义的规则。

3.2 支持的网络层类型:KPU的“能力边界”

KPU并非支持所有神经网络层,它只覆盖了CNN中最核心、计算量最大的算子。下表列出了KPU支持与不支持的层类型,这是模型转换前必须对照检查的清单。

层类型 KPU支持情况 硬件实现方式 注意事项
标准卷积(Conv2D) ✅ 支持 专用卷积引擎 卷积核尺寸支持1×1、3×3,其他尺寸需转换或拆分
深度可分离卷积(Depthwise Separable Conv) ✅ 支持 分两步执行:逐通道卷积 + 1×1逐点卷积 MobileNet等轻量网络的核心算子,KPU原生支持
最大池化(Max Pooling) ✅ 支持 池化引擎 支持2×2、4×4池化核,步长需与核大小一致
平均池化(Average Pooling) ✅ 支持 池化引擎 全局平均池化(GAP)常用于分类网络末端
全连接层(Fully Connected) ⚠️ 间接支持 需转换为1×1卷积 转换后输入输出通道数需满足对齐要求
ReLU激活 ✅ 支持 激活单元 ReLU6等变体需在转换时映射为标准ReLU
Softmax / Sigmoid / Tanh ❌ 不支持 需在CPU上执行 分类网络的Softmax层通常放在KPU输出之后
Batch Normalization(BN) ⚠️ 需融合 转换时融合进卷积层 训练时BN参数需在转换工具中合并到卷积权重
残差连接(Residual Add) ✅ 支持 特征图加法器 ResNet等结构可直接支持
上采样 / 反卷积 ❌ 不支持 需在CPU上实现或改用其他结构 语义分割模型需特别注意
LSTM / GRU / Transformer ❌ 不支持 无对应硬件 时序模型无法部署到KPU

从表中可以看出,KPU的设计目标非常明确:覆盖图像分类、目标检测、关键点检测等主流视觉任务。如果你的模型包含KPU不支持的层,通常有两种解决方案:一是修改模型结构,用支持的操作替代;二是将不支持的部分拆分到CPU上执行,但这样会增加推理延迟。

3.3 输入输出格式:数据如何“喂”给KPU

KPU的输入输出格式有严格的规范,理解这些规范是避免运行时错误的关键。KPU的输入数据格式为RGB888或GRAY888,即每个像素由3字节(RGB)或1字节(灰度)表示。在送入KPU之前,图像数据需要按以下规则排列:

  • 数据排列顺序:按通道交错排列(NHWC格式),即先存储第一个像素的R、G、B值,再存储第二个像素的R、G、B值,以此类推。
  • 尺寸对齐:输入图像的宽度和高度必须是4的倍数。如果原始图像尺寸不满足,需要在预处理阶段进行填充(Padding)或缩放。
  • 通道数限制:输入通道数必须为1、3或4。当使用4通道输入时,第4个通道通常填充为0或用于其他目的(如RGBD数据)。

输出特征图的格式与输入类似,但通道数必须为8的倍数。例如,一个分类网络的输出层如果有10个类别,KPU会将其对齐为16个通道,多余的6个通道填充为0。因此,在解析输出时,需要根据实际的类别数截取有效部分。

以下是一个典型的输入数据预处理代码示例(使用MicroPython):

import sensor
import image

# 初始化摄像头
sensor.reset()
sensor.set_pixformat(sensor.RGB565)  # 摄像头输出RGB565格式
sensor.set_framesize(sensor.QVGA)    # 320x240

# 转换为KPU需要的RGB888格式
img = sensor.snapshot()
img_rgb888 = img.to_rgb888()  # 转换为RGB888

# 尺寸对齐:确保宽高为4的倍数
width = img_rgb888.width()
height = img_rgb888.height()
if width % 4 != 0:
    width = (width // 4 + 1) * 4
if height % 4 != 0:
    height = (height // 4 + 1) * 4
img_rgb888.resize(width, height)

# 将图像数据传递给KPU
kpu_input = img_rgb888.to_bytes()  # 获取字节流

常见错误:很多初学者在部署时忽略输入尺寸对齐要求,导致KPU推理结果异常或直接报错。务必在预处理阶段检查并调整图像尺寸。

3.4 内存限制:KPU的“生存空间”

K210芯片的内存资源非常有限,这是边缘部署最大的约束之一。KPU可用的内存分为两部分:内部2MB SRAM(专用于AI运算)和外部6MB通用SRAM(通过DMA访问)。在模型部署时,需要同时考虑权重存储、中间特征图存储和输入输出缓冲区的占用。

具体的内存分配策略如下:

  • 权重存储:KPU的权重参数经过量化压缩后(通常为8位整数),存储在内部SRAM中。一个典型的MobileNet模型(约4.2M参数)量化后约占用4.2MB,这已经超出了内部SRAM的容量,因此需要将权重分块加载,或者使用更小的模型结构。
  • 中间特征图:每一层卷积的输出特征图需要暂存在内部SRAM中,作为下一层的输入。特征图的内存占用计算公式为:宽度 × 高度 × 通道数 × 1字节(8位量化)。例如,一个112×112×32的特征图占用约400KB。
  • 输入输出缓冲区:KPU需要预留一部分SRAM作为输入和输出的DMA缓冲区,通常各占几十KB。

下表展示了不同模型在KPU上的内存占用估算(以8位量化为例):

模型结构 参数量 量化后权重大小 最大中间特征图 是否可部署
LeNet-5(小型CNN) 约60K 约60KB 约50KB ✅ 轻松部署
MobileNet v1(α=0.25) 约0.5M 约500KB 约200KB ✅ 可部署
MobileNet v1(α=0.5) 约1.3M 约1.3MB 约400KB ⚠️ 需优化内存分配
ResNet-18 约11M 约11MB 约1MB ❌ 超出内存限制

从表中可以看出,KPU更适合部署轻量级网络。如果模型过大,需要采取以下优化策略:

  1. 模型剪枝:移除不重要的通道或层,减少参数量。
  2. 通道压缩:将模型的通道数调整为8的倍数,避免填充浪费。
  3. 分块推理:将大模型拆分为多个子模型,依次在KPU上执行,中间结果暂存到外部SRAM。

经验之谈:在实际项目中,建议将模型参数量控制在1MB以内,最大中间特征图控制在500KB以内,这样可以获得较好的推理速度和稳定性。

3.5 实战案例:YOLOv2-Tiny目标检测模型的KPU适配

为了将上述概念串联起来,我们以一个实际案例——在K210上部署YOLOv2-Tiny目标检测模型——来展示KPU核心概念的应用。

背景与约束:YOLOv2-Tiny是一个轻量级目标检测模型,输入尺寸为416×416×3,包含9个卷积层和6个池化层,参数量约15.8M(原始浮点模型)。K210的KPU内存限制为2MB内部SRAM,外部6MB SRAM可用作DMA缓冲。

初始现象:直接使用官方转换工具将YOLOv2-Tiny转换为kmodel时,转换失败,提示“权重大小超出KPU内部SRAM容量”。

分析过程:

  1. 计算量化后权重大小:15.8M参数 × 1字节(8位量化)≈ 15.8MB,远超2MB限制。
  2. 检查中间特征图:最大特征图出现在第6层卷积后,尺寸为13×13×256,占用约43KB,可以接受。
  3. 确认瓶颈在权重存储,而非特征图。

解决方案:

  1. 模型压缩:将YOLOv2-Tiny的通道数减半(每层卷积的滤波器数量减半),参数量降至约4M,量化后约4MB。
  2. 权重分块加载:利用KPU的“权重流式加载”功能,将权重分为多个块,在推理过程中按需从外部Flash加载到内部SRAM。这需要修改kmodel生成时的配置参数。
  3. 输入尺寸调整:将输入尺寸从416×416改为320×320(满足4的倍数要求),进一步减少计算量和内存占用。

结果对比:

指标 原始模型 压缩后模型
参数量 15.8M 4.0M
量化后权重 15.8MB 4.0MB
推理延迟 无法部署 约180ms
mAP(精度) 基准 下降约3.2%

风险与代价:模型压缩导致检测精度下降,特别是对小目标的检测效果变差。此外,权重分块加载增加了Flash读取次数,导致推理延迟从理论上的120ms增加到180ms。

复盘经验:在KPU上部署目标检测模型时,不能简单地将PC端模型直接转换。必须从内存预算出发,反向设计模型结构。建议在训练阶段就考虑KPU的通道对齐要求(输出通道为8的倍数),并使用深度可分离卷积替代标准卷积以降低参数量。

3.6 常见误区与总结复盘

在KPU开发过程中,以下误区最为常见,需要特别注意:

误区一:忽略通道对齐要求。很多开发者从PyTorch导出模型时,输出通道数为16、24等非8倍数,导致KPU推理结果错误。解决方法是使用Netron工具检查模型结构,在转换前手动调整通道数。

误区二:将Softmax放在KPU内执行。KPU不支持Softmax,但很多分类模型在最后一层使用了Softmax。正确做法是让KPU输出原始logits(未经过Softmax的得分),然后在CPU上执行Softmax操作。

误区三:忽视BN层的融合。Batch Normalization在训练时是独立层,但在推理时需要融合到卷积层中。如果转换工具没有自动完成融合,会导致推理结果偏差。

总结复盘:KPU的核心概念可以归纳为“一个引擎、两类存储、三层限制”。一个引擎指卷积运算引擎,两类存储指内部SRAM和外部SRAM,三层限制指通道对齐限制、内存容量限制和算子类型限制。在模型部署前,建议按照以下流程进行自查:

  1. 使用Netron可视化模型结构,列出所有层类型。
  2. 对照KPU支持层列表,标记出不支持的层。
  3. 计算量化后权重大小和最大中间特征图大小。
  4. 检查输入输出通道数是否满足对齐要求。
  5. 根据检查结果决定是调整模型结构还是优化内存分配策略。