如何在Python中使用PyTorch 2.0的动态形状特性进行推理？-Python教程-PHP中文网

pytorch 2.0 中动态形状需通过 torch.export.export 显式声明 dynamic_shapes（配合 dim 对象），而非 torch.compile 直接启用 dynamic=true；仅此方式才能安全、高效支持变长输入。

如何在python中使用pytorch 2.0的动态形状特性进行推理？

PyTorch 2.0 的 `torch.compile` 默认不支持动态形状

直接调用 torch.compile(model) 后传入不同尺寸的输入（比如 batch=1 和 batch=4 交替推理），会触发重新编译或报错 RuntimeError: Input tensor shapes changed。这不是 bug，而是默认模式（mode="default"）下，编译器将输入 shape 视为静态常量参与图优化。

真正支持动态 shape 的路径只有一条：显式启用 dynamic=True，且必须配合 torch.export 或 torch._dynamo.config 级别控制——但注意，torch.compile 本身不接受 dynamic 参数，它靠后端（如 aot_eager、inductor）和前端配置协同生效。

仅当使用 torch.export.export 导出带 dynamic_shapes 描述的模型时，生成的 ExportedProgram 才能安全接收变长输入
torch.compile + dynamic=True 实际是通过 torch._dynamo.config.dynamic_shapes = True 全局开关 + torch.compile(..., backend="inductor") 触发的，不是函数参数
Inductor 后端对 dynamic shape 的支持仍有限制：目前仅支持维度值为“symbolic”（如 s0）且满足单调约束（如 batch size 可变，但 H/W 不可同时任意变）

用 `torch.export.export` 声明动态维度才是可靠方式

这是 PyTorch 2.0 推荐的、生产级支持动态 shape 的标准流程。核心是用 torch.export.Dim 显式标记哪些维度可变，并在 export 时传入 dynamic_shapes 字典。

例如对一个图像分类模型，想让 batch size 和 height/width 都可变：

python-script-generator

快速生成专业的 Python 脚本和应用代码。一键创建完整项目结构，支持CLI、API、爬虫、Bot、Django等多种项目类型，包含完整的项目结构、配置文件、依赖管理、测试、README和文档。

下载

from torch.export import Dim
import torch
<p>model = YourModel().eval()
x = torch.randn(1, 3, 224, 224)</p><h1>声明动态维度：batch 和 height 可变，width 绑定到 height</h1><p>batch_dim = Dim("batch", min=1, max=32)
height_dim = Dim("height", min=128, max=480)
width_dim = height_dim  # 强制宽高一致</p><p>dynamic_shapes = {"x": {0: batch_dim, 2: height_dim, 3: width_dim}}
exported = torch.export.export(model, (x,), dynamic_shapes=dynamic_shapes)
</p>

必须用 Dim 对象（不能用字符串或整数）；min/max 是运行时校验边界，不是编译优化提示
dynamic_shapes 字典的 key 是模型 forward 的参数名（这里是 "x"），value 是该张量各维度的映射
导出后得到的 ExportedProgram 可直接用于推理：exported.module()(...) ，它内部已处理 shape 变化路径

推理时传入非法 shape 会立即报错，不是延迟失败

一旦用了 torch.export 并声明了 dynamic_shapes，运行时对输入 shape 的检查非常严格。超出 Dim 指定的 min/max 范围，或违反绑定关系（如 height=256 但 width=300），会在第一次调用时抛出清晰错误，例如：

ValueError: Input tensor x has size 300 at dimension 3, but expected size in range [128, 480] (bound to height)

这种失败发生在执行前，不进入 kernel，避免隐式 fallback 或静默降级
如果没设 min/max，则仅做符号一致性检查（比如所有绑定维度是否相等），不校验数值范围
注意：torch.compile 单独使用时无此校验机制，shape 错误往往导致 segfault 或 CUDA illegal memory access，更难调试

Inductor 编译后的动态 shape 性能不等于静态 shape

即使成功启用 dynamic_shapes=True 并通过 torch.compile 编译，实际性能仍取决于 shape 变化的“粒度”。Inductor 当前对 dynamic shape 的优化仍以“缓存多个常见 shape 的 kernel”为主，而非生成真正泛化的 kernel。

若 batch size 在 [1, 2, 4, 8] 间切换，Inductor 通常能命中缓存，延迟可控
若每次输入 shape 都完全不同（如 batch=17、23、41…），会频繁触发新 kernel 编译，首次运行极慢，且显存占用持续增长
真实部署中，建议配合 shape padding 或 bucketing：把输入 resize 到预设档位（如 1→1、2→2、3→4、5→8…），再喂给 compiled 模型

动态 shape 推理不是“设个 flag 就万事大吉”，关键在导出时的维度建模是否贴合实际业务分布，以及运行时是否做好 shape 归一化。否则容易陷入“能跑但很慢、内存暴涨”的陷阱。