Question about the moco implementation

[mingkai-zheng]

Hello, I'm a bit confused about the moco implementation in this paper. Since moco only has one forward pass for the teacher network, so I guess that the lazy update is not required for moco right? In this case, did you include the bn statistics for the current batch during the forward pass?

To be more specific, do you update the running_mean and running_var before calculating x?

with torch.no_grad():
    self.running_mean = self.momentum * mean + (1 - self.momentum) * self.running_mean
    self.running_var = self.momentum * var * n / (n - 1) + (1 - self.momentum) * self.running_var

x = (x - self.running_mean[None, :, None, None].detach()) / (
    torch.sqrt(self.running_var[None, :, None, None].detach() + self.eps)
)

or you calculate x first

x = (x - self.running_mean[None, :, None, None].detach()) / (
    torch.sqrt(self.running_var[None, :, None, None].detach() + self.eps)
)

with torch.no_grad():
    self.running_mean = self.momentum * mean + (1 - self.momentum) * self.running_mean
    self.running_var = self.momentum * var * n / (n - 1) + (1 - self.momentum) * self.running_var