PyTorch 写 Model 可以用 IfElse? 幻觉
吐个小槽. 很久以前有次我在知乎上的一个回答里夸了 TensorFlow 1.x, 然后被人抱怨说 graph mode 写不了 IfElse 不能忍.
然而, PyTorch 就可以写 IfElse 了?
最近 detectron2 遇到的产品 / 部署的需求越来越多, 看看 PyTorch 五花八门的部署 / 加速方案里对 IfElse 都有什么限制吧:
- Tracing: 遇到 IfElse 只留一支, 另一支直接消失
- Scripting: IfElse 支持的还可以 (但是循环就限制很多). 另外由于 Feature 支持的太多, 速度和优化都会更有难度.
- FX: 最近 PyTorch 新发明的用来做 quantization 的一个工具. 直接说不支持 IfElse.
- LazyTensor: pytorch/xla 里用来把模型跑在 TPU 上的 eager + graph mode 组合机制. 敢用 IfElse, 就慢 10 倍给你看.
虽然 researcher 用 PyTorch 写 IfElse 很开心, 欠的技术债终究要换一种方式还回来的. 坚持对 researcher 友好的后果就是对产品不友好.
本质上说, 如果用户写了 IfElse, 就意味着这段代码只能在单一进程的 Python 解释器里运行 -- 这本身就是一个巨大无比的 limitation. 需要用各种方式来 workaround:
- Tracing / FX: 先用 Python 跑一下试试再编译, 不保证对或者根本跑不了
- Scripting: 实现一个半成品的 Python 编译器
- LazyTensor: 用 Python 吧, 慢也没办法
虽然 TensorFlow 的 autograph 没深入用过, 但是从原理上看比以上方案都更合理. 当然, autograph 的实现是建立在 TensorFlow
已经有了足够多的 control flop operator 的前提下, 可以把 IfElse 变成 tf.cond.
而 PyTorch 在面向用户的 API 里仍然一个这样的 operator 都没有 (虽然 torchscript IR 里有), 并且可能以后也不会有.
UPDATE: 关于这件事写了一篇详细的文章: TorchScript: Tracing vs. Scripting