4、模型准备：PyTorch/TensorFlow/ONNX模型导出动态维度的方法、动态轴标记、batch维度设置

在RKNN模型转换与部署的完整链路中，模型导出是决定后续量化精度与推理灵活性的关键前置环节。许多开发者在完成模型训练后，直接使用静态shape导出模型，导致在部署阶段遇到batch size固定、输入分辨率不可变、多路视频流无法复用同一模型等问题。本章将系统讲解在PyTorch、TensorFlow、ONNX三种主流框架下导出动态维度模型的具体方法，重点剖析动态轴标记（dynamic axes）的语义、batch维度设置的工程陷阱，以及如何为RKNN工具链准备一份“既灵活又可控”的中间模型。

上图展示了从三种主流框架导出动态维度模型，到ONNX中间表示，再到RKNN工具链适配的完整链路。核心要点在于：动态维度并非“无限自由”，而是通过显式标记告诉推理引擎哪些轴可以在运行时变化。RKNN工具链在加载模型时，会依据这些标记生成对应的动态shape处理逻辑，但实际可变的轴范围和步长受限于NPU硬件调度策略。因此，理解动态轴的声明方式与RKNN的约束条件，是避免“导出成功但转换失败”的关键。

4.1 核心原理：动态维度与静态维度的本质差异

深度学习模型在编译与部署时，张量的shape信息决定了内存分配、算子调度和流水线优化策略。静态shape模型在编译期即可确定所有中间张量的大小，NPU可以提前规划好缓冲区并应用最激进的优化；而动态shape模型则需要在运行时根据实际输入尺寸重新计算部分调度策略，这会引入额外的开销。但动态维度的价值在于：同一份模型可以适配不同batch size、不同分辨率的输入，无需为每种规格单独导出和存储模型文件，这在多路视频分析、动态分辨率检测等场景中能显著减少开发与维护成本。

从ONNX的视角看，动态维度通过dim_param（维度参数名）或dim_value（固定值）来区分。当某个维度被赋予一个字符串名称（如"N"、"batch"）时，该维度即为动态轴；若使用整数，则为静态轴。RKNN-Toolkit2在解析ONNX模型时，会读取这些标记，并在内部构建动态shape的推理图。值得注意的是，动态轴并非所有算子都支持，例如某些硬件加速的卷积实现可能要求输入通道数为固定值，而Reshape、Transpose等算子对动态shape的传播也有额外限制。

4.2 PyTorch模型导出动态维度的方法

PyTorch通过torch.onnx.export函数导出ONNX模型时，dynamic_axes参数是控制动态维度的核心。该参数接受一个字典，键为输入/输出名称，值为一个子字典，子字典的键为维度索引，值为该维度的符号名称。以下是一个典型的动态batch导出示例：

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

# 准备一个静态shape的dummy输入，仅用于确定网络结构
dummy_input = torch.randn(1, 3, 224, 224)

# 定义动态轴：batch维度（索引0）标记为'N'，高度和宽度维度（索引2,3）标记为'H'和'W'
dynamic_axes = {
    'input': {0: 'N', 2: 'H', 3: 'W'},
    'output': {0: 'N'}
}

torch.onnx.export(
    model,
    dummy_input,
    'resnet18_dynamic.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes,
    opset_version=12,
    do_constant_folding=True
)
print("动态模型导出完成")

在上述代码中，dynamic_axes将输入张量的第0维（batch）、第2维（高度）、第3维（宽度）均声明为动态轴。这意味着导出的ONNX模型可以接受任意batch size和任意H×W分辨率的输入（只要满足算子的约束）。但需要注意：opset_version的选择会影响动态shape的表达能力，建议使用opset 11及以上版本，因为早期opset对动态shape的支持不完整，可能导致某些算子无法正确传播动态维度。

对于更复杂的模型，如果存在多个输入或中间张量需要动态化，可以在dynamic_axes中为每个张量单独声明。此外，PyTorch 2.x版本引入了torch.export和torch.onnx.dynamo_export等新接口，它们对动态shape的支持更加完善，但API略有不同。以下是一个使用dynamo_export的示例：

from torch.onnx import dynamo_export

# 使用动态shape的输入，通过dynamic_shapes参数指定
dynamic_shapes = {
    'input': {0: torch.export.Dim("batch", min=1, max=16), 2: torch.export.Dim("height"), 3: torch.export.Dim("width")}
}
dynamo_export(model, dummy_input, dynamic_shapes=dynamic_shapes).save("resnet18_dynamo.onnx")

这里torch.export.Dim允许指定动态维度的范围（如batch从1到16），这为后续RKNN转换提供了更明确的约束信息，有助于工具链生成更高效的调度策略。

4.3 TensorFlow模型导出动态维度的方法

TensorFlow中动态维度的处理方式与PyTorch不同。在Keras或tf.keras中，定义模型输入时使用None作为维度占位符即可表示动态轴。例如：

import tensorflow as tf

# 定义动态输入：batch维度为None，高度和宽度也为None
inputs = tf.keras.Input(shape=(None, None, 3), batch_size=None, name='input')
# 构建模型...
model = tf.keras.Model(inputs, outputs)

# 保存为SavedModel格式
model.save('dynamic_model_savedmodel')

在转换为TFLite或ONNX时，需要额外指定动态维度。对于TFLite，转换器默认会保留动态维度，但某些优化可能会将其固定。推荐使用以下方式：

converter = tf.lite.TFLiteConverter.from_saved_model('dynamic_model_savedmodel')
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 允许动态shape
converter._experimental_disable_per_channel = False
tflite_model = converter.convert()

如果希望将TensorFlow模型导出为ONNX，推荐使用tf2onnx工具。在转换时，需要显式指定动态维度：

import tf2onnx

# 指定输入签名，将batch维度标记为动态
spec = [tf.TensorSpec([None, None, None, 3], tf.float32, name='input')]
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13, output_path="model.onnx")

在tf.TensorSpec中，None即表示该维度为动态。tf2onnx会自动将这些None转换为ONNX的dim_param。需要注意的是，TensorFlow的None维度在转换后可能被命名为None或自动生成名称，建议在转换后使用onnx库检查并重命名动态轴，以便后续RKNN配置时清晰识别。

4.4 ONNX模型动态轴标记与手动修改

有时我们拿到的ONNX模型已经是静态shape，或者需要手动调整动态轴的名称和范围。此时可以使用onnx库直接修改模型文件。以下是一个典型的操作流程：

import onnx

model = onnx.load("static_model.onnx")

# 获取输入张量
input_tensor = model.graph.input[0]
print("原始shape:", input_tensor.type.tensor_type.shape)

# 修改batch维度为动态，命名为'N'
dim = input_tensor.type.tensor_type.shape.dim[0]
dim.dim_param = 'N'  # 将固定值改为动态参数名

# 修改高度和宽度维度为动态
dim_h = input_tensor.type.tensor_type.shape.dim[2]
dim_h.dim_param = 'H'
dim_w = input_tensor.type.tensor_type.shape.dim[3]
dim_w.dim_param = 'W'

# 检查模型合法性
onnx.checker.check_model(model)
onnx.save(model, "dynamic_model_modified.onnx")
print("动态轴标记完成")

除了修改输入张量，还需要检查输出张量以及中间可能存在的Shape相关算子。如果模型中有Reshape、Flatten等操作，它们可能会将动态维度传播到其他位置，需要确保这些算子支持动态shape。此外，onnxruntime提供了动态shape的推理验证功能，可以在导出后立即测试：

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("dynamic_model_modified.onnx", providers=['CPUExecutionProvider'])
# 使用不同的batch size进行推理
for batch in [1, 4, 8]:
    dummy = np.random.randn(batch, 3, 224, 224).astype(np.float32)
    outputs = sess.run(None, {'input': dummy})
    print(f"batch={batch}, output shape={outputs[0].shape}")

如果所有batch size均能正常推理且输出shape正确，说明动态轴标记成功。若出现错误，通常是因为某些算子不支持动态shape，需要回到原始框架中修改模型结构或使用onnx-simplifier等工具进行优化。

4.5 案例拆解：动态batch模型在RK3588上的部署实战

某安防企业需要在RK3588平台上部署一个行人检测模型，要求同一模型能够处理1路、4路、8路视频流，且输入分辨率在960×540到1920×1080之间动态变化。初始方案是分别导出三个静态模型（batch=1、4、8），但这样会导致模型文件冗余且切换时需要重新加载，延迟较高。

关键约束与目标：模型为YOLOv5s，PyTorch训练，需转换为RKNN格式。目标是在不损失精度的前提下，实现动态batch和动态分辨率支持，且单次推理延迟不超过15ms（NPU）。

实施步骤：

1. 导出动态ONNX模型：使用PyTorch的dynamic_axes将batch、高度、宽度均设为动态轴，opset=12。
2. 验证动态性：使用onnxruntime测试batch=1、4、8以及不同分辨率下的推理，确认输出shape正确。
3. RKNN转换配置：在rknn.config中设置target_platform='rk3588'，并在load_onnx时指定input_size_list为[[1,3,640,640]]作为基准shape。
4. 动态shape范围设置：在RKNN-Toolkit2中，通过rknn.build(do_quantization=True, dataset='dataset.txt')时，工具会自动识别ONNX中的动态轴。但需要额外设置dynamic_input参数，指定batch的可选范围（如[1,4,8]）和分辨率范围。
5. 板端部署：在RK3588上使用rknn_init时传入动态shape配置，运行时根据实际输入调用rknn_run并传入对应的shape信息。

结果对比：动态模型在batch=1时推理延迟为8ms，batch=4时为12ms，batch=8时为18ms（略超目标，但通过NPU调度优化后降至14ms）。相比静态多模型方案，模型文件从3个减少到1个，内存占用降低约40%，切换视频流时无需重新加载模型，延迟降低90%以上。

风险与代价：动态shape会导致NPU无法完全预分配缓冲区，首次推理时会有额外的内存分配开销。此外，某些算子（如YOLOv5的Detect层中的Reshape）在动态shape下可能触发额外的数据拷贝，需通过rknn.config中的optimization_level参数进行权衡。

4.6 常见误区与总结复盘

总结本章，动态维度模型导出的核心在于“声明”与“验证”两个环节。声明阶段需要根据部署场景明确哪些轴必须可变，并选择合适的框架API进行标记；验证阶段则需通过onnxruntime或RKNN模拟器进行多shape测试，确保模型在目标范围内稳定运行。对于RKNN部署，建议在导出时就将动态轴的范围限制在硬件友好的数值（如batch为1、4、8，分辨率为16的倍数），这能显著降低NPU调度复杂度，提升实际推理效率。