Late reply but in case anyone else is looking for a solution to quickly collect statistics I’ve created some code to do this at https://gist.github.com/thomasbrandon/ad5b1218fc573c10ea4e1f0c63658469.
You can run it off any iterator that returns tensors. Something like:
>>> DATA = untar_data(URLs.MNIST_SAMPLE)
... src = (ImageList.from_folder(DATA)
... .split_by_folder(valid='valid'))
... stats = collect_stats(src.train)
... stats
RunningStatistics(n=9718464, mean=[0.128,0.128,0.128], std=[0.305,0.305,0.305])
That not split so src.train, you’d use src.train.x after splitting or you can do it from a databunch with data.train_ds.x. It defaults to collapsing the last 2 dimensions, as appropriate for image data. You can pass n_dims to change this and generate stats of arbitrary channels of arbitrary shapes.
Performance on the code above was [12396/12396 00:03<00:00] - so 3secs for ~12000 28x28 images (off a reasonably fast NVMe and having run before so probably largely cached, disk IO obviously a limit). Numerical stability seems good, it will be off from the true value by a bit but is stable (assert_allclose(rtol=0.001, atol=0.01) across 4000 batches of randn).