Repository: pytorch/functorch (language: python)

Issue:
Why is using vmap(jacrev) for BatchNorm2d in non-tracking mode not working?
Hi, experts.
I am trying to use vmap(jacrev) to calculate the per-sample jacobian in a batch for my network during inference. However, when there is BatchNorm2d, it does not work. Because during inference, BatchNorm2d is simply applying the statistics previously tracked (and not doing any inter-sample operations), I think it should work just as any other simple operation from my understanding. Is there a way for me to make it work, or is there anything I am misunderstanding?

Below is my minimal code:
```
from functorch import jacrev, vmap
import torch
from torch import nn

layers = nn.Sequential(
        nn.Conv2d(3, 3, kernel_size=(3, 3)),
        nn.BatchNorm2d(3, track_running_stats=False),
    )

x = torch.randn(4, 3, 30, 30)
j = vmap(jacrev(layers))(x)
```

And I get this error in the bn layer
`ValueError: expected 4D input (got 3D input)`

I think this should fundamentally be doable, and just might be because of how vmap and jacrev is implemented.
Is there any simple workaround, or am I misunderstanding anything?

Thank you for any help


Discussion / hints:
Hi @kwmaeng91! 

What vmap is doing here is assuming that the first dimension is like a batch dimension that other operations shouldn't be seeing. So, BatchNorm2d sees a tensor of size `(3, ConvOut, ConvOut)` instead of `(4, 3, ConvOut, ConvOut)` which produces this error.

Conv2d has a unbatched case (allowing for 3D cases where it's assumed that there's only one sample) that batch norm doesn't have, so that's why batch norm is getting the error instead of Conv2d. To fix this, we can introduce a custom forward:

```
class Custom(nn.Module):
  def __init__(self):
    super().__init__()
    self.conv = nn.Conv2d(3, 3, kernel_size=(3, 3))
    self.bn = nn.BatchNorm2d(3, track_running_stats=False)
    
  def forward(self, x):
    x = self.conv(x)
    x = x.unsqueeze(0)  # tensor will now be of size (1, 3, ConvOut, ConvOut)
    return self.bn(x)
```
Dear @samdow, Thank you for your reply! It is really helpful.
I have a followup question. Now it works with `track_running_stats=False` with your help, but it still does not with `track_running_stats=True`, even when the model is in eval mode (i.e., `net.eval()`). I thought something similar will work as well for eval(), but I guess it does not?

From here (https://pytorch.org/functorch/stable/batch_norm.html), it says "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)." as the reason why this does not work. However, doesn't running_mean and running_var **not get updated in eval() mode**?? I am not sure why vmap cannot work with eval() mode.

Will there be yet another fix for `track_running_stats=True` with eval mode (i.e., `net.eval()`, or is this fundamentally not doable right now unlike my original question?
I can make this a separate issue if that makes things more clear.

Thank you for the help!
Hi @kwmaeng91! Great catch! I believe that this is a bug that we should fix. I may not have bandwidth to do this for a bit, but given your use case, I'll try to prioritize it. The fix itself should be simple but we'll want to test it and the way we're testing batch norm is a little subtle to avoid the error case
Hi, @samdow, Thank you for the reply! I will be awaiting your fix :) In the meantime, is there any quick fix I might want to try? I wonder if I just comment out the assertion it would work for me.
If you are able to building from source, adjusting this check [here](https://github.com/pytorch/functorch/blob/main/functorch/csrc/BatchRulesNorm.cpp#L61) to do
`if (training && (!input_bdim || <rest of the check as it appears right now>))`

it should work and you should no longer see the assertion. If you can't build from source, please let me know and I think our CI builds wheels and I can try to make a PR with the change so you can have access to it while I work on fixing the testing 
@samdow Thank you for the help. I have recompiled the modified source successfully. Still, it would be great if your fix gets committed in the future.

Fix patch:
diff --git a/functorch/csrc/BatchRulesNorm.cpp b/functorch/csrc/BatchRulesNorm.cpp
index c5b273652..e78538329 100644
--- a/functorch/csrc/BatchRulesNorm.cpp
+++ b/functorch/csrc/BatchRulesNorm.cpp
@@ -58,7 +58,7 @@ batch_norm_batch_rule(
   auto running_mean = *running_mean_maybe_owned;
   c10::MaybeOwned<Tensor> running_var_maybe_owned = at::borrow_from_optional_tensor(running_var_opt);
   auto running_var = *running_var_maybe_owned;
-  TORCH_CHECK(!input_bdim || ((!running_mean.defined() || running_mean_bdim) && (!running_var.defined() || running_var_bdim)),
+  TORCH_CHECK(!training || (!input_bdim || ((!running_mean.defined() || running_mean_bdim) && (!running_var.defined() || running_var_bdim))),
       "Batch norm got a batched tensor as input while the running_mean or running_var, which will be updated in place, ",
       "were not batched.\nIf you are using a module and do not need eval mode, please set `track_running_stats` to be False.",
       "If you are using a prebuilt module and do not need eval mode, please see the functorch website for resources on ",
@@ -85,18 +85,12 @@ batch_norm_batch_rule(
     if (running_mean.defined()) {
       running_mean_ = moveBatchDimToFront(running_mean, running_mean_bdim);
       running_mean_ = ensure_has_bdim(*running_mean_, running_mean_bdim.has_value(), bdim_size.value());
-      running_mean_ = reshape_dim_into(0, 0, *running_mean_);
-      if (training) {
-        running_mean_ = running_mean_->contiguous();
-      }
+      running_mean_ = reshape_dim_into(0, 0, *running_mean_).contiguous();
     }
     if (running_var.defined()) {
       running_var_ = moveBatchDimToFront(running_var, running_var_bdim);
       running_var_ = ensure_has_bdim(*running_var_, running_var_bdim.has_value(), bdim_size.value());
-      running_var_ = reshape_dim_into(0, 0, *running_var_);
-      if (training) {
-        running_var_ = running_var_->contiguous();
-      }
+      running_var_ = reshape_dim_into(0, 0, *running_var_).contiguous();
     }
 
     const auto dummy_weight = at::ones(input_.size(1), input_.options());  // cudnn and miopen require a weight


Test patch:
diff --git a/test/common_utils.py b/test/common_utils.py
index f45dc4704..332cae090 100644
--- a/test/common_utils.py
+++ b/test/common_utils.py
@@ -144,7 +144,22 @@ def get_batched_arg(arg, bdim):
         yield batched_args, in_dims, kwarg_values
 
 
-def get_exhaustive_batched_inputs_batch_norm(arg_values, kwarg_values, batch_size=2):
+def is_batch_norm_training(op_name, kwarg_values):
+    batch_norm_fns = ("nn.functional.batch_norm", "nn.functional.instance_norm")  # instance norm calls batch norm
+    if op_name not in batch_norm_fns:
+        return False
+
+    # batch norm and instance norm require the value to be a plain bool
+    default_training = op_name == "nn.functional.instance_norm"  # instance norm defaults to training, batch norm doesn't
+    is_training = tuple(arg for arg in tuple(kwarg_values.values()) if isinstance(arg, bool))
+    if len(is_training) == 0:
+        return default_training
+    else:
+        assert len(is_training) == 1
+        return is_training[0]
+
+
+def get_exhaustive_batched_inputs_batch_norm_is_training(arg_values, kwarg_values, batch_size=2):
     flat_args, arg_spec = pytree.tree_flatten(tuple(arg_values))
     is_tensors = [isinstance(a, torch.Tensor) for a in flat_args]
     num_tensors = sum(is_tensors)
@@ -169,13 +184,12 @@ def get_batched_arg(arg, bdim):
         yield batched_args, in_dims, kwarg_values
 
 
-def get_fallback_and_vmap_exhaustive(op, arg_values, kwarg_values, opinfo=None, compute_loop_out=True):
+def get_fallback_and_vmap_exhaustive(op, arg_values, kwarg_values, is_batch_norm_and_training=False, compute_loop_out=True):
     out_dim = 0
     batch_size = 2
-    batch_norm_fns = ("nn.functional.batch_norm", "nn.functional.instance_norm")  # instance norm calls batch norm
 
-    if opinfo is not None and opinfo.name in batch_norm_fns:
-        generator = get_exhaustive_batched_inputs_batch_norm(arg_values, kwarg_values, batch_size)
+    if is_batch_norm_and_training:
+        generator = get_exhaustive_batched_inputs_batch_norm_is_training(arg_values, kwarg_values, batch_size)
     else:
         generator = get_exhaustive_batched_inputs(arg_values, kwarg_values, batch_size)
 
diff --git a/test/test_ops.py b/test/test_ops.py
index 251ea3bb8..1d2539f69 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -19,7 +19,7 @@
 from common_utils import (
     get_fallback_and_vmap_exhaustive,
     get_exhaustive_batched_inputs,
-    get_exhaustive_batched_inputs_batch_norm,
+    get_exhaustive_batched_inputs_batch_norm_is_training,
     xfail,
     skip,
     skipOps,
@@ -27,6 +27,7 @@
     # tol2,
     opsToleranceOverride,
     check_vmap_fallback,
+    is_batch_norm_training,
 )
 from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
 from functorch import grad, vjp, vmap, jacrev, jacfwd
@@ -553,7 +554,9 @@ def vjp_of_vjp(*args_and_cotangents):
                 result_vjps, _ = tree_flatten(result_vjps)
                 return (*result, *result_vjps)
 
-            generator = get_fallback_and_vmap_exhaustive(vjp_of_vjp, args_and_cotangents, {}, opinfo=op)
+            is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs)
+            generator = get_fallback_and_vmap_exhaustive(
+                vjp_of_vjp, args_and_cotangents, {}, is_batch_norm_and_training=is_batch_norm_and_training)
             for loop_out, batched_out in generator:
                 self.assertEqual(loop_out, batched_out)
 
@@ -628,7 +631,10 @@ def test_vmapvjp(self, device, dtype, op):
         for sample in samples:
             cotangents = get_sample_cotangents(op, sample)
             fn, args = get_vjp_fn_and_args_with_cotangents(op, sample, cotangents)
-            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(fn, args, {}, opinfo=op):
+            is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs)
+            generator = get_fallback_and_vmap_exhaustive(
+                fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training)
+            for loop_out, batched_out in generator:
                 self.assertEqual(loop_out, batched_out)
 
     vmapjvpall_fail = {
@@ -709,7 +715,10 @@ def test_vmapjvpall(self, device, dtype, op):
             kwarg_values = sample.kwargs
             args = tuple(arg_values) + tuple(kwarg_values)
             fn, args = get_jvp_variant_primals_tangents(op, sample)
-            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(fn, args, {}, opinfo=op):
+            is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values)
+            generator = get_fallback_and_vmap_exhaustive(
+                fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training)
+            for loop_out, batched_out in generator:
                 self.assertEqual(loop_out, batched_out)
 
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
@@ -793,8 +802,9 @@ def test():
                 kwarg_values = sample.kwargs
                 args = tuple(arg_values) + tuple(kwarg_values)
                 fn, args = get_jvp_variant_primals_tangents(op, sample)
+                is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values)
                 for loop_out, batched_out in get_fallback_and_vmap_exhaustive(
-                        fn, args, {}, opinfo=op, compute_loop_out=False):
+                        fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False):
                     pass
         check_vmap_fallback(self, test, op, dry_run=False)
 
@@ -911,13 +921,14 @@ def test():
             for sample in samples:
                 cotangents = get_sample_cotangents(op, sample)
                 fn, args = get_vjp_fn_and_args_with_cotangents(op, sample, cotangents)
+                is_batch_norm_and_training = is_batch_norm_training(op.name, sample.kwargs)
                 for loop_out, batched_out in get_fallback_and_vmap_exhaustive(
-                        fn, args, {}, opinfo=op, compute_loop_out=False):
+                        fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False):
                     pass
                 for a_op in op.aliases:
                     fn, args = get_vjp_fn_and_args_with_cotangents(a_op, sample, cotangents)
                     for loop_out, batched_out in get_fallback_and_vmap_exhaustive(
-                            fn, args, {}, opinfo=op, compute_loop_out=False):
+                            fn, args, {}, is_batch_norm_and_training=is_batch_norm_and_training, compute_loop_out=False):
                         pass
 
         check_vmap_fallback(self, test, op, dry_run=False)
@@ -979,8 +990,8 @@ def test_vjpvmap(self, device, dtype, op):
         for sample in samples:
             args = [sample.input] + list(sample.args)
             kwargs = sample.kwargs
-            if is_batch_norm:
-                generator = get_exhaustive_batched_inputs_batch_norm(args, kwargs)
+            if is_batch_norm and is_batch_norm_training(op.name, kwargs):
+                generator = get_exhaustive_batched_inputs_batch_norm_is_training(args, kwargs)
             else:
                 generator = get_exhaustive_batched_inputs(args, kwargs)
 
@@ -1311,7 +1322,10 @@ def was_skipped_from_batched_tensors(batched_out, batch_size):
         for sample_input in sample_inputs:
             cotangents = get_sample_cotangents(op, sample_input)
             f, args = get_autograd_fn_and_args_with_cotangents(op, sample_input, cotangents)
-            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(f, args, {}, opinfo=op):
+            is_batch_norm_and_training = is_batch_norm_training(op, sample_input.kwargs)
+            generator = get_fallback_and_vmap_exhaustive(
+                f, args, {}, is_batch_norm_and_training=is_batch_norm_and_training)
+            for loop_out, batched_out in generator:
                 if all(was_skipped_from_batched_tensors(bo, lo.shape[0]) for (bo, lo) in zip(batched_out, loop_out)):
                     continue  # we weren't able to use the batched tensor in autograd.grad
                 self.assertEqual(loop_out, batched_out)
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 1d4cf867c..c5c255217 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -34,6 +34,7 @@
     check_vmap_fallback,
     tol1,
     opsToleranceOverride,
+    is_batch_norm_training,
 )
 import types
 from collections import namedtuple
@@ -3159,8 +3160,10 @@ def test_vmap_exhaustive(self, device, dtype, op):
         for sample_input in sample_inputs_itr:
             arg_values = [sample_input.input] + list(sample_input.args)
             kwarg_values = sample_input.kwargs
+            is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values)
             try:
-                generator = get_fallback_and_vmap_exhaustive(op.op, arg_values, kwarg_values, opinfo=op)
+                generator = get_fallback_and_vmap_exhaustive(
+                    op.op, arg_values, kwarg_values, is_batch_norm_and_training=is_batch_norm_and_training)
                 for loop_out, batched_out in generator:
                     # empty_like and new_empty produce garbage values so we just check the shapes.
                     if op.name == 'empty_like' or op.name == 'new_empty':
@@ -3168,7 +3171,8 @@ def test_vmap_exhaustive(self, device, dtype, op):
                         continue
                     self.assertEqual(loop_out, batched_out)
                 for a_op in op.aliases:
-                    a_generator = get_fallback_and_vmap_exhaustive(a_op, arg_values, kwarg_values, opinfo=op)
+                    a_generator = get_fallback_and_vmap_exhaustive(
+                        a_op, arg_values, kwarg_values, is_batch_norm_and_training=is_batch_norm_and_training)
                     for loop_out, batched_out in a_generator:
                         self.assertEqual(loop_out, batched_out)
             # todo(chilli): Garbage hack I added to deal with indexing not working
@@ -3303,7 +3307,9 @@ def test():
             for sample_input in sample_inputs_itr:
                 arg_values = [sample_input.input] + list(sample_input.args)
                 kwarg_values = sample_input.kwargs
-                generator = get_fallback_and_vmap_exhaustive(op.op, arg_values, kwarg_values, opinfo=op)
+                is_batch_norm_and_training = is_batch_norm_training(op.name, kwarg_values)
+                generator = get_fallback_and_vmap_exhaustive(
+                    op.op, arg_values, kwarg_values, is_batch_norm_and_training=is_batch_norm_and_training)
                 for loop_out, batched_out in generator:
                     # empty_like and new_empty produce garbage values so we just check the shapes.
                     if op.name == 'empty_like' or op.name == 'new_empty':
@@ -3311,7 +3317,8 @@ def test():
                         continue
                     self.assertEqual(loop_out, batched_out)
                 for a_op in op.aliases:
-                    a_generator = get_fallback_and_vmap_exhaustive(a_op, arg_values, kwarg_values, opinfo=op)
+                    a_generator = get_fallback_and_vmap_exhaustive(
+                        a_op, arg_values, kwarg_values, is_batch_norm_and_training=is_batch_norm_and_training)
                     for loop_out, batched_out in a_generator:
                         self.assertEqual(loop_out, batched_out)
         check_vmap_fallback(self, test, op)