3、构建动态形状网络：使用Python API定义动态输入、设置动态维度、动态输出

好，咱们进入正题。这一章我带你手写一个动态形状网络。说白了，就是让TensorRT在推理时能接受不同大小的输入，而不是被绑死在一个固定尺寸上。

我个人习惯把动态形状网络比作「变形金刚」。你给它一个224x224的图，它能跑；给它一个512x512的图，它也能跑。这在实际项目中太重要了——比如你部署一个目标检测模型，用户上传的图片尺寸五花八门，总不能每张都resize到固定大小吧？

3.1 动态输入的定义方式

在Python API里，定义动态输入其实就一行代码的事。但这一行背后，藏着不少门道。

先看最基础的写法：

import tensorrt as trt

# 创建builder和network
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 定义动态输入
input_tensor = network.add_input(
    name="input",
    dtype=trt.float32,
    shape=(-1, 3, -1, -1)  # 关键在这里：-1表示动态维度
)

看到那个 (-1, 3, -1, -1) 了吗？这就是动态形状的核心。第一个-1是batch size，后面两个-1是高度和宽度。中间的3是通道数，固定不变。

我在项目中遇到过一个问题：有人把所有的维度都设成-1，结果引擎构建失败。为什么？因为TensorRT需要至少一个维度是确定的，才能做内存规划和算子选择。你想想看，如果连通道数都是动态的，那卷积核怎么选？

3.2 设置动态维度的范围

光写-1还不够。你得告诉TensorRT：「这个动态维度，最小是多少，最大是多少，常见是多少。」

这就用到 OptimizationProfile 了：

# 创建优化配置
profile = builder.create_optimization_profile()

# 设置输入"input"的三个形状
profile.set_shape(
    "input",
    min=(1, 3, 224, 224),   # 最小尺寸：1张224x224的图
    opt=(4, 3, 512, 512),   # 常见尺寸：4张512x512的图
    max=(8, 3, 1024, 1024)  # 最大尺寸：8张1024x1024的图
)

# 把profile绑定到config
config = builder.create_builder_config()
config.add_optimization_profile(profile)

嗯，这里要注意三个值的含义：

• min：推理时不会小于这个尺寸。我一般设成实际业务中最小的输入。
• opt：TensorRT会针对这个尺寸做极致优化。选你最常用的尺寸。
• max：不能超过这个尺寸，否则会报错。

3.3 动态输出的自动推导

输入设好了，输出怎么办？好消息是：大部分情况下，TensorRT会自动推导输出的形状。

比如一个简单的卷积网络：

# 添加一个卷积层
conv = network.add_convolution(
    input=input_tensor,
    num_output_maps=64,
    kernel_shape=(3, 3),
    kernel=weights,
    bias=bias
)

# 输出形状会自动变成 (-1, 64, -1, -1)
# 因为卷积不改变空间尺寸（padding=1的情况下）
output = conv.get_output(0)
output.name = "output"
network.mark_output(output)

但有些层会导致输出形状变化，比如池化层、步长卷积、上采样等。这时候你需要手动确认一下：

# 添加一个步长为2的卷积
conv_stride = network.add_convolution(
    input=input_tensor,
    num_output_maps=128,
    kernel_shape=(3, 3),
    stride=(2, 2),
    kernel=weights,
    bias=bias
)

# 输出形状会变成 (-1, 128, -1, -1)
# 但实际尺寸是输入的一半，TensorRT会自动处理

3.4 完整示例：构建一个动态CNN

来，咱们把上面这些串起来，写一个完整的动态形状网络构建流程：

import tensorrt as trt
import numpy as np

def build_dynamic_network():
    # 1. 创建builder和network
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    
    # 2. 定义动态输入
    input_tensor = network.add_input(
        name="input",
        dtype=trt.float32,
        shape=(-1, 3, -1, -1)
    )
    
    # 3. 构建网络结构（这里用简单的卷积+池化）
    # 第一层卷积
    conv1 = network.add_convolution(
        input=input_tensor,
        num_output_maps=32,
        kernel_shape=(3, 3),
        padding=(1, 1),
        kernel=np.random.randn(32, 3, 3, 3).astype(np.float32),
        bias=np.zeros(32, dtype=np.float32)
    )
    conv1.name = "conv1"
    
    # ReLU激活
    relu1 = network.add_activation(
        input=conv1.get_output(0),
        type=trt.ActivationType.RELU
    )
    
    # 池化层
    pool1 = network.add_pooling(
        input=relu1.get_output(0),
        type=trt.PoolingType.MAX,
        window_size=(2, 2),
        stride=(2, 2)
    )
    
    # 4. 标记输出
    output = pool1.get_output(0)
    output.name = "output"
    network.mark_output(output)
    
    # 5. 设置动态profile
    config = builder.create_builder_config()
    profile = builder.create_optimization_profile()
    profile.set_shape(
        "input",
        min=(1, 3, 224, 224),
        opt=(4, 3, 512, 512),
        max=(8, 3, 1024, 1024)
    )
    config.add_optimization_profile(profile)
    
    # 6. 构建引擎
    engine = builder.build_engine(network, config)
    
    return engine

# 调用
engine = build_dynamic_network()
print("动态形状网络构建成功！")

这段代码跑通后，你就有了一个能处理224到1024任意尺寸输入的引擎。是不是很爽？

3.5 动态形状的注意事项

最后，我总结几个实战中容易翻车的地方：

问题	原因	解决方案
构建引擎超慢	动态形状需要探索更多算子组合	缩小min-max范围，或优化opt值
推理时形状不匹配	输入尺寸超出profile范围	加预处理检查，或resize到合法范围
显存占用飙升	TensorRT为最大尺寸预留了显存	使用动态显存分配，或限制max尺寸
某些层不支持动态	比如全连接层需要固定输入维度	改用卷积或GAP替代

下一章，我会带你看看如何在推理时动态绑定输入数据，以及如何处理多个profile的情况。到时候你会发现，动态形状的真正威力，在推理阶段才完全释放出来。