Freezing Batch norm layers while keeping them in the graph

amey: <p>Hi<br>
If we set requires_grad to False for batch norm layers of a model, the batch norm layers do not remain in the graph. In this case, I cant fine tune these layers later if I want to. Is there any way we can freeze the layers, yet keep them in the graph so that they can be trained later?</p>

ptrblck: <p>You could set the requires_grad attributes of the batchnorm parameters to True again if you want to train them later.</p>

amey: <p>Thank you very much. I tried this. Setting them to requires_grad = False in the initial training causes the batch norm layers to not train and they later can be trained for finetuning with requires_grad = True.<br>
However, I noticed that even if the batch norm layers are not being trained with requires_grad = False, I think all the BN gradient calculations are still happening and a lot of memory is being occupied. In fact, I cant train a complete model because I am getting out of memory error very often. I am guessing if there is a way to shut down the mean and variance calculations.<br>
BTW, I was using the SyncBatchNorm class because I am using 2 GPUs for training with DataParallel.</p>

ptrblck: <p>The gradients might still be needed if previous layers depend on them.<br>
To disable the stats updates, call <code>.eval()</code> on the batchnorm layers or use <code>track_running_stats=False</code> if you want to use the batch stats during training and evaluation.</p>

amey: <p>Thanks. I will try this asap. Meanwhile, I was wondering if there is a way to set the batch stats - mean and variance to 0s and 1s so that all gradients are based on these values and we altogether avoid calculating the means and variances(which might be taking up memory)?<br>
Also, I was following your manual implementation of the batch normalization layer(<a href="https://github.com/ptrblck/pytorch_misc/blob/master/batch_norm_manual.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch_misc/batch_norm_manual.py at master · ptrblck/pytorch_misc · GitHub</a>). Could you tell me if it supports multi GPUs?</p>

ptrblck: <aside class="quote no-group" data-username="amey" data-post="5" data-topic="141420">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/amey/48/45954_2.png" class="avatar"> amey:</div>
<blockquote>
<p>Meanwhile, I was wondering if there is a way to set the batch stats - mean and variance to 0s and 1s so that all gradients are based on these values and we altogether avoid calculating the means and variances(which might be taking up memory)?</p>
</blockquote>
</aside>
<p>The stats will be initialized with these values, so you could call <code>.eval()</code> directly on the batchnorm layers after initializing the model.<br>
However, note that freezing the affine parameters and disabling the normalization via <code>.eval()</code> would mean that you are basically adding “noise” to the model and I’m unsure if that’s your use case or why you would want to do it.</p>
<aside class="quote no-group" data-username="amey" data-post="5" data-topic="141420">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/amey/48/45954_2.png" class="avatar"> amey:</div>
<blockquote>
<p>Could you tell me if it supports multi GPUs?</p>
</blockquote>
</aside>
<p>It might work with e.g. <code>DDP</code>, but note that I’ve written it more for educational reasons to play around with the manual implementation in Python directly.</p>