最近有人告诉我,在 Apache TVM 的帮助下集成/编译我的 PyTorch 模型应该可以加快推理速度,但我对 TVM 的使用感到非常困惑,因为似乎有多种使用方式。
他们网站上的这个博客说
用法很简单:
import torch_tvm
torch_tvm.enable()
而已!然后,PyTorch 将在其 JIT 编译过程中尝试将所有可以转换的运算符转换为已知的 Relay 运算符。
如果工作就这么多,为什么其他具体教程要全面得多?在经历了相当多的教程之后,我假设的另一件事(不确定我是否正确)是 TVM 不能真正使用默认的火炬代码运行方式,需要“jit”编译器(运行火炬代码的非默认方式)来研究。这个假设是基于 github repo 教程显示以下片段的事实
import torch
import torch_tvm
torch_tvm.enable()
# The following function will be compiled with TVM
@torch.jit.script
def my_func(a, b, c):
return a * b + c
函数my_func
用一个装饰器包装,看起来是用 jit 编译函数。但是使用这个完全相同的函数并在有和没有 wrapper + TVM 的情况下计时它的使用表明,这个函数的正常使用要高效得多。如果它没有加快速度,那么 jit+tvm 到底有什么帮助?如果它应该加快速度,为什么这里不是这样?
PS 对于我的问题的全面描述,我深表歉意。即使在使用 torch_tvm 阅读和修改了相当多的内容之后,我也无法弄清楚这一点并了解发生了什么。将不胜感激任何资源链接或任何可能对我有所帮助的解释。