I am using ddp to train my model now. In my model I want to calculate the standard deviation across the batch. However since I am wondering if i can calculate the standard deviation across the entire batch instead of within each device. The standard deviation will be part of my computation graph.
I feel that this is similar to synchronize batchnorm and should be doable. How would I go about doing this? Here is an example of what I want to do
def forward(self, input):
        feats = conv(input)
        batch, channel, height, width = feats.shape
        stddev = feats.view(batch, -1, 1, channel, height, width)
        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
        stddev = stddev.repeat(batch, 1, height, width)
        feats = torch.cat([feats, stddev], 1)
        output = conv_last(feats)
        return output

Basically when I compute stddev, I want it to do it over entire batch.