如何在Coral Dev Board上基于TensorFlow Lite同时调度CPU、GPU、TPU执行推理负载

首先明确核心逻辑差异：TFLite推理阶段的设备分配规则和Keras训练阶段完全不同，无法直接复用训练时通过层级tf.device()上下文做设备绑定的逻辑，需要按推理运行时的特性调整实现思路。

核心认知前提 #

• TFLite运行时的委托机制本质是按优先级分流算子：单解释器挂载多委托时，高优先级委托先匹配自身能支持的算子，剩余算子自动落到低优先级委托上顺序执行，不属于多设备并行，无法达到同时跑任务降延迟的效果。
• Coral Dev Board上三类设备的算子支持边界：Edge TPU仅支持经过量化编译的专属格式模型做全图加速，不支持单算子级拆分分配；GPU后端基于OpenCL实现，仅兼容部分常见浮点/半量化算子，无法运行TPU专属的 int8 量化算子；CPU端通过XNNPACK实现全算子兼容兜底。

可落地的实操步骤 #

• 按设备能力拆分独立子模型
  • 基于你现有的Keras模型，对照三类设备的算子支持列表做切割：把可适配Edge TPU的量化算子子集单独切出，编译为Edge TPU专属格式的子模型；把TPU不支持但GPU后端兼容的算子单独切出，生成可挂载GPU委托的子模型；剩余不兼容GPU的算子生成CPU兜底的子模型。
  • 切割时要保证相邻子模型的输入输出张量形状严格匹配，在衔接位置提前预留固定内存缓存区，避免跨子模型传递张量时产生额外拷贝开销。
• 为每个子模型初始化对应专属后端的解释器实例
  • TPU子模型：加载Edge TPU编译产物，初始化解释器时挂载edgetpu委托，为该实例分配独立的实时线程优先级，避免被其他低优先级进程抢占算力。
  • GPU子模型：初始化解释器时挂载gpu委托，按板载GPU的内存容量设置独立的内存池，避免和TPU、CPU的资源池抢共享内存带宽。
  • CPU子模型：用XNNPACK委托初始化解释器，把推理线程绑定到空闲的CPU大核上，不要占用小核资源导致调度延迟。
• 独立调度流实现并行推理
  • 不要依赖单解释器的内置线程做并行，直接为每个解释器开独立的工作线程，分别触发对应解释器的invoke()调用。
  • 调度逻辑采用流水线模式：前序子模型完成张量计算后，直接把缓存区的指针传给下游衔接的子模型，不要等所有并行分支全部计算完成再拼接结果，减少线程空等的时间开销。

常见问题规避 #

• 不要为了省事尝试给单解释器同时挂载TPU、GPU、CPU三个委托，这种模式本质是串行分流算子，不仅没法实现多设备并行，还会因为委托切换的额外开销拉高推理延迟。
• 模型拆分时不要做无意义的切割，比如把本来可以在TPU上高效运行的层硬拆分到GPU，跨设备的张量拷贝开销往往会抵消并行带来的收益。
• Coral Dev Board的内存带宽为三类设备共享，并行运行时要合理限制每个解释器的内存池大小，避免触发OOM导致推理卡顿。

内容的提问来源于stack exchange，提问作者Shraddha