在采用即时编译(JIT)时,首要步骤是从高级Python代码中获取模型的计算图。这个获取到的图作为JIT编译器进行优化和代码生成的输入。在计算图获取阶段,主要有两种策略:追踪和脚本。每种方法都有其独特的优点和局限,影响着它们能很好地处理的模型类型以及随后可进行的优化。追踪:记录执行路径追踪的运作方式是使用示例输入执行模型函数,并记录在该次特定执行过程中对张量对象进行的操作序列。可以将其想象成一个专门记录机器学习操作及其数据依赖的分析器。工作方式:您向JIT系统提供模型函数(例如,torch.nn.Module的forward方法,或用@tf.function装饰的TensorFlow函数)。您使用示例张量输入调用该函数(例如,model(example_input))。JIT机制会拦截涉及张量的框架操作调用(如卷积、矩阵乘法、激活函数)。它会构建一个图表示(通常是中间表示,如TorchScript IR、TensorFlow GraphDef或HLO),该表示精确反映了这些特定输入所执行的操作序列。示例:考虑一个简单的Python函数:def simple_op(a, b): c = a + b d = c * 2 return d如果使用a = tensor([1])和b = tensor([2])进行追踪,追踪器会记录:操作:add(输入:a,b;输出:c)操作:mul(输入:c,常数2;输出:d)返回:d生成的图捕获了这种线性序列。优点:简便性: 对现有Python代码应用追踪通常很直接,无需大量修改。如果代码以即时执行模式运行,则很可能可以进行追踪。Python兼容性: 追踪自然地处理框架操作之间执行的任意Python代码,只要这些代码不影响图结构本身(例如,使用标准Python数据结构,打印值)。缺点:静态控制流: 追踪在处理依赖于数据的控制流时存在根本性困难。如果模型包含Python的if、for或while语句,且条件或循环边界取决于张量值,则追踪只捕获追踪时使用的特定示例输入所经过的路径。生成的图将不包含其他分支,也无法通用地表示循环结构。def conditional_op(x, threshold): if x.sum() > threshold: # Data-dependent condition return x * 2 else: return x + 1使用满足条件的x追踪conditional_op会生成一个只包含x * 2路径的图。x + 1路径将完全缺失。输入依赖性: 被追踪的图与追踪期间使用的输入的属性(如形状、数据类型)内在相关。虽然一些JIT系统随后可以处理有限的动态性,但初始追踪可能过于专门化。副作用: 追踪可能无法正确捕获Python的副作用,或者可能以意想不到的方式将其固化到图中。脚本:解析源代码脚本采用一种不同的方法。它不执行代码,而是直接解析模型函数(或其子集)的Python源代码,并将其转换为图表示,其中包含控制流结构。工作方式:通常,您需要明确标记要进行脚本化的函数(例如,在PyTorch中使用@torch.jit.script),或者使用脚本编译器可以理解的受限Python子集来编写函数。脚本编译器解析函数的源代码,构建抽象语法树(AST)。它分析AST,将Python构造(赋值、运算符、控制流)转换为其中间表示(IR)中相应的图节点和结构。示例:使用相同的conditional_op函数:@torch.jit.script # 示例装饰器 def conditional_op(x, threshold): # 脚本编译器解析此结构 if x.sum() > threshold: result = x * 2 else: result = x + 1 return result脚本编译器分析if/else结构,并生成一个图,其中包含表示条件(x.sum() > threshold)的节点,以及true分支(x * 2)和false分支(x + 1),以及在运行时选择适当路径的控制流机制。digraph G { rankdir=TB; node [shape=box, style=filled, fillcolor="#e9ecef", fontname="Arial"]; edge [fontname="Arial"]; X [label="输入 x", shape=ellipse, fillcolor="#a5d8ff"]; Threshold [label="输入 阈值", shape=ellipse, fillcolor="#a5d8ff"]; Sum [label="x.sum()"]; Compare [label="> threshold"]; If [label="如果", shape=diamond, fillcolor="#ffe066"]; Mul [label="x * 2"]; Add [label="x + 1"]; Result [label="结果", shape=ellipse, fillcolor="#b2f2bb"]; X -> Sum; Threshold -> Compare; Sum -> Compare; Compare -> If; If -> Mul [label="真"]; If -> Add [label="假"]; X -> Mul; X -> Add; Mul -> Result; Add -> Result; }脚本化conditional_op函数生成的图表示,清晰地展示了条件分支。优点:控制流处理: 脚本化擅长捕获依赖于张量值的Python控制流(if、for、while),并将其直接在图中表示出来。输入独立性: 生成的图表示函数的实际逻辑,独立于“编译”步骤中使用的任何特定输入值。清晰的图: 提供对预期计算更完整的表示,可能允许在控制流路径上进行更高级的图级优化。缺点:语言子集: 脚本编译器通常只支持完整Python语言的一个子集。您可能需要重构复杂的Python逻辑、列表推导或对不支持库的调用,以符合可脚本化的子集。代码修改: 通常需要明确的注解(装饰器)或遵守特定的编码模式。它通常无法透明地处理任意的、未经修改的Python代码。更高的学习曲线: 开发者需要理解脚本语言子集的限制和要求。调试脚本错误有时不如调试标准Python直观。追踪与脚本的选择追踪与脚本之间的选择通常取决于模型性质和开发流程:特点追踪脚本易用性通常对现有Python代码更简单需要代码调整/注解控制流较差(只捕获一条路径)良好(明确捕获分支/循环)Python功能处理大多数操作之间的Python代码限于语言子集输入依赖性高(图与追踪输入相关联)低(图表示代码逻辑)稳定性如果控制流改变,可能脆弱表示更可靠适用场景简单模型、快速原型、静态图具有数据依赖控制流的模型、部署现代JIT系统通常同时提供这两种选项。例如,PyTorch的TorchScript允许用户选择@torch.jit.trace或@torch.jit.script,甚至可以组合追踪和脚本化模块。TensorFlow的@tf.function主要采用追踪机制(“autograph”会隐式转换一些Python控制流,略微模糊了界限,但其核心捕获方式是基于追踪的)。了解观察执行路径(追踪)与解析代码逻辑(脚本)之间的主要区别,对于有效使用JIT编译器十分重要。追踪提供了便利,但限制了表达能力,尤其是在动态控制流方面。脚本化要求开发者投入更多精力来适应其限制,但能生成更完整的表示,可以处理复杂的程序结构,这有助于在JIT编译器中进行更全面的优化。