Better error messages for BatchNorm

backpack/extensions/backprop_extension.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import abc
-import warnings
 from abc import ABC
 from typing import Any, Dict, List, Tuple, Type, Union
 
@@ -60,7 +59,10 @@ def __init__(
             AssertionError: if fail_mode is not valid
         """
         if fail_mode not in (FAIL_WARN, FAIL_ERROR, FAIL_SILENT):
-            raise AssertionError(f"no valid fail mode: {fail_mode}")
+            raise AssertionError(
+                f"Invalid failure mode: {fail_mode}."
+                f" Should be one of {(FAIL_WARN, FAIL_ERROR, FAIL_SILENT)}"
+            )
         self.saved_quantities: SavedQuantities = SavedQuantities()
         self.savefield: str = savefield
         self.__module_extensions: Dict[Type[Module], ModuleExtension] = module_exts
@@ -94,24 +96,17 @@ def __get_module_extension(self, module: Module) -> Union[ModuleExtension, None]
         module_extension = self.__module_extensions.get(module.__class__)
 
         if module_extension is None:
-            if self._fail_mode is FAIL_ERROR:
-                # PyTorch converts this Error into a RuntimeError for torch<1.7.0
-                raise NotImplementedError(
-                    f"Extension saving to {self.savefield} "
-                    "does not have an extension for "
-                    f"Module {module.__class__}"
-                )
-            elif self._fail_mode == FAIL_WARN:
-                for _ in module.parameters():
-                    warnings.warn(
-                        f"Extension saving to {self.savefield} does not have an "
-                        f"extension for Module {module.__class__} "
-                        f"although the module has parameters"
-                    )
-                    break
+            self._handle_missing_module_extension(module)
 
         return module_extension
 
+    def _handle_missing_module_extension(self, module: Module) -> None:
+        """What to do if module does not have an extension (default: raise exception)"""
+        raise NotImplementedError(
+            f"Extension saving to {self.savefield} "
+            f"does not have an extension for Module {module.__class__}."
+        )
+
     def __call__(
         self, module: Module, g_inp: Tuple[Tensor], g_out: Tuple[Tensor]
     ) -> None:

backpack/extensions/firstorder/base.py

@@ -1,9 +1,12 @@
 """Base class for first order extensions."""
+import warnings
 from typing import Dict, List, Type
 
+from backpack.utils.errors import change_error_to_warn_message
 from torch.nn import Module
 
-from backpack.extensions.backprop_extension import FAIL_WARN, BackpropExtension
+from backpack.extensions.backprop_extension import BackpropExtension
+from backpack.extensions.backprop_extension import FAIL_ERROR, FAIL_WARN
 from backpack.extensions.module_extension import ModuleExtension
 
 
@@ -18,7 +21,7 @@ def __init__(
         self,
         savefield: str,
         module_exts: Dict[Type[Module], ModuleExtension],
-        fail_mode: str = FAIL_WARN,
+        fail_mode: str = FAIL_ERROR,
         subsampling: List[int] = None,
     ):  # noqa: D107
         super().__init__(
@@ -27,3 +30,19 @@ def __init__(
 
     def expects_backpropagation_quantities(self) -> bool:  # noqa: D102
         return False
+
+    def _handle_missing_module_extension(self, module: Module) -> None:
+        message = (
+            f"Extension saving to {self.savefield} "
+            f"does not have an extension for Module {module.__class__}, "
+            "but it has parameters. "
+            f"Those parameters will not have their field {self.savefield} set."
+        )
+
+        for _ in module.parameters():
+            if self._fail_mode is FAIL_ERROR:
+                # PyTorch converts this Error into a RuntimeError for torch<1.7.0
+                raise NotImplementedError(message + " " + change_error_to_warn_message)
+            elif self._fail_mode == FAIL_WARN:
+                warnings.warn(message)
+            break

backpack/extensions/firstorder/batch_grad/__init__.py

@@ -4,6 +4,7 @@
 """
 from typing import List
 
+from backpack.extensions.backprop_extension import FAIL_ERROR
 from torch.nn import (
     LSTM,
     RNN,
@@ -60,7 +61,7 @@ class BatchGrad(FirstOrderBackpropExtension):
     objective is a sum of independent functions (no batchnorm).
     """
 
-    def __init__(self, subsampling: List[int] = None):
+    def __init__(self, subsampling: List[int] = None, fail_mode: str = FAIL_ERROR):
         """Initialization.
 
         Defines extension for each module.
@@ -87,4 +88,5 @@ def __init__(self, subsampling: List[int] = None):
                 Embedding: embedding.BatchGradEmbedding(),
             },
             subsampling=subsampling,
+            fail_mode=fail_mode,
         )

backpack/extensions/firstorder/batch_grad/batchnorm_nd.py

@@ -26,4 +26,4 @@ def check_hyperparameters_module_extension(
         g_inp: Tuple[Tensor],
         g_out: Tuple[Tensor],
     ) -> None:  # noqa: D102
-        batch_norm_raise_error_if_train(module, raise_error=False)
+        batch_norm_raise_error_if_train(module, ext)

backpack/extensions/firstorder/batch_l2_grad/__init__.py

@@ -3,6 +3,7 @@
 Defines the backpropagation extension.
 Within it, define the extension for each module.
 """
+from backpack.extensions.backprop_extension import FAIL_ERROR
 from torch.nn import (
     LSTM,
     RNN,
@@ -48,7 +49,7 @@ class BatchL2Grad(FirstOrderBackpropExtension):
         - ``[¹/ₙ g₁, …, ¹/ₙ gₙ]`` if the loss is a mean, ``¹/ₙ ∑ᵢ₌₁ⁿ fᵢ``.
     """
 
-    def __init__(self):
+    def __init__(self, fail_mode: str = FAIL_ERROR):
         """Initialization.
 
         Define the extensions for each module.
@@ -70,4 +71,5 @@ def __init__(self):
                 BatchNorm3d: batchnorm_nd.BatchL2BatchNorm(),
                 Embedding: embedding.BatchL2Embedding(),
             },
+            fail_mode=fail_mode,
         )

backpack/extensions/firstorder/batch_l2_grad/batchnorm_nd.py

@@ -24,4 +24,4 @@ def check_hyperparameters_module_extension(
         g_inp: Tuple[Tensor],
         g_out: Tuple[Tensor],
     ) -> None:  # noqa: D102
-        batch_norm_raise_error_if_train(module)
+        batch_norm_raise_error_if_train(module, ext)

backpack/extensions/firstorder/gradient/batchnorm_nd.py

@@ -27,4 +27,4 @@ def check_hyperparameters_module_extension(
         g_inp: Tuple[Tensor],
         g_out: Tuple[Tensor],
     ) -> None:  # noqa: D102
-        batch_norm_raise_error_if_train(module)
+        batch_norm_raise_error_if_train(module, ext)

backpack/extensions/firstorder/sum_grad_squared/__init__.py

@@ -2,6 +2,7 @@
 
 Defines module extension for each module.
 """
+from backpack.extensions.backprop_extension import FAIL_ERROR
 from torch.nn import (
     LSTM,
     RNN,
@@ -51,7 +52,7 @@ class SumGradSquared(FirstOrderBackpropExtension):
         - ``[¹/ₙ g₁, …, ¹/ₙ gₙ]`` if the loss is a mean, ``¹/ₙ ∑ᵢ₌₁ⁿ fᵢ``.
     """
 
-    def __init__(self):
+    def __init__(self, fail_mode: str = FAIL_ERROR):
         """Initialization.
 
         Defines module extension for each module.
@@ -73,4 +74,5 @@ def __init__(self):
                 BatchNorm3d: batchnorm_nd.SGSBatchNormNd(),
                 Embedding: embedding.SGSEmbedding(),
             },
+            fail_mode=fail_mode,
         )

backpack/extensions/firstorder/sum_grad_squared/batchnorm_nd.py

@@ -24,4 +24,4 @@ def check_hyperparameters_module_extension(
         g_inp: Tuple[Tensor],
         g_out: Tuple[Tensor],
     ) -> None:  # noqa: D102
-        batch_norm_raise_error_if_train(module)
+        batch_norm_raise_error_if_train(module, ext)

backpack/extensions/firstorder/variance/__init__.py

@@ -2,6 +2,7 @@
 
 Defines module extension for each module.
 """
+from backpack.extensions.backprop_extension import FAIL_ERROR
 from torch.nn import (
     LSTM,
     RNN,
@@ -51,7 +52,7 @@ class Variance(FirstOrderBackpropExtension):
         - ``[¹/ₙ g₁, …, ¹/ₙ gₙ]`` if the loss is a mean, ``¹/ₙ ∑ᵢ₌₁ⁿ fᵢ``.
     """
 
-    def __init__(self):
+    def __init__(self, fail_mode: str = FAIL_ERROR):
         """Initialization.
 
         Defines module extension for each module.
@@ -73,4 +74,5 @@ def __init__(self):
                 BatchNorm3d: batchnorm_nd.VarianceBatchNormNd(),
                 Embedding: embedding.VarianceEmbedding(),
             },
+            fail_mode=fail_mode,
         )

backpack/extensions/firstorder/variance/batchnorm_nd.py

@@ -25,4 +25,4 @@ def check_hyperparameters_module_extension(
         g_inp: Tuple[Tensor],
         g_out: Tuple[Tensor],
     ) -> None:  # noqa: D102
-        batch_norm_raise_error_if_train(module)
+        batch_norm_raise_error_if_train(module, ext)

backpack/extensions/secondorder/base.py

@@ -1,9 +1,28 @@
 """Contains base classes for second order extensions."""
+import warnings
+
 from backpack.extensions.backprop_extension import BackpropExtension
+from backpack.extensions.backprop_extension import FAIL_ERROR, FAIL_WARN
+from backpack.utils.errors import change_error_to_warn_message
+from torch.nn import Module
 
 
 class SecondOrderBackpropExtension(BackpropExtension):
     """Base backpropagation extension for second order."""
 
+    def _handle_missing_module_extension(self, module: Module) -> None:
+        message = (
+            f"Extension saving to {self.savefield} "
+            f"does not have an extension for Module {module.__class__}. "
+            "Further computations will likely fail as second-order quantities "
+            "will not be backpropagated."
+        )
+
+        if self._fail_mode is FAIL_ERROR:
+            # PyTorch converts this Error into a RuntimeError for torch<1.7.0
+            raise NotImplementedError(message + " " + change_error_to_warn_message)
+        elif self._fail_mode == FAIL_WARN:
+            warnings.warn(message)
+
     def expects_backpropagation_quantities(self) -> bool:  # noqa: D102
         return True

backpack/extensions/secondorder/diag_ggn/batchnorm_nd.py

@@ -41,4 +41,4 @@ def check_hyperparameters_module_extension(
         g_inp: Tuple[Tensor],
         g_out: Tuple[Tensor],
     ) -> None:  # noqa: D102
-        batch_norm_raise_error_if_train(module)
+        batch_norm_raise_error_if_train(module, ext)

backpack/utils/cleanup.py

@@ -0,0 +1,37 @@
+import inspect
+from backpack import BackpropExtension
+
+
+def _check_extensions(func_name, extensions):
+    for ext in extensions:
+        if not isinstance(ext, BackpropExtension):
+            if inspect.isclass(ext) and issubclass(ext, BackpropExtension):
+                raise ValueError(
+                    "{} expect instances of BackpropExtension,".format(func_name)
+                    + " but received a class instead [{}].".format(ext)
+                    + " Instantiate it before passing it to backpack."
+                )
+            else:
+                raise ValueError(
+                    "{} expects instances of BackpropExtension,".format(func_name)
+                    + " but received [{}].".format(ext)
+                )
+
+
+def zero_backpack(model, *extensions):
+    """Clears the backpack-computed quantities of the model.
+
+    Removes references from the model to quantities computed by BackPACK
+    using the given extensions. Can be used to free memory or remove additional
+    tensors from model parameters for saving data with pickle or deepcopy.
+
+    Args:
+        model: A :py:class:`Module <torch.nn.Module>`
+        *extensions: Instances of BackpropExtensions that have been called
+    """
+    _check_extensions("zero_backpack", extensions)
+
+    for p in model.parameters():
+        for ext in extensions:
+            if hasattr(p, ext.savefield):
+                delattr(p, ext.savefield)

backpack/utils/errors.py

@@ -2,28 +2,45 @@
 from typing import Union
 from warnings import warn
 
+from backpack.extensions.backprop_extension import BackpropExtension
+from backpack.extensions.backprop_extension import FAIL_ERROR, FAIL_WARN, FAIL_SILENT
 from torch.nn import BatchNorm1d, BatchNorm2d, BatchNorm3d
 
+change_error_to_warn_message = (
+    "To only raise a warning, change the failure mode of the BackPACK extension "
+    "(for example, `BatchGrad(fail_mode='WARNING')`)."
+)
+
 
 def batch_norm_raise_error_if_train(
-    module: Union[BatchNorm1d, BatchNorm2d, BatchNorm3d], raise_error: bool = True
+    module: Union[BatchNorm1d, BatchNorm2d, BatchNorm3d], ext: BackpropExtension
 ) -> None:
     """Check if BatchNorm module is in training mode.
 
     Args:
         module: BatchNorm module to check
-        raise_error: whether to raise an error, alternatively warn. Default: True.
+        ext: The BackpropExtension checking for errors
 
     Raises:
-        NotImplementedError: if module is in training mode
+        ValueError: if module is in training mode and BackPACK extension's
+            fail_mode is FAIL_ERROR (default)
     """
     if module.training:
         message = (
-            "Encountered BatchNorm module in training mode. BackPACK's computation "
-            "will pass, but results like individual gradients may not be meaningful, "
-            "as BatchNorm mixes samples. Only proceed if you know what you are doing."
+            "Encountered BatchNorm module in training mode."
+            "Quantity to compute is undefined as BatchNorm mixes samples. "
+            "You should most likely use another type of normalization. "
+            "Concepts like individual gradients are not meaningful with BatchNorm. "
+            "The code to compute the requested quantity may not raise an error, "
+            "but the quantity will not match its definition. "
+            "Advanced users: If you are specifically interested in what this code "
+            "would return for a BatchNorm network, change the failure mode of "
+            f"the BackPACK extension. {change_error_to_warn_message} "
+            "This is not supported behavior."
         )
-        if raise_error:
-            raise NotImplementedError(message)
-        else:
+        if ext._fail_mode == FAIL_ERROR:
+            raise ValueError(message)
+        if ext._fail_mode == FAIL_WARN:
             warn(message)
+        if ext._fail_mode == FAIL_SILENT:
+            return