极智AI | 谈谈算子开发 - 广播机制 boradcast

Posted 2022-12-10 极智视界

  欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享

  大家好，我是极智视界，本文介绍一下 谈谈算子开发 - 广播机制 boradcast。

  广播机制在 reduce机制 的算子开发的时候可以说十分常见，所谓 reduce机制 是指像 sum、add、concat之类的输入是多头的，输出是单头类型的算子。拿 sum(a, b) 来说，就需要 a 和 b 具有相同的 shape，这时就是广播机制发挥作用的时候了。

  那么什么是广播机制呢？广播机制主要功能是 升维，可以将一个数组的每一个维度扩展为一个固定的 shape，需要被扩展的数组的每个维度的大小或者与目标 shape 相等，或者为1，广播会在元素个数为 1 的维度上进行。例如，原来的数组为 (1, 3, 512, 512)，目标 shape 为 (32, 3, 512, 512)，则可以通过广播将 a 的维度扩展为目标 shape。

  再举个具体的例子，比如有 Tensor A 的 shape 为 (32, 3, 1, 7)，Tensor B 的 shape 为 (1, 1, 5, 1)，你想执行 sum(Tensor A，Tensor B)，那么内部会是怎么样一个执行过程呢？

• 一般的思路就是：先把 Tensor A 和 Tensor B 的维度弄成一致的，然后再进行加法。一致一般就是指每个维度上的 max，也即 (32, 3, 5, 7)；
• 首先把 Tensor A (32, 3, 1, 7) -> (32, 3, 5, 7)，也即在 axis=2 上从 1 广播至 5；
• 然后把 Tensor B (1, 1, 5, 1) -> (32, 3, 5, 7)，也即在 axis=0 上从 1 广播至 32，在 axis=1 上从 1 广播至 3，在 axis=3 上从 1 广播至 7；
• 然后再将广播后的 Tensor A 和 Tensor B 相加。

  这里需要注意的是，如果需要广播的 tensor 的 shape 不满足每一个维度 (同一个轴 axis) 的大小与目标 shape 相等或者为 1 的要求，则无法进行广播。比如 Tensor C (32, 3, 3, 7) 与 Tensor D (1, 1, 5, 1)，这个时候的 Tensor C 和 Tensor D 就无法进行广播了。

  当然在用的时候你也许并不会察觉有什么广播操作，这是因为有些框架、有些接口它内部给你实现好了 自动广播机制，所以这个时候的广播，对于用户来说是不感知的，这也没有什么特别好疑惑的。

  好了，以上分享了 谈谈算子开发 - 广播机制 boradcast，希望我的分享能对你的学习有一点帮助。

 【极智视界】

《极智AI | 谈谈算子开发 - 广播机制 boradcast》

搜索关注我的微信公众号【极智视界】，获取我的更多经验分享，让我们用极致+极客的心态来迎接AI ！