Problems for USBatchNorm2d and fedalign loss

[calvin89029]

Hi, thanks for the great work.
However, I have some problems while reading code. Not sure if I miss something.

In slimmable_ops.py line 146-147, should they by

self.running_mean[:c] 
self.running_var[:c]

By the way, for line 122, 136, and 137, I can't figure out why you need an additional self.bn=nn.BatchNorm2d here?

[calvin89029]

One more question. In methods/fedalign.py line 54, not sure why you multiply (loss_CE/loss.item()) before loss?

loss = self.args.mu*(loss_CE/loss.item())*loss

[mmendiet]

Hello @calvin89029! Thank you for your interest.
Here, track_running_stats=False, and therefore self.running_mean and self.running_var are None. This means that the batchnorm operation on lines 144-152 will use batch statistics, which is what we want for the subnetworks. For the full network, we maintain the running mean and variance, which is why we have self.bn (track_running_stats=True by default). This is based on the repos I referenced here. Nonetheless, it could probably be simplified by setting track_running_stats=True here, remove self.bn, replace lines 136 and 137 with

self.running_mean[:c] 
self.running_var[:c]

And replace lines 146 and 147 with

None,
None,

But in the end it should be the same :)
For your second question, I discuss this in Section E of the supplemental material. Basically, it's just to ensure that loss is on relatively the same scale with loss_CE.
Hope this helps!

[calvin89029]

Hi @mmendiet, thank you for the reply.
For the first question, yeah I think that will be much clearer.
For the second one, I just reimplement your work and train on some other dataset. However, if I train without gradient clip and loss_CE scaling, the Lipschitz loss seems to get larger and larger along the training process and finally explode. Did you encounter this problem too? And is this the reason why you add these two tricks?