How to simplify function for standardising images?
I have a function to calculate the mean and standard deviation of my dataset.
Is there a simpler way to do this? As it takes a while to compute.
def get_mean_std(loader):
  sum = 0
  sum_sq_err = 0
  for data, _ in loader:
    sum += torch.mean(data, dim=[0,2,3]) 
    sum_sq_err += torch.mean(data**2, dim=[0,2,3])
  mean = sum/len(loader)
  std = (sum_sq_err/(len(loader)) - mean**2)**0.5
  return mean, std


Note that this approach is not even correct in general, as the mean of a set is not the mean of the means of some subsets in general (it is when all the subsets have the same length, but that may or may not be the case here).
Provided that every batch is of the same size, what I would do is to call torch.sum in the loop, but rather than already accumulating it into a sum, appending it into a list, and then reducing it via torch.sum + a division afterwards. Note that torch.sum implements a highly non-trivial algorithm that is more precise in general than the naïve iterative sum.
