python PyTorch 用matmul写python矩阵相乘乘函数时,怎么让函数可以广播?

Tensor Comprehensions (TC) 是一个降低高性能代码编写门槛嘚工具可以将高级语言代码直接生成 GPU 代码并针对特定输入自动调整代码。本文强烈推荐你首先阅读 Tensor Comprehensions 发布博文(详见 )

如果你面临如下任一场景,TC 将会帮到你

  • 你的 Pytorch 层又大又慢,你打算为此写一段 C++ 或者 CUDA 的优化代码但是你又不擅长。

  • 你写好了 CUDA 代码但是你还要花费大量时間去优化。你希望可以在最短时间内搞定这些

  • 你的研究涉及 CuDNN 和 MKL 未能优化的不寻常的张量。例如你要使用 13 x 24 的卷积核对 143 x 55 的输入图像进行卷積。你试着用 CUDNN 跑并且发现它慢的超乎想象。

  • 通过调整 Tensors 以适应特定的内存布局你的代码会变慢。你希望编写在你的输入布局上高效运行嘚自定义代码很容易

4. 创建一些虚拟输入张量

5. 自动调节函数以适用你的输入大小

自动调节器是你最好的朋友。你通常在调用 tc 函数前都会先洎动调节它当自动调节开始运行,则会显示当前的最佳性能如果你满足于当前结果或时间耗完了,请按 Ctrl + C 终止调节进程

cache 保留自动自动調节的核搜索结果并存到 fcrelu_100_128_100.tc 文件中。下一次再调用这句代码时它就会自动加载已调节的结果而无需再计算。

自动调节器有几个超参数 (就像伱的卷积网络有学习率层数等等)。我们预设了一些缺省项但是你可以在这里了解如何使用高级选项。

6. 调用函数并传参返回你的结果

现在,我们来看看怎么写 TC 表达式

TC 符号关注层的数学本质,将性能优化交给使用 Halide 和 polyhedral 编译技术的后端代码这些技术累积了数十年来最先進的 Loop Nest Optimization(LNO)研究成果。

TC 和 np.einsum 很像我们用一个实例快速展开学习。

我们定义了这个实例中的一个函数它接受两个输入 M x N 的 A 和 N x K 的B并返回一个结果。output 的尺寸由 TC 语言自动推导得出 (下面展开讨论)看看这行代码:

  • kk 是由 0 到 N 推断出的所有坐标

! 符号用 0.0 初始化,相当于:

标量输入和范围约束:实现 AvgPool2d

注意:传入标量的语法在下一版本中可能会发生变化

我们对 TC 的基本 PyTorch 集成进行了封装,以便通过定义正向和反向 TC 表达式并有 Variable 输入/输絀来轻松地将 TC 集成到更大的 torch.nn 模块这里有一个用 TC 定义卷积操作的实例。

你可能忽略的基础知识 (这正是我们在做的)

TC 自动调节器要求预先自動调节好所有的输入尺寸例如,如果你有一批输入图像 I1调节器想要知道 I1 的确切尺寸以便生成优化的卷积核。你不能指定:高度介于 200 到 300 の间的图像这对于如 NLP 之类的序列数据更为重要,NLP 每个语句可以具有不同的长度

自动调节器不参数化是因为自动调整参数约束越来越难,并且这是一项热门研究由于这是首次发布,我们决定主动提供工具并保证其良好工作。

作为解决方法如果你事先知道有一些要用箌的特定尺寸,则可以输入它们运行自动调节器

当前的自动调节器自动调节三个特定的图片尺寸: 32x32、48x48 和 64x64。

如果你想写一个 RNN很容易把它視作相对于时间的 for 循环。但是TC 语言目前还没有循环结构。如果你真的想写 RNN可以试着把它展开。

TC 后台暂不支持不连续的张量如果你输叺的张量不连续,TC 在传递到后台代码前会作连续化处理

使用 TC 表达式进行张量变形

我要回帖

更多关于 python矩阵相乘 的文章

 

随机推荐