TensorFlow Model Garden：面向 Google Edge TPU（Pixel 6）的 NAS 优化模型实战指南

本文基于 TensorFlow Model Garden 仓库中的 official/projects/edgetpu 项目，系统讲解针对 Pixel 6 内置 Google Tensor SoC 中 Edge TPU 加速器优化的视觉与 NLP 模型家族：包括 MobileNetEdgeTPUV2 图像分类、DeepLabV3+/Autoseg 语义分割以及 MobileBERT-EdgeTPU 语言模型的架构设计、训练命令、TFLite 量化导出参数与真机基准测试方法。读完本文后，你可以完整复现“预注册实验配置 → 训练 → 后训练量化导出 → Pixel 6 真机 NNAPI 基准测试”的全链路，并理解每个关键命令行参数背后的源码实现。

一、项目定位与环境要求

该目录下的模型均针对 Pixel 6 SoC（Google Tensor）中的 Edge TPU 进行优化。项目使用**神经架构搜索（Neural Architecture Search, NAS）**自动化模型设计流程：搜索算法的目标是在网络结构的巨大空间中，联合最小化 TPU 延迟（以及能耗）并最大化精度，从而为不同延迟目标生成一族模型。这些模型通过 TensorFlow Model Garden 与 TensorFlow Hub 公开发布，供开发者在 Pixel 6 上构建设备端应用。

根据 README 的声明，运行环境要求为：

• TensorFlow 2.4
• Python 3.7

整个项目分为两大方向，仓库目录与文档一一对应：

方向	目录	说明文档
视觉（分类/分割/检测）	official/projects/edgetpu/vision/	vision/README.md
自然语言理解（MobileBERT-EdgeTPU）	official/projects/edgetpu/nlp/	nlp/README.md

其中对象检测相关的 EdgeTPU 模型托管在仓库的 research/object_detection/ 目录中，视觉 README 中亦有对应指引。

二、视觉方向：MobileNetEdgeTPUV2 与 NAS 搜索空间

2.1 核心设计：基于分组卷积的倒瓶颈块（GC-IBN）

MobileNetEdgeTPUV2 的搜索空间围绕一个关键构件展开——分组卷积倒瓶颈块（Group Convolution based Inverted Bottleneck, GC-IBN）。倒瓶颈块（IBN）是移动端视觉网络的经典构件：常规 IBN 在深度卷积前后使用逐点卷积做扩展/投影；而全卷积替换逐点扩展虽然参数量更大，但面对视觉网络后段的窄而深的张量时，延迟与内存代价过高，限制了"融合式"全卷积 IBN 在整个模型中的应用。

GC-IBN 用分组卷积（GC）作为融合扩展的一部分，通过 group size 参数在"常规深度卷积 IBN"与"融合全卷积 IBN"之间打开一个连续的设计空间：可以逐步增加可训练参数量而不承担全卷积 IBN 的延迟代价，同时仍可能比常规 IBN 更快。这一权衡高度依赖张量形状、无法手工泛化到整个网络，因此项目使用 AutoML 技术逐层搜索最优块决策，针对不同延迟目标产出一族网络。

2.2 分类模型族：Int8 量化基准（224x224 输入）

以下为 Pixel 6 Edge TPU 上的真机基准测试结果（int8 量化模型，来自 vision/README.md）：

模型	精度 (int8)	Pixel 6 Edge TPU 延迟 (ms)
MobileNetEdgeTPUv2-Tiny	74.66%	0.78
MobileNetEdgeTPUv2-XS	75.79%	0.82
MobileNetEdgeTPUv2-S	77.36%	1.03
MobileNetEdgeTPUv2-M	78.43%	1.35
MobileNetEdgeTPUv2-L	79.00%	1.64
MobileNetEdgeTPU dm1.0 (V1)	75.6%	0.92
MobileNetEdgeTPU dm1.25 (V1)	77.06%	1.20
MobileNetEdgeTPU dm1.5 (V1)	75.9%	1.42
MobileNetEdgeTPU dm1.75 (V1)	78.6%	1.93

可以看到 V2 家族在同精度区间内全面低于 V1（dm 为 depth multiplier）的延迟，改善了设备端模型的延迟-精度帕累托前沿。

2.3 模型在源码中的注册与构建

从源码结构看，所有 V2 变体与 Autoseg 主干都注册在 mobilenet_edgetpu_v2_model.py 的 MODEL_CONFIGS 字典中：mobilenet_edgetpu_v2、mobilenet_edgetpu_v2_tiny、mobilenet_edgetpu_v2_xs、mobilenet_edgetpu_v2_s、mobilenet_edgetpu_v2_m、mobilenet_edgetpu_v2_l 以及 autoseg_edgetpu_backbone_xs/s/m。MobilenetEdgeTPUV2 是一个 Keras 模型包装类：

• __init__ 接收 model_config_name 与 overrides，通过 MODEL_CONFIGS[model_config_name] 取出块级配置并做 replace(**overrides)，再调用 mobilenet_edgetpu_v2() 函数式构建网络；
• from_name(model_name, model_weights_path, checkpoint_format, overrides) 类方法支持按预定义名称构建模型并加载权重（tf_checkpoint 或 keras_checkpoint 两种格式），还支持通过 overrides['model_config'] 注入自定义配置；
• 当配置 features_as_dict=True 时输出为字典（供 ASPP 解码器使用），否则输出特征列表（供分割任务使用）。

2.4 实验配置：预注册的实验工厂

训练入口使用的实验配置在 mobilenet_edgetpu_config.py 中通过 @exp_factory.register_config_factory 注册，每个工厂名即训练/导出命令中的 EXPERIMENT_TYPE（或 --model_name）取值，包括：

• V2 分类：mobilenet_edgetpu_v2_tiny / _xs / _s / _m / _l
• V1 分类：mobilenet_edgetpu、mobilenet_edgetpu_dm0p75 / _dm1p25 / _dm1p5 / _dm1p75
• 搜索模型：mobilenet_edgetpu_search（需额外通过 --params_override=task.saved_model_path='...' 指定 saved model）
• Autoseg 主干：autoseg_edgetpu_backbone_xs / _s / _m

其中基础实验配置 mobilenet_edgetpu_base_experiment_config 的关键超参（全部为 ImageNet 分类场景）：

配置项	取值	说明
train_batch_size	4096	全局训练 batch
输入路径	imagenet-2012-tfrecord/train*、valid*	TFRecord 格式 ImageNet
输入尺寸	224 x 224 x 3，1001 类	num_classes=1001
数据增强	AutoAugment（aug_type='autoaug'）	训练集启用
精度	dtype='bfloat16'	训练/评估均使用 bfloat16
优化器	RMSProp（rho=0.9, momentum=0.9, epsilon=0.001）+ EMA（average_decay=0.99）
学习率	exponential：0.008 * (batch_size // 128)，decay_steps=2.4 个 epoch，decay_rate=0.97，阶梯衰减
预热	线性预热 5 个 epoch，从 0 开始
总步数	550 * steps_per_epoch	steps_per_epoch = 1281167 // 4096

三、训练流程：train.py 命令行实操

视觉训练入口为 vision/train.py。分类与分割的通用训练命令如下（EXPERIMENT_TYPE 必须取上表中的预注册名，如分类用 mobilenet_edgetpu_v2_xs，分割用 seg_deeplabv3plus_mobilenet_edgetpuv2_s_ade20k、autoseg_edgetpu_xs）：

EXPERIMENT_NAME=xxx  # 例如 'mobilenet-edgetpu-test-run'
EXPERIMENT_TYPE=xxx  # 例如 'mobilenet_edgetpu_v2_xs'
python3 train.py \
  --experiment_name=${EXPERIMENT_NAME} \
  --experiment_type=${EXPERIMENT_TYPE} \
  --mode=train_and_eval

从源码看，train.py 的 main() 完成以下工作：

1. gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_params) 加载可选的 gin 配置；
2. train_utils.parse_configuration(FLAGS) 按 experiment_type 解析出 ExperimentConfig（即第二节注册的工厂产物）；
3. 若 mode 含 train，用 train_utils.serialize_config 把配置序列化写入 model_dir（纯 eval 模式跳过，避免连续 eval 作业与训练作业写文件竞争）；
4. 若配置了 params.runtime.mixed_precision_dtype，调用 performance.set_mixed_precision_policy 设置混合精度；
5. distribute_utils.get_distribution_strategy(...) 依据 distribution_strategy/num_gpus/tpu_address 构建分发策略，在其作用域内由 task_factory.get_task(params.task) 构造任务（分类任务实现于 tasks/image_classification.py，分割任务实现于 tasks/semantic_segmentation.py）；
6. 最终交给 train_lib.run_experiment(...) 执行训练/评估循环。

分割模型的分割头配置分别由 semantic_segmentation_config.py（DeepLabV3+ 组合）与 semantic_segmentation_searched_config.py（Autoseg 搜索模型）注册。

四、语义分割：DeepLabV3+ 主干与 Autoseg-EdgeTPU

4.1 架构

分割模型分两条路线：

• DeepLabV3+ 路线：以预训练 MobileNetEdgeTPUV2 为特征编码器，配合 ASPP 解码器与 DeepLabV3+ 头，并使用可分离卷积压缩模型规模；
• Autoseg-EdgeTPU 路线：用强化学习算法联合搜索特征提取器与分割头。特征提取器来自 Edge TPU 搜索空间（IBN 与融合 IBN 混合），逐层自动寻找最优 kernel 大小、通道数、扩展率与分组数；分割头是 BiFPN 头的优化版本，带自定义重复次数与特征选择。

4.2 Argmax 融合：消除尾部延迟

模型最后两级（双线性 resize 与 Argmax）在设备上贡献了显著延迟，原因是其间激活尺寸巨大（512 x 512 x 类别数）。由于可以先缩小 Argmax 的计算尺寸、再用最近邻把类别图放大到目标分辨率而不明显影响质量，这两个算子可以融合。对应的导出参数见下一节 --finalize_method。

4.3 分割模型基准（ADE20K 32 类）

默认 Argmax（来自 vision/README.md）：

模型	主干	分割头	参数量 (M)	ADE20K 32 类 mIOU	Pixel 6 EdgeTPU 延迟 (ms)
deeplabv3plus_mobilenet_edgetpuv2_baseline	MobileNetV2 (baseline)	DeepLabV3+	2.34	54.06%	7.5
deeplabv3plus_mobilenet_edgetpuv2_xs	MobileNetEdgeTPUV2-XS	DeepLabV3+	3.6	56.02%	5.2
deeplabv3plus_mobilenet_edgetpuv2_s	MobileNetEdgeTPUV2-S	DeepLabV3+	5.2	59.43%	5.9
deeplabv3plus_mobilenet_edgetpuv2_m	MobileNetEdgeTPUV2-M	DeepLabV3+	7.7	59.81%	7.2
autoseg_edgetpu_xs	AutosegEdgeTPU-XS	BiFPN	2.9	59.64%	5.4
autoseg_edgetpu_s	AutosegEdgeTPU-S	BiFPN	3.1	61.31%	5.7

启用 Argmax 融合后（同一 checkpoint，仅图结构不同，无参数变化）：

模型	ADE20K 32 类 mIOU	延迟 (ms)
deeplabv3plus_mobilenet_edgetpuv2_xs	56%	3.4
deeplabv3plus_mobilenet_edgetpuv2_s	59.41%	4.2
deeplabv3plus_mobilenet_edgetpuv2_m	59.79%	5.5
autoseg_edgetpu_xs	59.62%	3.6
autoseg_edgetpu_s	61.28%	3.9

AutosegEdgeTPU-XS/S 以更少的参数（2.9M/3.1M）显著超过 MobileNetV2 基线（2.34M，54.06% mIOU），并在融合 Argmax 后延迟降到 3.4~5.5 ms。

五、从训练到量化部署：serving/export_tflite.py

导出入口为 vision/serving/export_tflite.py。它从 checkpoint 构建实验模型、加载权重、按 finalize_method 追加服务化层，然后用 tf.lite.TFLiteConverter 转换并（可选）做后训练量化，最后同时输出 .tflite 与 SavedModel。

5.1 分类模型导出

python3 serving/export_tflite.py \
  --model_name=${EXPERIMENT_TYPE} \
  --ckpt_path=${CHECKPOINT} \
  --dataset_dir=/path/to/calibration/dataset \
  --output_dir=/tmp \
  --quantize \
  --image_size=224

--model_name 必须是预注册的分类配置名（如 mobilenet_edgetpu_v2_xs）。

5.2 分割模型导出

python3 serving/export_tflite.py \
  --model_name=${EXPERIMENT_TYPE} \
  --ckpt_path=${CHECKPOINT} \
  --dataset_dir=/path/to/calibration/dataset \
  --output_dir=/tmp \
  --quantize \
  --quantize_less_restrictive \
  --image_size=512 \
  --finalize_method=${ARGMAX_FUSION}

--model_name 需为预注册分割配置（如 deeplabv3plus_mobilenet_edgetpuv2_s_ade20k_32）。--finalize_method 控制 Argmax 融合策略，取值：

• resize512,argmax：先双线性放大到 512x512 再做 Argmax；
• resize256,argmax,resize512,squeeze：缩到 256x256 → Argmax → 最近邻放大到 512x512 → 去掉维度 1；
• resize128,argmax,resize512,squeeze：缩到 128x128 → Argmax → 最近邻放大到 512x512 → 去掉维度 1。

从源码看，finalize_method 的解析为 FLAGS.finalize_method.lower().split(',')，支持的语法是 none|(argmax|resize<?>)[,...]，其中 resize<?> 表示"双线性或最近邻 resize 到指定尺寸"，squeeze 去掉 size-1 维度，例如 resize128,argmax,resize512,squeeze 即"双线性缩到 128x128 → argmax → 最近邻放大 512x512"。

5.3 完整参数速查

参数	默认值	说明
--model_name	必填	实验配置工厂名，用于构建模型
--ckpt_path	无	checkpoint 路径；不提供则导出随机权重
--ckpt_format	tf_checkpoint	tf_checkpoint（tf.train.Checkpoint）或 keras_checkpoint
--output_dir	无	导出输出目录
--image_size	224	输入图像边长，应与训练配置一致
--fix_batch_size	True	True 时导出固定 batch=1 的模型
--output_layer	None	从中间层取输出（如 probs）；未找到时打印模型 summary 辅助选择
--export_keras_model	False	True 用 tf_keras.Model 保存 SavedModel，否则用 tf.saved_model API
--quantize	False	导出 TFLite 前做量化（SavedModel 不量化）
--quantize_less_restrictive	False	允许非 int8 中间类型、自动选择模型输出类型
--use_experimental_quantizer	True	启用 TFLiteConverter 2.0 实验性量化器
--num_calibration_steps	100	PTQ 校准步数
--dataset_name	imagenet2012	量化校准数据集名
--dataset_dir	无	校准数据集位置
--dataset_split	train	校准所用数据 split

导出流程的调用链为：get_export_config_from_flags() 组装 ExportConfig → export_util.build_experiment_model(experiment_type) 按实验名建图 → common_modules.load_weights 加载权重 → 构造 model_for_inference 并执行 finalize_serving 追加服务层 → configure_tflite_converter 注入量化配置后 convert()。相关工具实现在 serving/export_util.py，并有配套测试 serving/export_tflite_test.py。

5.4 Pixel 6 真机基准测试

导出的模型兼容 NNAPI，可用 TFLite benchmark 工具在 Pixel 6 上压测 Edge TPU 延迟：

bazel build -c opt --config=android_arm64 tensorflow/lite/tools/benchmark:benchmark_model
# 推送到设备
adb push bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model /data/local/tmp
adb push /path/to/model.tflite /data/local/tmp/
# 设备端运行基准（启用 NNAPI 并指定 Edge TPU 加速器）
adb shell /data/local/tmp/benchmark_model --graph=/data/local/tmp/model.tflite \
  --use_nnapi=true --nnapi_accelerator_name=google-edgetpu

关键点：--use_nnapi=true 且 nnapi_accelerator 指定为 google-edgetpu。视觉 README 同时给出推理可视化的 notebook（serving/inference_visualization_tool.ipynb）与 ImageNet TFLite 评估脚本（serving/tflite_imagenet_evaluator.py）。

六、NLP 方向：MobileBERT-EdgeTPU

6.1 设计动机与量化结论

MobileBERT-EdgeTPU 以 MobileBERT 为搜索空间基础，通过 AutoML 找到硬件利用率最高提升 2 倍的架构，使更大、更准的模型能跑上芯片且延迟仍优于基线 MobileBERT。项目构建了定制蒸馏训练流水线并做了详尽的超参搜索（学习率、dropout 等）。一个重要的工程结论是：与多数视觉模型不同，MobileBERT 若直接做普通后训练量化（PTQ）或量化感知训练（QAT），精度会显著掉点——需要裁剪 mask 值等针对性修改才能保住量化精度。因此项目同时提供一组 Edge TPU 友好的 float 模型，其中 float 版 MobileBERT-EdgeTPU-M 的精度甚至逼近 1.3GB 的 BERT_large，量化从"必需"变成"可选优化"。

6.2 预训练模型一览

SQUAD 分数使用 SQuAD V1.1 数据集、通过 BertSpanLabeler 任务头测得（nlp/README.md）：

模型	参数量	MLM	SQUAD (float)	SQUAD (int8)
MobileBERT (baseline)	24.6M	71.4%	89.02%	87.95%
MobileBERT-EdgeTPU-XS	27.1M	71.2%	88.20%	87.15%
MobileBERT-EdgeTPU-S	38.3M	72.8%	89.97%	89.40%
MobileBERT-EdgeTPU-M	50.9M	73.8%	90.24%	89.50%

注意：float 基线 MobileBERT 在 NNAPI 下部分算子被委托给 CPU，延迟明显更高；BERT_base/BERT_large 精度引自训练论文，模型过大不可行于设备端。

6.3 从 checkpoint 恢复模型

按 README 的示例（或参考 nlp/serving/export_tflite_squad.py）：

import tensorflow as tf
from official.nlp.projects.mobilebert_edgetpu import params

bert_config_file = ...
model_checkpoint_path = ...

# 建立实验参数并加载配置
experiment_params = params.EdgeTPUBERTCustomParams()

# 将输入 mask 类型改为 tf.float32，避免额外的 cast 算子
experiment_params.student_model.encoder.mobilebert.input_mask_dtype = 'float32'
pretrainer_model = model_builder.build_bert_pretrainer(
    experiment_params.student_model,
    name='pretrainer',
    quantization_friendly=True)

checkpoint_dict = {'model': pretrainer_model}
checkpoint = tf.train.Checkpoint(**checkpoint_dict)
checkpoint.restore(FLAGS.model_checkpoint).assert_existing_objects_matched()

关键细节：input_mask_dtype='float32' 是 Edge TPU 友好化改造的一部分（消除额外 cast）；quantization_friendly=True 启用量化友好构建。参数结构定义在 nlp/configs/params.py 的 EdgeTPUBERTCustomParams 中，包含教师/学生模型配置、层间蒸馏（LayerWiseDistillationParams：hidden 蒸馏因子 100.0、beta 5000.0、gamma 5.0、attention 1.0）与端到端蒸馏（EndToEndDistillationParams：58 万步、2 万步预热、distill_ground_truth_ratio=0.5）两套流水线配置。

6.4 使用 TF-Hub 模型做下游微调

import tensorflow as tf
import tensorflow_hub as hub
from official.nlp.modeling import models

encoder_network = hub.KerasLayer(
    'https://tfhub.dev/google/edgetpu/nlp/mobilebert-edgetpu/s/1',
    trainable=True)
squad_model = models.BertSpanLabeler(
    network=encoder_network,
    initializer=tf.keras.initializers.TruncatedNormal(stddev=0.01))

6.5 蒸馏训练入口

预训练入口 nlp/run_mobilebert_edgetpu_train.py 的流程从源码看为：EdgeTPUBERTCustomParams + utils.config_override 解析命令行参数 → 构建分发策略 → 分别构建 teacher（quantization_friendly=False）与 student（quantization_friendly=True）pretrainer，教师 checkpoint 必填、学生 checkpoint 可选（缺失时仅告警）→ 交给 MobileBERTEdgeTPUDistillationTrainer（nlp/mobilebert_edgetpu_trainer.py）驱动 Orbit Controller 训练，checkpoint 每 2 万步由 CheckpointManager 保存（max_to_keep=5）以支持抢占恢复。下游任务（如 GLUE-MNLI、SQuAD）的 gin 实验配置位于 nlp/experiments/ 目录，含 mobilebert_baseline.yaml、mobilebert_edgetpu_{xxs,xs,s,m}.yaml 及 downstream_tasks/ 子目录。

七、小结

official/projects/edgetpu 展示了一条完整的"为特定硬件加速器做模型共同设计"的工程路线：用 NAS 打开由 GC-IBN、逐层超参构成的搜索空间，为不同延迟目标产出 Tiny/XS/S/M/L 模型族；用预注册的实验工厂（register_config_factory）统一训练与导出的配置入口；用 TFLite 导出脚本统一后训练量化与 Argmax 融合等服务化变换；最终以 NNAPI benchmark 在 Pixel 6 真机闭环验证。视觉侧覆盖分类（0.78~1.64 ms）与分割（融合后 3.4~5.5 ms），NLP 侧则以 MobileBERT-EdgeTPU 的 S/M 模型在 SQuAD 上超过 float BERT_base 精度，为设备端 Transformer 提供了量化可选的部署路径。