第三章：TorchScript 与图模式编程

在自动驾驶和具身智能系统中，模型部署面临着严格的实时性和资源约束。TorchScript 提供了一种将 PyTorch 模型转换为可优化、可序列化的中间表示（IR）的方法，使得模型能够脱离 Python 运行时独立执行。本章将深入探讨 TorchScript 的核心机制，包括两种主要的转换方式、控制流处理、自定义算子开发，以及在多模态融合网络中的应用实践。

学习目标

完成本章学习后，您将能够：

理解 torch.jit.script 和 torch.jit.trace 的本质区别和适用场景
处理包含复杂控制流的模型脚本化
开发和集成高性能的自定义 TorchScript 算子
优化多模态融合网络的编译和部署
识别并避免 TorchScript 编程中的常见陷阱

3.1 torch.jit.script vs torch.jit.trace

TorchScript 提供了两种将 PyTorch 模型转换为图表示的方法：tracing 和 scripting。理解它们的工作原理和适用场景对于选择正确的编译策略至关重要。

3.1.1 基本概念与使用场景

Tracing (torch.jit.trace) 通过运行模型并记录执行的操作来创建计算图。它就像一个"录像机"，记录下模型对特定输入的处理过程。

输入张量 --> [Trace 记录器] --> 执行路径 --> 静态计算图
              |                    |
              v                    v
         记录每个操作          固定的执行流程

Tracing 的工作原理基于操作记录机制。当我们调用 torch.jit.trace(model, example_input) 时，PyTorch 会创建一个特殊的追踪环境，在这个环境中执行模型的前向传播。每当执行到一个张量操作（如矩阵乘法、卷积、激活函数等），追踪器就会记录这个操作的类型、输入输出形状、以及操作参数。这个过程类似于程序执行的动态分析，通过实际运行来理解程序行为。

追踪完成后，这些记录的操作会被组织成一个有向无环图（DAG），其中节点代表操作，边代表数据依赖关系。这个图表示是静态的，意味着无论输入数据如何变化，执行路径都是固定的。这种特性使得 trace 模式能够进行激进的优化，如操作融合、常量传播、死代码消除等。

Scripting (torch.jit.script) 则通过分析 Python 源代码的抽象语法树（AST）来生成 TorchScript 代码。它更像一个"翻译器"，将 Python 代码转换为 TorchScript 的类型化中间表示。

Python 代码 --> [AST 分析器] --> 类型推断 --> TorchScript IR
                |                   |
                v                   v
           语法解析            保留控制流

Script 模式的核心是静态分析和类型推断。编译器首先解析 Python 函数的源代码，构建抽象语法树。然后通过数据流分析推断每个变量和表达式的类型。这个过程比 trace 复杂得多，因为需要理解 Python 的语义并将其转换为更受限但性能更好的 TorchScript 语言。

TorchScript 是 Python 的一个子集，它去除了动态特性但保留了控制流结构。这意味着 if-else 条件、for/while 循环、函数调用等都会被保留在最终的计算图中。运行时，这些控制流会根据实际输入动态执行不同的路径，提供了比 trace 更大的灵活性。

在自动驾驶场景中，感知模型（如目标检测网络）通常具有固定的前向传播路径，适合使用 trace；而决策规划模块包含大量条件判断，必须使用 script。例如，一个典型的自动驾驶系统可能包含：

感知层（适合 trace）： - 图像编码器：ResNet、EfficientNet 等 CNN 架构 - 特征金字塔网络（FPN）：固定的特征融合路径 - 检测头的卷积层：规则的卷积-BN-ReLU 序列
决策层（需要 script）： - 轨迹预测：根据目标类型选择不同的预测模型 - 风险评估：基于场景的条件判断逻辑 - 行为规划：包含复杂的状态机和决策树

理解这两种模式的本质差异对于设计高效的部署方案至关重要。在实践中，往往需要将模型分解为多个部分，对每部分选择最合适的编译策略。

3.1.2 Trace 模式的优势与限制

优势：

简单直观：无需修改代码，直接对现有模型进行 tracing
性能优越：生成的图更加紧凑，没有控制流开销
优化充分：静态图便于进行常量折叠、算子融合等优化

让我们深入理解 trace 模式的优势来源。静态图的最大优势在于编译时优化的空间。当整个计算图在编译时完全确定，编译器可以进行全局分析和优化：

常量折叠（Constant Folding）：编译器可以在编译时计算所有常量表达式。例如，如果模型中有 x * 2.0 * 0.5 的操作序列，编译器会直接将其优化为 x * 1.0，甚至进一步优化为 x。这种优化在包含批归一化层的网络中特别有效，因为在推理模式下，BN 的参数都是常量。

算子融合（Operator Fusion）：相邻的逐元素操作可以融合成一个内核。典型的例子是 Conv-BN-ReLU 序列，原本需要三次内存访问（写卷积结果、读写 BN、读写 ReLU），融合后只需要一次内存往返。在内存带宽受限的场景下，这种优化可以带来 2-3 倍的性能提升。

内存规划（Memory Planning）：静态图允许编译器预先分配所有中间张量的内存，并且可以复用不再使用的内存空间。这种内存池化策略可以显著减少内存分配的开销和碎片化。

限制：

无法处理控制流：条件分支和循环会被"烘焙"成固定路径
输入形状敏感：只能处理与 trace 时相同形状的输入
动态行为丢失：依赖于输入值的动态计算无法正确捕获

这些限制的根本原因在于 trace 的"录制-回放"机制。让我们通过具体例子理解每个限制：

控制流固化问题：假设模型中有一个阈值判断逻辑，当置信度大于 0.5 时执行路径 A，否则执行路径 B。在 trace 时，如果示例输入导致置信度为 0.7，那么只有路径 A 会被记录。即使后续输入的置信度为 0.3，模型仍然会执行路径 A，导致错误的结果。

形状固化问题：Trace 会记录具体的张量形状而非符号形状。如果 trace 时输入形状是 [1, 3, 224, 224]，那么推理时输入 [2, 3, 224, 224] 会导致形状不匹配错误。虽然可以通过 torch.jit.trace 的 check_inputs 参数提供多个示例来缓解，但这仍然是一个根本性限制。

动态计算丢失：某些操作的结果依赖于张量的具体值而非形状。例如，torch.nonzero（找出非零元素的索引）返回的张量大小取决于输入中非零元素的数量。Trace 模式会固定返回特定大小的张量，无法适应不同输入。

考虑一个简化的车道线检测后处理模块：

模型输出 --> NMS处理 --> 动态数量的检测框
           |
           v
      根据置信度阈值筛选（控制流）

这种情况下，trace 会固定检测框的数量，无法适应不同场景下检测结果数量的变化。

实际案例：YOLO 后处理的挑战

在 YOLO 等目标检测模型中，后处理步骤包含大量动态操作：

置信度筛选：保留置信度大于阈值的预测框
NMS 处理：根据 IoU 阈值去除重复检测
类别筛选：可能只保留特定类别的检测结果

这些操作的输出数量都是动态的，取决于场景复杂度。在高速公路场景可能只有几个目标，而在拥挤的城市街道可能有上百个目标。使用 trace 会将检测数量固定在 trace 时的值，严重限制模型的实用性。

缓解策略：

分离静态和动态部分：将模型分为静态的特征提取部分（使用 trace）和动态的后处理部分（使用 script 或保持为 Python）
填充到固定大小：始终返回固定数量的检测框，用无效值填充，但这会浪费计算和内存
多版本编译：为不同的输入配置编译多个版本，运行时根据输入特征选择

3.1.3 Script 模式的类型系统

TorchScript 实现了一个静态类型系统，支持以下核心类型：

基础类型：

Tensor：张量类型，可指定设备和数据类型
int, float, bool：标量类型
str：字符串类型
None：空类型

TorchScript 的类型系统设计借鉴了静态类型语言的优点，同时保持了与 Python 的兼容性。这个类型系统的核心目标是在编译时捕获类型错误，生成高效的机器码，并提供清晰的接口定义。

张量类型的细节：在 TorchScript 中，Tensor 类型包含了丰富的元信息：

设备类型：CPU、CUDA、XLA 等
数据类型：float32、float16、int8 等
形状信息：可以是具体形状如 [3, 224, 224]，也可以是符号形状如 [N, C, H, W]
内存布局：连续（contiguous）或跨步（strided）

这些信息使得编译器能够选择最优的内核实现。例如，知道张量在 CUDA 上且是 float16 类型，编译器可以选择使用 Tensor Core 加速的内核。

容器类型：

List[T]：同质列表
Tuple[T1, T2, ...]：异质元组
Dict[K, V]：字典类型
Optional[T]：可选类型

容器类型的设计平衡了表达能力和性能。让我们深入理解每种容器的特性和使用场景：

List[T] - 同质列表：所有元素必须是相同类型 T。这个限制使得列表可以用连续内存表示，访问效率高。在自动驾驶中，List[Tensor] 常用于表示多尺度特征图或多个目标的特征向量。列表支持动态长度，可以在运行时添加或删除元素，但频繁的大小变化会导致内存重分配。

Tuple - 异质元组：可以包含不同类型的元素，但长度和类型在编译时固定。元组通常用于函数返回多个值，如检测模型返回 (boxes, scores, classes)。由于元组的结构在编译时已知，编译器可以完全展开元组操作，避免运行时开销。

Dict[K, V] - 字典类型：键和值都必须是同质的。TorchScript 的字典实现基于哈希表，提供 O(1) 的平均访问时间。在多任务学习中，Dict[str, Tensor] 常用于存储不同任务的输出。需要注意的是，字典的迭代顺序在 TorchScript 中是确定的（插入顺序），这与 Python 3.7+ 一致。

Optional[T] - 可选类型：表示值可能是 T 类型或 None。这在处理可能失败的操作时特别有用。例如，目标跟踪中，如果目标丢失，返回 Optional[Tensor] 可以优雅地处理这种情况。编译器会插入必要的空值检查，确保类型安全。

类型推断规则：

函数参数需要类型注解（除了 self 和 Tensor 类型）
返回值类型会自动推断
局部变量类型从初始化表达式推断
类属性必须在 init 中初始化

类型推断的工作原理基于数据流分析。编译器构建一个类型约束图，通过求解约束系统来推断未知类型。这个过程类似于 Hindley-Milner 类型推断，但针对 PyTorch 的特点进行了定制：

前向推断：从已知类型的表达式推断未知类型

x = torch.zeros(3, 4)  # x 的类型推断为 Tensor
y = x.sum()            # y 的类型推断为 Tensor（标量）
z = y.item()           # z 的类型推断为 float

后向推断：从使用上下文推断类型

def process(x: Tensor) -> Tensor:
    y = helper(x)      # helper 的参数类型推断为 Tensor
    return y * 2       # helper 的返回类型推断为 Tensor

泛型推断：对于泛型函数，根据实际参数推断类型参数

def first(x: List[T]) -> T:  # T 是类型变量
    return x[0]

result = first([1, 2, 3])    # T 推断为 int

类型系统确保了编译时的类型安全，避免了运行时类型错误。在具身智能的控制网络中，明确的类型定义有助于：

提前发现类型不匹配错误
生成更高效的机器码
支持硬件加速器的类型映射

高级类型特性：

Union 类型：TorchScript 支持有限的 Union 类型，主要用于 Optional。完整的 Union 类型支持（如 Union[int, float]）会使类型推断复杂化，因此被限制。

类型别名：可以定义类型别名来提高代码可读性：

DetectionOutput = Tuple[Tensor, Tensor, Tensor]  # (boxes, scores, classes)

递归类型：TorchScript 不支持直接的递归类型定义，但可以通过类来实现树形结构。

类型注解的最佳实践：

明确优于隐式：即使可以推断，添加类型注解可以提高代码可读性
使用类型别名：为复杂类型定义别名，特别是在接口边界
避免 Any 类型：TorchScript 不支持 Any，强制类型明确性
容器类型参数化：始终指定容器的元素类型，如 List[Tensor] 而非 List

3.1.4 混合使用策略

实际应用中，通常需要结合使用 trace 和 script 来获得最佳效果。常见的混合策略包括：

策略1：主干 trace + 后处理 script

CNN 特征提取（trace）--> 特征图 --> 检测头（script with NMS）
                                    |
                                    v
                              包含控制流的后处理

策略2：递归嵌套 通过 torch.jit.script 装饰器标记需要保留控制流的函数，然后在 trace 过程中调用这些函数。

策略3：分阶段编译

先对稳定的模块进行 trace
对包含动态逻辑的模块使用 script
最后组合成完整的 TorchScript 模型

在自动驾驶的多任务网络中，可以对共享的骨干网络使用 trace，而对不同的任务头（检测、分割、深度估计）分别选择合适的编译方式。

3.2 控制流与动态行为处理

控制流是区分 TorchScript 与静态图框架的关键特性。正确处理条件分支、循环和递归对于构建灵活的推理系统至关重要。

3.2.1 条件分支的编译优化

TorchScript 支持 Python 的 if-else 语句，但会进行特殊的编译优化：

分支预测与推测执行： 当条件可以在编译时确定时，TorchScript 会进行死代码消除（DCE）。对于运行时条件，编译器会生成两个分支的代码，并在执行时选择。

条件表达式 --> [静态分析] --> 编译时常量？
                |                    |
                v                    v
           生成两个分支          消除死分支
                |
                v
          运行时分支选择

条件计算的向量化： 对于简单的条件赋值，TorchScript 会尝试使用 torch.where 进行向量化：

标量条件 --> 逐元素判断 --> 性能开销大
    |
    v
向量化条件 --> torch.where --> SIMD 并行

在自动驾驶的目标跟踪中，根据目标的运动状态选择不同的预测模型是常见需求。合理的条件分支设计可以显著提升推理效率。

3.2.2 循环结构的展开与优化

TorchScript 支持 for 和 while 循环，但有特定的优化策略：

循环展开（Loop Unrolling）： 对于固定次数的小循环，编译器会自动展开以减少循环开销：

for i in range(4):     -->    操作0; 操作1; 操作2; 操作3
    操作(i)

循环不变量外提（Loop Invariant Code Motion）： 将循环内不变的计算移到循环外：

循环体 --> [依赖分析] --> 不变量？
            |                |
            v                v
        保留在循环内    提升到循环外

动态循环的限制：

循环次数必须是整数类型
循环体内不能修改循环变量
break 和 continue 支持但会影响优化

在点云处理中，对每个点进行独立处理的循环可以通过批量化操作替代，避免 Python 循环的开销。

3.2.3 动态形状的处理技巧

动态形状是部署中的常见挑战，TorchScript 提供了符号形状（symbolic shapes）机制：

形状传播规则：

保守传播：未知维度标记为 -1
形状特化：为常见形状生成特化代码
运行时检查：插入必要的形状断言

处理策略：

动态批次大小：
输入 [N, C, H, W] --> N 为符号维度
                      |
                      v
                编译多个特化版本
                (N=1, N=4, N=8, ...)

形状相关的优化：

内存预分配：基于最大可能形状
算子选择：根据实际形状选择最优实现
缓存策略：重用相同形状的中间结果

在具身智能的视觉输入处理中，图像分辨率可能因任务而异。使用形状特化可以在保持灵活性的同时获得接近静态图的性能。

3.2.4 递归函数的脚本化

TorchScript 支持递归函数，但需要注意：

递归深度限制： 默认递归深度限制为 256，可通过 torch._C._jit_set_max_recursion_depth 调整。

尾递归优化： TorchScript 不进行自动的尾递归优化，需要手动转换为循环。

递归类型推断： 递归函数必须有明确的返回类型注解，因为类型推断无法处理递归依赖。

递归调用 --> [深度检查] --> 超过限制？
              |                |
              v                v
          继续执行        抛出异常
              |
              v
        栈帧管理开销

最佳实践：

优先使用迭代而非递归
对于树形结构遍历，考虑使用显式栈
利用动态规划避免重复计算

在机器人的运动规划中，递归的路径搜索算法（如 RRT）应该转换为迭代版本以提高效率。

3.3 自定义 TorchScript 算子

当内置算子无法满足需求时，自定义算子成为必要选择。这在处理特定硬件加速、领域特定算法时尤为重要。

3.3.1 算子注册机制

TorchScript 的算子注册遵循分发（dispatch）机制：

算子调用 --> [Dispatcher] --> 后端选择
                |               |
                v               v
           (CPU/CUDA/...)   具体实现
                |
                v
           Schema 验证

注册流程：

定义算子 Schema：描述输入输出类型
实现算子逻辑：C++ 或 CUDA 代码
注册到分发器：关联 schema 和实现
Python 绑定：使算子可从 Python 调用

Schema 定义语法：

<返回类型> <算子名>(<参数列表>) -> <返回值名称>

支持的参数类型包括：

Tensor：张量参数
Scalar：标量参数（int, float, bool）
int[]：整数列表
Tensor?：可选张量
Tensor(a!)：原地修改的张量

在自动驾驶的激光雷达处理中，体素化（voxelization）是常见的自定义算子需求，将点云转换为规则的 3D 网格表示。

3.3.2 C++ 扩展开发

开发 C++ 扩展的关键步骤：

项目结构：

custom_op/
├── csrc/
│   ├── custom_op.cpp    # CPU 实现
│   └── custom_op_cuda.cu # CUDA 实现
├── setup.py              # 编译配置
└── __init__.py           # Python 接口

CPU 算子实现要点：

内存管理：使用 at::empty 分配输出张量
并行化：利用 at::parallel_for 进行多线程
类型分发：AT_DISPATCH_ALL_TYPES 宏处理不同数据类型
错误处理：TORCH_CHECK 进行运行时检查

性能优化技巧：

避免频繁的内存分配
利用连续内存布局
使用 SIMD 指令（通过 ATen 的向量化接口）
缓存友好的数据访问模式

在点云的最远点采样（FPS）算法中，自定义 C++ 实现可以比纯 Python 版本快 100 倍以上。

3.3.3 性能优化策略

算子融合（Operator Fusion）： 将多个小算子合并为一个大算子，减少内存访问：

Conv -> BN -> ReLU  -->  FusedConvBNReLU
    |      |     |              |
    v      v     v              v
  3次内存访问            1次内存访问

内存布局优化：

通道最后格式（NHWC）：对于某些硬件更高效
内存对齐：确保数据对齐到缓存行边界
零拷贝视图：利用 stride 操作避免数据复制

计算图优化：

常量折叠：编译时计算常量表达式
公共子表达式消除：重用相同的计算结果
死代码消除：移除未使用的计算

批处理策略：

动态批处理：
输入队列 --> [批处理器] --> 批次大小？
                |              |
                v              v
           延迟 vs 吞吐量   自适应调整

在具身智能的实时控制中，需要在延迟和吞吐量之间找到平衡点。

3.3.4 与 CUDA 内核集成

CUDA 内核开发的关键考虑：

线程组织：

Grid --> Blocks --> Threads
  |         |          |
  v         v          v
全局索引  共享内存   寄存器

内存层次优化：

全局内存：合并访问，避免 bank 冲突
共享内存：块内线程共享，低延迟
常量内存：广播读取，缓存优化
纹理内存：空间局部性优化

CUDA 算子实现模式：

主机代码：

1. 检查输入
2. 分配输出
3. 计算启动配置
4. 调用内核
5. 同步（可选）

设备代码：

1. 计算全局索引
2. 边界检查
3. 执行计算
4. 写入结果

性能调优工具：

nsys：系统级性能分析
ncu：内核级性能分析
cuda-memcheck：内存错误检测

常见优化技术：

Warp 级原语：__shfl_sync, __ballot_sync
原子操作优化：使用 warp 级归约减少原子操作
流并发：多流执行隐藏延迟
混合精度：利用 Tensor Core

在自动驾驶的 3D 检测中，NMS（非极大值抑制）的 CUDA 实现是关键性能瓶颈，通过优化的并行算法可以实现实时处理。

3.4 多模态融合网络的脚本化

多模态融合是自动驾驶和具身智能的核心技术，涉及视觉、语言、激光雷达等多种输入的联合处理。

3.4.1 视觉-语言模型的编译挑战

视觉-语言模型（VLM）结合了计算机视觉和自然语言处理，面临独特的编译挑战：

动态序列长度：

文本输入 --> [Tokenizer] --> 可变长度序列
                |                   |
                v                   v
           动态 padding         动态计算图

跨模态注意力： 视觉特征和文本特征的交互需要特殊处理：

视觉特征 [N, H*W, D_v] ──┐
                          ├──> 跨模态注意力 --> 融合特征
文本特征 [N, L, D_t] ────┘         |
                                   v
                              动态矩阵乘法

编译策略：

静态 padding：预定义最大序列长度
桶化（Bucketing）：为不同长度范围编译不同版本
动态批处理：运行时重组批次

在自动驾驶的场景理解中，需要将道路图像与导航指令进行融合，动态序列长度是必须处理的问题。

3.4.2 注意力机制的优化

注意力机制是多模态融合的核心，其优化至关重要：

Flash Attention 集成：

标准注意力：
Q @ K^T --> Softmax --> @ V
   |           |          |
   v           v          v
O(N²) 内存   数值不稳定  带宽瓶颈

Flash Attention：
分块计算 --> 在线 Softmax --> 融合内核
   |            |              |
   v            v              v
O(N) 内存   数值稳定     计算受限

多头注意力的并行化：

头间并行：不同注意力头独立计算
序列并行：长序列分割到多个设备
张量并行：矩阵乘法的分块计算

KV Cache 优化： 在自回归生成中，缓存键值对以避免重复计算：

增量计算：
新 token --> 计算 Q_new --> 与缓存的 K,V 交互
                |                    |
                v                    v
           仅计算新位置          重用历史计算

3.4.3 动态路由的实现

动态路由允许模型根据输入选择不同的处理路径：

专家混合（MoE）架构：

输入 --> [Gate] --> Top-K 选择 --> 专家网络
           |            |             |
           v            v             v
      路由概率    稀疏激活      并行处理

TorchScript 中的实现挑战：

动态分发：根据门控输出选择专家
负载均衡：确保专家间的计算均衡
批处理不规则性：不同样本选择不同专家

优化技术：

静态路由表：预计算可能的路由组合
专家并行：将不同专家放置在不同设备
容量因子：限制每个专家的最大处理量

在具身智能的任务规划中，不同任务类型（导航、操作、交互）可能需要激活不同的专家模块。

3.4.4 部署优化实践

模型分片策略：

完整模型 --> [分析依赖] --> 独立子图
                |              |
                v              v
           识别切分点      最小化通信

异构设备部署：

CPU 部分：预处理、后处理、控制逻辑
GPU 部分：密集计算、矩阵运算
专用加速器：特定算子（如 NPU 的卷积）

内存优化技术：

激活检查点：用计算换内存
量化感知脚本化：编译时考虑量化
动态内存池：重用中间缓冲区

延迟优化清单：

算子融合减少内核启动
预分配输出缓冲区
异步执行和流水线
编译时常量传播

吞吐量优化清单：

动态批处理提高利用率
多实例并行服务
请求级负载均衡
自适应批大小调整

在自动驾驶的端到端模型中，需要同时处理多个摄像头输入、激光雷达点云和高精地图，合理的模型分片和设备分配是实现实时性能的关键。

本章小结

本章深入探讨了 TorchScript 的核心概念和实践技术：

核心要点：

Trace vs Script：trace 适合静态计算图，script 保留动态控制流
类型系统：静态类型确保编译时安全和运行时性能
控制流优化：条件分支、循环和递归的编译策略
自定义算子：通过 C++/CUDA 扩展实现特定功能
多模态融合：处理不同输入模态的编译挑战

关键公式：

注意力复杂度：O(N²D) 其中 N 是序列长度，D 是特征维度
Flash Attention 内存：O(N) vs 标准 O(N²)
算子融合收益：减少 (n-1) 次内存往返，n 为融合算子数

性能提升经验值：

Trace 模式：通常获得 1.5-3x 加速
Script 模式：1.2-2x 加速（保留灵活性）
自定义 CUDA 算子：10-100x 加速（相比纯 Python）
算子融合：20-40% 内存带宽节省

练习题

基础题

练习 3.1：Trace vs Script 选择 给定以下场景，选择合适的 TorchScript 转换方式并说明理由： a) ResNet-50 图像分类模型 b) 带有动态 NMS 的 YOLO 检测器 c) Transformer 解码器的自回归生成 d) 根据置信度动态选择处理分支的融合网络

提示：考虑是否有控制流和动态行为

参考答案

a) Trace：ResNet-50 是纯前向传播网络，没有控制流 b) Script：NMS 包含动态循环和条件判断 c) Script：自回归生成有循环和 KV cache 更新 d) Script：动态分支选择需要保留控制流

练习 3.2：类型注解修正 以下 TorchScript 函数有类型错误，请修正：

@torch.jit.script
def process_boxes(boxes, scores, threshold):
    keep = []
    for i in range(len(scores)):
        if scores[i] > threshold:
            keep.append(boxes[i])
    return keep

提示：TorchScript 需要明确的类型注解

参考答案

需要添加类型注解：

boxes: Tensor
scores: Tensor
threshold: float
keep: List[Tensor]
返回类型: List[Tensor]

函数参数和返回值都需要正确的类型标注。

练习 3.3：循环优化识别 判断以下循环是否会被 TorchScript 自动展开：

# 循环 A
for i in range(4):
    x = x + weights[i]

# 循环 B  
for i in range(n):  # n 是输入参数
    x = x * scale

提示：考虑循环次数是否编译时已知

参考答案

循环 A：会展开，循环次数（4）是编译时常量循环 B：不会展开，循环次数 n 是运行时变量

编译器只能展开静态已知次数的小循环。

挑战题

练习 3.4：混合编译策略设计 设计一个自动驾驶感知模型的 TorchScript 编译策略。模型包含：

共享的 ResNet backbone
三个任务头：目标检测（带 NMS）、语义分割、深度估计
后处理融合模块（根据场景动态加权）

要求说明每个部分使用 trace 还是 script，以及理由。

提示：考虑计算特性和部署需求

参考答案

编译策略：

ResNet backbone：使用 trace，固定前向传播路径
检测头主体：使用 trace，卷积和全连接层
NMS 后处理：使用 script，包含动态循环
分割头：使用 trace，纯卷积上采样
深度估计头：使用 trace，固定解码器结构
融合模块：使用 script，动态加权逻辑

组合方式：backbone trace 后，各头部独立处理，最后 script 融合。

练习 3.5：自定义算子性能分析 某点云体素化算子的两种实现：

实现 A：纯 Python 循环，处理 10000 个点需要 100ms
实现 B：自定义 CUDA 算子，处理 10000 个点需要 2ms

计算： a) 加速比是多少？ b) 若点云有 100000 个点，估算各自用时 c) 考虑 GPU 内核启动开销约 0.5ms，重新评估小批量（100个点）的性能

提示：注意固定开销的影响

参考答案

a) 加速比 = 100ms / 2ms = 50x

b) 假设线性扩展：

Python：100ms × 10 = 1000ms
CUDA：2ms × 10 = 20ms

c) 100个点的情况：

Python：100ms / 100 = 1ms
CUDA：0.5ms（启动）+ 0.02ms（计算）≈ 0.52ms
此时加速比仅 1.9x，固定开销占主导

结论：CUDA 算子在大批量时优势明显，小批量时开销不可忽略。

练习 3.6：注意力机制内存优化 标准注意力的内存占用为 O(N²)，其中 N 是序列长度。若：

序列长度 N = 2048
隐藏维度 D = 512
批大小 B = 8
使用 float16

计算： a) 存储 attention scores 需要多少内存？ b) Flash Attention 将内存降至 O(N)，节省了多少？ c) 在 V100 (32GB) 上，标准注意力最大能处理多长序列？

提示：float16 占 2 字节

参考答案

a) Attention scores 形状：[B, N, N] 内存 = 8 × 2048 × 2048 × 2 bytes = 64 MB

b) Flash Attention O(N) 内存：约 8 × 2048 × 512 × 2 = 16 MB 节省：64 - 16 = 48 MB (75%节省)

c) 设最大序列长度为 L： 8 × L² × 2 ≤ 32 × 10⁹ L² ≤ 2 × 10⁹ L ≤ 44721

实际受其他因素限制，约 16K-32K

练习 3.7：多模态动态路由设计 设计一个 MoE 路由策略，有 8 个专家，每个样本选择 top-2 专家。若批大小为 32，如何优化以下指标： a) 负载均衡：每个专家处理的样本数尽量均匀 b) 批处理效率：相同专家的样本可以批处理 c) 内存占用：避免所有专家同时激活

提示：考虑容量因子和路由正则化

参考答案

优化策略：

a) 负载均衡：

添加负载均衡损失：L_balance = Var(expert_loads)
设置容量因子：每个专家最多处理 32×2/8×1.25 = 10 个样本

b) 批处理效率：

对路由结果排序，相同专家的样本组成批次
使用 block-sparse 矩阵乘法

c) 内存优化：

顺序执行：一次只激活 2-3 个专家
使用激活检查点，只保存路由结果
专家权重使用 CPU offload

实现要点：路由概率添加噪声促进探索，使用辅助损失平衡负载。

练习 3.8：编译优化效果评估 某具身智能模型优化前后对比：

原始 PyTorch：延迟 50ms，吞吐量 100 QPS
TorchScript：延迟 30ms，吞吐量 150 QPS
- 算子融合：延迟 25ms，吞吐量 180 QPS
- 量化 (INT8)：延迟 15ms，吞吐量 300 QPS

问：若需要满足 20ms 延迟约束，处理 10000 请求/秒的负载，需要多少个模型实例？

提示：考虑延迟和吞吐量约束

参考答案

分析各版本：

原始 PyTorch：50ms > 20ms，不满足延迟约束
TorchScript：30ms > 20ms，不满足延迟约束
TorchScript + 融合：25ms > 20ms，不满足延迟约束
TorchScript + 融合 + 量化： - 延迟 15ms < 20ms ✓ - 单实例吞吐量：300 QPS - 需要实例数：10000 / 300 = 34 个实例

结论：需要使用量化版本，部署 34 个实例。

常见陷阱与错误

1. Trace 模式的隐式固化

问题：trace 会将所有动态行为固化为静态路径

# 错误：条件分支被固化
traced_model = torch.jit.trace(model, example_input)
# 不同条件的输入会得到错误结果

解决：对包含控制流的模块使用 script

2. 类型推断失败

问题：TorchScript 无法推断复杂的类型

# 错误：返回类型不明确
def process(x):
    if x.sum() > 0:
        return x, True
    return None  # 类型不一致

解决：使用 Optional 类型和一致的返回类型

3. Python 特性不支持

问题：并非所有 Python 特性都被支持

不支持：yield、async/await、大部分内置库
受限：字符串操作、复杂数据结构解决：简化代码逻辑，使用支持的操作

4. 原地操作的陷阱

问题：某些原地操作会破坏自动微分

# 危险：原地修改可能导致梯度错误
x[mask] = 0  # 可能有问题

解决：使用 torch.where 等函数式操作

5. 动态形状的性能退化

问题：过度动态的形状导致无法优化解决：使用形状特化，为常见形状创建专门版本

6. 自定义算子的版本兼容

问题：C++ 扩展可能与 PyTorch 版本不兼容解决：使用条件编译，测试多版本兼容性

7. 递归深度溢出

问题：默认递归限制可能不够解决：转换为迭代或增加递归深度限制

8. 混合精度的数值问题

问题：某些操作在低精度下不稳定解决：关键操作保持 FP32，使用自动混合精度

最佳实践检查清单

设计阶段

[ ] 识别模型中的动态和静态部分
[ ] 确定 trace vs script 的边界
[ ] 评估控制流的必要性
[ ] 规划类型注解策略

实现阶段

[ ] 为所有函数参数添加类型注解
[ ] 避免使用不支持的 Python 特性
[ ] 优先使用向量化操作而非循环
[ ] 测试不同输入形状的正确性

优化阶段

[ ] 识别可融合的算子序列
[ ] 评估自定义算子的必要性
[ ] 配置合适的并行策略
[ ] 测试量化对精度的影响

部署阶段

[ ] 验证序列化模型的正确性
[ ] 测试边界条件和异常输入
[ ] 监控内存使用和泄漏
[ ] 准备回滚方案

性能调优

[ ] 使用 Profiler 识别瓶颈
[ ] 比较不同批大小的性能
[ ] 评估预热（warm-up）的影响
[ ] 测试多线程/多流并发

维护阶段

[ ] 记录编译配置和依赖
[ ] 建立性能基准测试
[ ] 监控生产环境指标
[ ] 定期更新和优化