How to Implement Shuffle BN for MoCo in Tensorflow 2

Hi, hope I can get some help here. I want to implement unsupervised contrastive learning model MoCo in TF2, but I have no idea how to implement the essential trick mentioned in the paper - Shuffling BN. I think I understand what shuffling BN does, but I don’t know any APIs to fetch different data slices from each GPU, shuffle them, and send new slices back to each GPU.

Hi @Zhihua_Che

Welcome to the TensorFlow Forum!

The given information is not enough to know about the issue. Please share minimal reproducible code to replicate and understand the issue. Thank you.