在 as_strided 之前进行整形以进行优化

Question

def forward(x, f, s):
    B, H, W, C = x.shape # e.g. 64, 16, 16, 3
    Fh, Fw, C, _ = f.shape # e.g. 4, 4, 3, 3 
    # C is redeclared to emphasise that the dimension is the same
    
    Sh, Sw = s # e.g. 2, 2

    strided_shape = B, 1 + (H - Fh) // Sh, 1 + (W - Fw) // Sw, Fh, Fw, C

    x = as_strided(x, strided_shape, strides=(
        x.strides[0], 
        Sh * x.strides[1], 
        Sw * x.strides[2], 
        x.strides[1], 
        x.strides[2], 
        x.strides[3]), 
    )

    # print(x.flags, f.flags)

    # The reshaping changes the einsum from 'wxyijk,ijkd' to 'wxyz,zd->wxyd'
    f = f.reshape(-1, f.shape[-1])
    x = x.reshape(*x.shape[:3], -1) # Bottleneck!
    
    return np.einsum('wxyz,zd->wxyd', x, f, optimize='optimal')

（相反，变体 without reshapes 使用return np.einsum('wxyijk,ijkd->wxyd', x, f)）

作为参考，这里是 x 和 f 在重塑之前的标志：

x.flags:

C_CONTIGUOUS : False
F_CONTIGUOUS : False
OWNDATA : False
WRITEABLE : True
ALIGNED : True
WRITEBACKIFCOPY : False
UPDATEIFCOPY : False


f.flags:

C_CONTIGUOUS : True
F_CONTIGUOUS : False
OWNDATA : True
WRITEABLE : True
ALIGNED : True
WRITEBACKIFCOPY : False
UPDATEIFCOPY : False

有趣的是，例程中的主要瓶颈不是einsum，而是x 的重塑（扁平化）。我知道f 不会遇到此类问题，因为它的内存是 C 连续的，因此重塑相当于在不更改数据的情况下进行快速内部修改 - 但由于 x 不是 C 连续的（并且不拥有它的数据，就此而言），重塑要昂贵得多，因为它涉及经常更改数据/获取非缓存对齐的数据。这反过来又是由在x 上执行的as_strided 函数导致的——步幅的修改必须以扰乱自然顺序的方式进行。 （仅供参考，as_strided 的速度非常快，无论跨度如何，都应该很快）

有没有办法在不产生瓶颈的情况下达到相同的结果？也许通过在使用as_strided 之前重塑x？

---

另请注意，对于几乎 100% 的应用程序： B：[1-64]，H，W：[1-60]，C：[1-8] 前锋，前锋：[1-12]

---

我还在这里包含了一些图表，用于随着张量尺寸 B（批量大小）以及我设备上 H, W（图像大小）的变化而变化的时间（如您所见，涉及 reshape 的那个已经与 Tensorflow 具有相当的竞争力）：

---

编辑：一个有趣的发现 - 重塑算法在 CPU 上以 5 倍击败非重塑算法，但是当我使用 GPU（即使用 CuPy 而不是 NumPy）时，两种算法都同样快（大约是 TensorFlow 前向传播速度的两倍）

Answer 1

由于您提到的原因（在非连续数组上复制），重新调整跨步数组的成本有点高，但并不像您想象的那么昂贵。 np.einsum 实际上可能是应用程序的瓶颈，具体取决于张量大小。正如Convolutional layer in Python using Numpy 中提到的，np.tensordot 可以很好地替代np.einsum。

只是给你一个简单的例子：

x = np.arange(64*221*221*3).reshape((64, 221, 221, 3))
f = np.arange(4*4*3*5).reshape((4, 4, 3, 5))
s = (2, 2)

B, H, W, C = x.shape # e.g. 64, 16, 16, 3
Fh, Fw, C, _ = f.shape # e.g. 4, 4, 3, 3 
Sh, Sw = s # e.g. 2, 2
strided_shape = B, 1 + (H - Fh) // Sh, 1 + (W - Fw) // Sw, Fh, Fw, C
print(strided_shape)
# (64, 109, 109, 4, 4, 3)

初始化变量后，我们可以测试代码部分的时序

%timeit x_strided = as_strided(x, strided_shape, strides=(x.strides[0], Sh * x.strides[1], Sw * x.strides[2], x.strides[1], x.strides[2], x.strides[3]), )
>>> 7.11 µs ± 118 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit f_reshaped = f.reshape(-1, f.shape[-1])
>>> 450 ns ± 7.43 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%timeit x_reshaped = x_strided.reshape(*x_strided.shape[:3], -1) # Bottleneck!
>>> 94.6 ms ± 896 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# einsum without reshape
%timeit np.einsum('wxy...,...d->wxyd', x_strided, f, optimize='optimal')
>>> 809 ms ± 1.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# einsum with reshape
%%timeit
f_reshaped = f.reshape(-1, f.shape[-1])
x_reshaped = x_strided.reshape(*x_strided.shape[:3], -1) # Bottleneck!
k = np.einsum('wxyz,zd->wxyd', x_reshaped, f_reshaped, optimize='optimal')
>>> 549 ms ± 3.05 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# tensordot without reshape
%timeit k = np.tensordot(x_strided, f, axes=3)
>>> 271 ms ± 4.89 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# tensordot with reshape
%%timeit
f_reshaped = f.reshape(-1, f.shape[-1])
x_reshaped = x_strided.reshape(*x_strided.shape[:3], -1) # Bottleneck!
k = np.tensordot(x_reshaped, f_reshaped, axes=(3, 0))
>>> 266 ms ± 3.15 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

我在您的代码中使用张量大小（即 64、16、16、3 和 4、4、3、3）得到了类似的结果。

如您所见，调整大小操作存在开销，但由于数据连续，它使矩阵操作更快。请注意，结果会因 CPU 速度、CPU 架构/代数等而异。