Computing std over each image channel

Hello, I am trying to compute the std for an image, but I want a different std for each channel. Suppose we have an image with the shape [channel,x,y], in numpy we can achieve this by using multi-axis, like: img.std(axis=(1,2))

Unfortunately, the torch std implementation does not support multi-axis, so I tryed to implement my own function:

def std_for_channel(img):
   if (torch.typename(img) == 'torch.autograd.variable.Variable'):
      num_channels = img.size()[0]
   else:
      num_channels = img.shape[0]
   out = torch.zeros(num_channels)
   for i in range(num_channels):
      out[i] = img[i,:,:].std(unbiased=False)
   return out

while this works when passing a tensor, I keep getting error when using this on a Variable (getting ‘RuntimeError: can’t assign a Variable to a scalar value of type float’). Since I want to use this std information in my custom loss, how can I implement it correctly so backpropagation will be able to train the weights of the network?

1 Like

img.contiguous().view(img.size(0), -1).std(-1)

3 Likes

Thank you very much.

Sorry for bumping this, but for anyone who comes across this and it seems like magic…

img.contiguous() makes sure the tensor is in a contiguous chunk of memory for efficient access.
.view(img.size(0), -1) resizes from [c, x, y, z…etc] to [c, n] where n = xyz…
.std(-1) takes the standard deviation when you collapse the last channel so you end up with the channel mean.

If the last command is confusing, it’s because pytorch dims can be confusing, here is a nice explanation
Understanding dimensions in pytorch

1 Like