Calculating Mean & STD for Batch [Python/Numpy]

Question

Looking to calculate Mean and STD per channel over a batch efficiently.

---

Details:

• batch size: 128
• images: 32x32
• 3 channels (RGB)

So each batch is of size [128, 32, 32, 3].

There are lots of batches (naive method takes ~4min over all batches).

And I would like to output 2 arrays: (meanR, meanG, meanB) and (stdR, stdG, stdB)

---

(Also if there is an efficient way to perform arithmetic operations on the batches after calculating this, then that would be helpful. For example, subtracting the mean of the whole dataset from each image)

Answer 1

If I understood you correctly and you want to calculate mean and std values for all images:

Demo: 2 images of (2,2,3) shape each (for the sake of simplicity):

In [189]: a
Out[189]:
array([[[[ 1,  2,  3],
         [ 4,  5,  6]],

        [[ 7,  8,  9],
         [10, 11, 12]]],


       [[[13, 14, 15],
         [16, 17, 18]],

        [[19, 20, 21],
         [22, 23, 24]]]])

In [190]: a.shape
Out[190]: (2, 2, 2, 3)

In [191]: np.mean(a, axis=(0,1,2))
Out[191]: array([ 11.5,  12.5,  13.5])

In [192]: np.einsum('ijkl->l', a)/float(np.prod(a.shape[:3]))
Out[192]: array([ 11.5,  12.5,  13.5])

Speed measurements:

In [202]: a = np.random.randint(255, size=(128,32,32,3))

In [203]: %timeit np.mean(a, axis=(0,1,2))
9.48 ms ± 822 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [204]: %timeit np.einsum('ijkl->l', a)/float(np.prod(a.shape[:3]))
1.82 ms ± 22.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 2

Assume you want to get the mean of multiple axis(if I didn't get you wrong). numpy.mean(a, axis=None) already supports multiple axis mean if axis is a tuple.

I'm not so sure what you mean by naive method.