Github项目推荐 | Pytorch TVM 扩展

将TVM集成到PyTorch中。

* TVM:深度学习编译器

Github项目链接:

https://github.com/pytorch/tvm

构建

你需要在这个PR的基础上构建PyTorch:https://github.com/pytorch/pytorch/pull/18588

cd pytorch
git fetch origin pull/18588/head:tvm_dev
git checkout tvm_dev
python setup.py install

然后,你需要单独构建本仓库:

# Make sure the right llvm-config is in your PATH
  python setup.py install

测试

python setup.py test

使用

这个包非常显然地挂钩到 PyTorch 的 JIT 中,因此适用相同的工具(可查看 @torch.jit.script,torch.jit.trace 和 graph_for)。以下是使用示例:

from tvm import relay # This imports all the topi operators
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

如果要禁用JIT挂钩,请使用 torch_tvm.disable() 。

代码布局

  • register.cpp:设置pybind绑定并调用TVM后端的注册。
  • compiler.{h,cpp}:用TVM编译PyTorch JIT图的主要逻辑。
  • operators.{h,cpp}:从JIT IR映射到TVM操作符的位置。

TODO

  • 添加从Python中将不透明op名称的翻译注册到TVM中(如在operator.cpp中完成)的功能。
  • 零拷贝 set_input
  • 纾困机制(调用PyTorch JIT后备)
  • Threadpool 集成
  • 分配器集成
  • 操作符翻译

    卷积

    BatchNorm

    RELU

    AveragePool

    MaxPool

    线性

  • 张量操作

    重塑

    查看

原文发布于微信公众号 - AI研习社(okweiwu)

原文发表时间:2019-05-05

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

扫码关注云+社区

领取腾讯云代金券