In [1]:
import numpy as np

channels_last = np.random.random((150, 150, 40))
channels_first = np.random.random((40, 150, 150))
In [2]:
def gradient_channels_first(image_data):
    g = np.empty((image_data.shape[0] * 2,) + image_data.shape[1:])
    for i in xrange(image_data.shape[0]):
        gx, gy = np.gradient(image_data[i])
        g[2 * i] = gx
        g[2 * i + 1] = gy
    return g
In [3]:
def gradient_channels_last(image_data):
    g = np.empty(image_data.shape[:2] + (image_data.shape[2] * 2,))
    for i in xrange(image_data.shape[2]):
        gx, gy = np.gradient(image_data[..., i])
        g[..., 2 * i] = gx
        g[..., 2 * i + 1] = gy
    return g
In [5]:
%timeit gradient_channels_first(channels_first)
100 loops, best of 3: 13 ms per loop
In [6]:
%timeit gradient_channels_last(channels_last)
10 loops, best of 3: 22.7 ms per loop