pytorch的 jit功能不应该叫aot功能吗?

发布时间 2023-10-16 15:17:41作者: ZephyrYin

jit简介

pytorch的jit功能是为了生成类似于计算图的功能,能加快模型速度,便于部署,属于是取tensorflow之长补自己之短。

jit由来

之所以叫jit,是因为他仍然属于实时解释的范畴,但是能预先生成中间量,“审时度势”地实时运行。

参考

可以看这位帖子:https://zhuanlan.zhihu.com/p/650885922
以下均为摘抄:

对于 C/C++ 之类的编程语言,我们会将所有的代码提前“翻译”成机器语言,这样在实际运行时就不会有额外的开销,可以全速进行推理。但是这样做也有一些缺点:
编译花费的时间很长,对于需要频繁修改代码的场景,欠缺灵活性
编译时无法感知运行环境,也就无法采取一些针对硬件环境的优化

另一种做法是使用解释器,不对代码进行提前编译,而是在运行时“边解释边执行”,代码更灵活,也可以尝试作一些针对当前环境的优化。这种语言被称为“解释性语言”,比如 Lisp、Pascal 等。当然代价通常就是更差的性能,毕竟“编译”需要占用运行时的时间,而且由于不能得到全部上下文,所以无法进行依赖上下文的优化。

那么如果“我都要”呢?既希望代码可以充分利用当前环境,又不希望牺牲性能的话,可以尝试第三种方案,也就是 JIT。JIT 是 just-in-time 的缩写,它不会将编译的过程一口气完成,而是先对代码进行一些处理,存储成某种序列化表示(比如计算图);然后在实际的运行时环境中,通过 profiling 的方式,进行针对环境的优化并执行代码。

torch jit 的名字就来源于此,PyTorch 使用 trace 或 script 之类的方法将模型转换成计算图,然后在运行时 "just in time" 的优化和执行推理过程。