welcome to my blog
pytorch中的广播机制和numpy中的广播机制一样, 因为都是数组的广播机制
两个维度不同的Tensor可以相乘, 示例
如何理解数组的广播机制
以数组A和数组B的相加为例, 其余数学运算同理
核心:如果相加的两个数组的shape不同, 就会触发广播机制, 1)程序会自动执行操作使得A.shape==B.shape, 2)对应位置进行相加
运算结果的shape是:A.shape和B.shape对应位置的最大值,比如:A.shape=(1,9,4),B.shape=(15,1,4),那么A+B的shape是(15,9,4)
有两种情况能够进行广播
- A.ndim > B.ndim, 并且A.shape最后几个元素包含B.shape, 比如下面三种情况, 注意不要混淆ndim和shape这两个基本概念
- A.shape=(2,3,4,5), B.shape=(3,4,5)
- A.shape=(2,3,4,5), B.shape=(4,5)
- A.shape=(2,3,4,5), B.shape=(5)
- A.ndim == B.ndim, 并且A.shape和B.shape对应位置的元素要么相同要么其中一个是1, 比如
- A.shape=(1,9,4), B.shape=(15,1,4)
- A.shape=(1,9,4), B.shape=(15,1,1)
下面分别进行举例
A.ndim 大于 B.ndim
A.ndim 等于 B.ndim