mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
a3d495bd4e
This PR fixes typos in `.md` and `.rst` files under functorch directory Pull Request resolved: https://github.com/pytorch/pytorch/pull/87663 Approved by: https://github.com/kit1980
49 lines
1.9 KiB
ReStructuredText
49 lines
1.9 KiB
ReStructuredText
Patching Batch Norm
|
|
===================
|
|
|
|
What's happening?
|
|
-----------------
|
|
Batch Norm requires in-place updates to running_mean and running_var of the same size as the input.
|
|
Functorch does not support inplace update to a regular tensor that takes in a batched tensor (i.e.
|
|
``regular.add_(batched)`` is not allowed). So when vmaping over a batch of inputs to a single module,
|
|
we end up with this error
|
|
|
|
How to fix
|
|
----------
|
|
All of these options assume that you don't need running stats. If you're using a module this means
|
|
that it's assumed you won't use batch norm in evaluation mode. If you have a use case that involves
|
|
running batch norm with vmap in evaluation mode, please file an issue
|
|
|
|
Option 1: Change the BatchNorm
|
|
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
|
If you've built the module yourself, you can change the module to not use running stats. In other
|
|
words, anywhere that there's a BatchNorm module, set the ``track_running_stats`` flag to be False
|
|
|
|
.. code-block:: python
|
|
|
|
BatchNorm2d(64, track_running_stats=False)
|
|
|
|
|
|
Option 2: torchvision parameter
|
|
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
|
Some torchvision models, like resnet and regnet, can take in a ``norm_layer`` parameter. These are
|
|
often defaulted to be BatchNorm2d if they've been defaulted. Instead you can set it to BatchNorm
|
|
that doesn't use running stats
|
|
|
|
.. code-block:: python
|
|
|
|
import torchvision
|
|
from functools import partial
|
|
torchvision.models.resnet18(norm_layer=partial(BatchNorm2d, track_running_stats=False))
|
|
|
|
Option 3: functorch's patching
|
|
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
|
functorch has added some functionality to allow for quick, in-place patching of the module. If you
|
|
have a net that you want to change, you can run ``replace_all_batch_norm_modules_`` to update the
|
|
module in-place to not use running stats
|
|
|
|
.. code-block:: python
|
|
|
|
from functorch.experimental import replace_all_batch_norm_modules_
|
|
replace_all_batch_norm_modules_(net)
|