Get Batch statistics over multiple GPUs

Hey Folks,

I was looking to implement class imbalanced loss for a classification model, where the cross entropy loss would be weighted with class frequency.(Inverse frequency). Now, Since I want to implement this at batch level, I need to know the frequency of each class in all the batches. But at a given time, i know this information one a given gpu only. Is there a way to get this information, during runtime.

You can use hooks. They are called at the time of computing backward or forward. If you use multi-gpu they will be called once for each gpu.