Чистая NumPy 2D средняя производная от свертки входного изображения

Question

Я имею b 2d m x n полутоновые изображения, которые я сворачиваю с p x q фильтр, а затем делать скупой на. С чисто NumPy я хотел бы вычислить производную входного изображения и фильтра, но у меня возникли проблемы с вычислением производной входного изображения:

def conv2d_derivatives(x, f, dy):
    """
    dimensions:
        b = batch size
        m = input image height
        n = input image width
        p = filter height
        q = filter width
        r = output height
        s = output width

    input:
        x = input image                       (b x m x n)
        f = filter                            (p x q)
        dy = derivative of some loss w.r.t. y (b x r x s)

    output:
        df = derivative of loss w.r.t. f      (p x q)
        dx = derivative of loss w.r.t. x      (b x m x n)

    notes:
        wx = windowed version of x s.t. wx[b, r, s] = the window of x to compute y[b, r, s]
        vdx = a view of dx 
    """
    b, m, n = x.shape
    p, q = f.shape
    r = m - p + 1
    s = n - q + 1
    wx = as_strided(x, (b, r, s, p, q), np.array([m * n, 1, q, 1, n]) * x.itemsize)

    # This derivative is correct
    df = 1 / (p * q) * np.einsum('brspq,brs->pq', wx, dy)

    # Method 1: this derivative is incorrect
    dx = np.zeros_like(x)
    vdx = as_strided(dx, (b, r, s, p, q), np.array([m * n, 1, q, 1, n]) * dx.itemsize)
    np.einsum('pq,brs->brspq', f, dy, out=vdx)
    dx /= (p * q)

    # Method 2: this derivative is correct, but it's slow and memory-intensive
    dx = np.zeros_like(x)
    vdx = as_strided(dx, (b, r, s, p, q), np.array([m * n, 1, q, 1, n]) * dx.itemsize)
    prod = f[None, None, None, :, :] * dy[:, :, :, None, None]
    for index in np.ndindex(*vdx.shape):
        vdx[index] += prod[index]
    dx /= (p * q)

    return df, dx

Я знаю, что производная убытка по отношению к w[b,r,s,p,q] просто 1/(p*q) * f[p,q] * dy[b,r,s], Однако я не хочу явно вычислять производные для w и хранить их в памяти, потому что этот массив будет огромным.

Я думал, что мог бы сделать Einsum зрения dx, vdxпохож на оконный wdxи надеюсь, что Einsum будет увеличиваться vdx[b,r,s,p,q] += f[p,q] * dy[b,r,s], но это на самом деле назначает vdx[b,r,s,p,q] = f[p,q] * dy[b,r,s], Если бы был способ указать out_add_to в einsum, тогда моя проблема будет решена.

Как мне вычислить dx без хранения большого b x r x s x p x q матрица в чистом NumPy? Я не могу использовать scipy или любую другую зависимость для этой проблемы.