Pytorch学习（十一）--- Tensor Comprehensions 初探

最新推荐文章于 2025-05-12 10:52:16 发布

原创

最新推荐文章于 2025-05-12 10:52:16 发布 · 2.3k 阅读

7 ·

CC 4.0 BY-SA版权

文章标签：

#Ｄeep Learning #pytorch

初衷

其实, 有时候你会很难受. 因为你有一些想法, 在python端写会很慢, 写成cuda代码又有难度, 还要想着各种优化. 这时候你就不爽了,好不容易想到骚操作, 竟然因为写不来代码, 就泯灭这个想法吗. 这时候, 你可以看看Tensor Comprehesions(TC)这个包.

先吹一波: TC是一个让你不用写高性能代码的包, 它会直接根据简单的语法来生成GPU代码.
如果你还在为一下事情烦恼:
- 当你的pytorch 层很慢, 然后你想写CUDA代码, 这时候你没必要真的写CUDA代码.
- 你有一个CUDA层, 你花了一周写啊,调试啊,优化啊.但是你现在想一小时搞定.
- 其他..

使用

安装

这个没啥好说的:

conda install -c pytorch -c tensorcomp tensor_comprehensions

大概会安装这些包:

   package                    |            build
    ---------------------------|-----------------
    protobuf-3.4.1             |       h21cfbc1_2         5.6 MB  tensorcomp
    llvm-tapir50-0.1.0         |       h186cc49_2       331.5 MB  tensorcomp
    cudatoolkit-8.0            |                3       322.4 MB
    rhash-1.3.5                |       hbf7ad62_1         178 KB
    unzip-6.0                  |       h611a1e1_0          91 KB
    halide-0.1.0               |       h9df8326_2        37.9 MB  tensorcomp
    isl-tc-0.1.0               |       h9c8d533_2         1.3 MB  tensorcomp
    libuv-1.19.2               |       h14c3975_0         713 KB
    pytorch-0.3.1              |py36_cuda8.0.61_cudnn7.0.5_2       205.1 MB  pytorch
    gfla