[ADD] AutomaticDerivatives for autograd-based support of new layers

README.md

@@ -1,8 +1,8 @@
 # <img alt="BackPACK" src="./logo/backpack_logo_torch.svg" height="90"> BackPACK: Packing more into backprop
 
-[![Travis](https://travis-ci.org/f-dangel/backpack.svg?branch=master)](https://travis-ci.org/f-dangel/backpack)
+[![RTD](https://readthedocs.org/projects/backpack/badge/?version=master)]()
 [![Coveralls](https://coveralls.io/repos/github/f-dangel/backpack/badge.svg?branch=master)](https://coveralls.io/github/f-dangel/backpack)
-[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/release/python-370/)
+[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/release/python-390/)
 
 BackPACK is built on top of [PyTorch](https://github.com/pytorch/pytorch). It efficiently computes quantities other than the gradient.
 

backpack/core/derivatives/automatic.py

@@ -0,0 +1,181 @@
+"""Automatic derivative implementation via ``torch.func``."""
+
+from copy import deepcopy
+from functools import partial
+from typing import Callable, List, Optional, Tuple, Union, Any
+
+from torch import Tensor
+from torch.func import functional_call, vjp, vmap
+from torch.nn import Module
+
+from backpack.core.derivatives.basederivatives import BaseParameterDerivatives
+from backpack.utils.subsampling import subsample
+
+
+class AutomaticDerivatives(BaseParameterDerivatives):
+    """Implements derivatives for an arbitrary layer using ``torch.func``.
+
+    This class can be used to support new layers without implementing their
+    derivatives by hand. However, this comes at the cost of performance, since the
+    autograd-based implementation is often not as efficient as a hand-crafted one
+    and re-evaluates the forward pass through the layer.
+
+    Attributes:
+        BATCH_AXIS: Index of the layer input's batch axis. Default: ``0``.
+    """
+
+    BATCH_AXIS: int = 0
+
+    @staticmethod
+    def clone_without_hooks(module: Module) -> Module:
+        """Create a copy of module without BackPACK hooks.
+
+        Args:
+            module: Module to be cloned.
+
+        Returns:
+            Cloned module without BackPACK hooks.
+        """
+        # Temporarily remove input and output to avoid error when calling deepcopy
+        input0, output = module.input0, module.output
+        delattr(module, "input0")
+        delattr(module, "output")
+
+        # Clone the module and remove BackPACK's hooks
+        clean = deepcopy(module)
+        clean._forward_hooks.clear()
+        clean._backward_hooks.clear()
+
+        # Restore input and output in the original module
+        module.input0, module.output = input0, output
+
+        return clean
+
+    def as_functional(
+        self, module: Module, param_name: Optional[str] = None
+    ) -> Union[Callable[[Tensor, Tensor], Tensor], Callable[[Tensor], Tensor]]:
+        """Return a function that performs the layer's forward pass on a single datum.
+
+        Args:
+            module: Layer for which to return the forward function.
+            param_name: If specified, the name of a parameter of the module that
+                will be passed as second argument to the returned function.
+                Default: ``None`` (all parameters are frozen in the module).
+
+        Returns:
+            Function that performs the forward pass of the layer and returns a tensor
+            representing the result. First argument is the un-batched input tensor.
+            If `param_name` is specified, the second argument corresponds to the
+            parameter.
+        """
+        module = self.clone_without_hooks(module)
+        parameters = dict(module.named_parameters())
+        buffers = dict(module.named_buffers())
+
+        if param_name is None:
+
+            def f(x: Tensor) -> Tensor:
+                return functional_call(module, {**parameters, **buffers}, x)
+
+        else:
+            parameters.pop(param_name)
+
+            def f(x: Tensor, param: Tensor) -> Tensor:
+                return functional_call(
+                    module, {**parameters, **buffers, param_name: param}, x
+                )
+
+        return f
+
+    def _jac_t_mat_prod(
+        self,
+        module: Module,
+        g_inp: Tuple[Tensor],
+        g_out: Tuple[Tensor],
+        mat: Tensor,
+        subsampling: Optional[List[int]] = None,
+    ) -> Tensor:
+        f = vmap(self.as_functional(module))  # {x_n} -> {f(x_n)}
+
+        X = subsample(module.input0, dim=self.BATCH_AXIS, subsampling=subsampling)
+        _, vjp_func = vjp(f, X)  # {v_n} -> {Jf(x_n)^T v_n}
+        # vmap over matrix columns
+        vmp_func = vmap(vjp_func)
+        (vmp,) = vmp_func(mat)
+
+        return vmp
+
+    def __getattribute__(self, name: str) -> Any:
+        """Dynamically generate parameter MJP methods if their attributes are accessed.
+
+        Args:
+            name: Name of the requested attribute.
+
+        Returns:
+            The requested attribute.
+        """
+        suffix = "_jac_t_mat_prod"
+        if name.endswith(suffix) and name.split(suffix)[0] != "":
+            param_str, _ = name.split(suffix)
+            param_str = param_str[1:] # remove leading underscore
+            return partial(self._param_mjp, param_str)
+
+        return super().__getattribute__(name)
+
+
+    def _param_mjp(
+        self,
+        param_str: str,
+        module: Module,
+        g_inp: Tuple[Tensor],
+        g_out: Tuple[Tensor],
+        mat: Tensor,
+        sum_batch: bool = True,
+        subsampling: Optional[List[int]] = None,
+    ) -> Tensor:
+        """Compute matrix-Jacobian products (MJPs) of the module w.r.t. a parameter.
+
+        Handles both vector and matrix inputs. Preserves input format in output.
+
+        Args:
+            param_str: Attribute name under which the parameter is stored in the module.
+            module: Module whose Jacobian will be applied. Must provide access to IO.
+            g_inp: Gradients w.r.t. module input.
+            g_out: Gradients w.r.t. module output.
+            mat: Matrix the Jacobian will be applied to. Has shape
+                ``[V, *module.output.shape]`` (matrix case) or same shape as
+                ``module.output`` (vector case). If used with subsampling, has dimension
+                len(subsampling) instead of batch size along the batch axis.
+            sum_batch: Sum out the MJP's batch axis. Default: ``True``.
+            subsampling: Indices of samples along the output's batch dimension that
+                should be considered. Defaults to ``None`` (use all samples).
+
+        Returns:
+            Matrix-Jacobian products. Has shape ``[V, *param_shape]`` when batch
+            summation is enabled (same shape as parameter in the vector case). Without
+            batch summation, the result has shape ``[V, N, *param_shape]`` (vector case
+            has shape ``[N, *param_shape]``). If used with subsampling, the batch size N
+            is replaced by len(subsampling).
+        """
+        f = self.as_functional(
+            module, param_name=param_str
+        )  # (x, param) -> f(x, param)
+
+        def param_vjp(x, param, v) -> Tensor:
+            f_x = partial(f, x)  # param -> f(x, param)
+            _, vjp_func = vjp(f_x, param)  # v -> Jf(x, param)^T v
+            (mjp,) = vjp_func(v)
+            return mjp
+
+        # vectorize over data points: ({x_n}, param, {v_n}) -> {Jf(x_n, param)^T v_n}
+        param_vjp = vmap(param_vjp, in_dims=(self.BATCH_AXIS, None, self.BATCH_AXIS))
+        # vectorize over matrix columns
+        param_mjp = vmap(param_vjp, in_dims=(None, None, 0))
+
+        X = subsample(module.input0, dim=self.BATCH_AXIS, subsampling=subsampling)
+        mjp = param_mjp(X, getattr(module, param_str), mat)
+
+        if sum_batch:
+            mjp = mjp.sum(dim=self.BATCH_AXIS + 1)
+
+        return mjp

backpack/utils/subsampling.py

@@ -1,11 +1,13 @@
 """Utility functions to enable mini-batch subsampling in extensions."""
 
-from typing import List
+from typing import List, Optional
 
 from torch import Tensor
 
 
-def subsample(tensor: Tensor, dim: int = 0, subsampling: List[int] = None) -> Tensor:
+def subsample(
+    tensor: Tensor, dim: int = 0, subsampling: Optional[List[int]] = None
+) -> Tensor:
     """Select samples from a tensor along a dimension.
 
     Args:

fully_documented.txt

@@ -19,6 +19,7 @@ backpack/core/derivatives/sum_module.py
 backpack/core/derivatives/dropout.py
 backpack/core/derivatives/slicing.py
 backpack/core/derivatives/bcewithlogitsloss.py
+backpack/core/derivatives/automatic.py
 
 backpack/extensions/__init__.py
 backpack/extensions/backprop_extension.py
@@ -130,3 +131,5 @@ test/utils/conv_transpose.py
 test/custom_module/
 test/test_retain_graph.py
 test/test_batch_first.py
+test/test_automatic_support.py
+test/automatic_extensions.py

test/automatic_extensions.py

@@ -0,0 +1,49 @@
+"""Define layer extensions with derivatives based on autodiff."""
+
+from backpack.core.derivatives.automatic import AutomaticDerivatives
+from backpack.extensions.firstorder.batch_grad.batch_grad_base import BatchGradBase
+from backpack.extensions.secondorder.diag_ggn.diag_ggn_base import DiagGGNBaseModule
+
+
+class DiagGGNExactAutomatic(DiagGGNBaseModule):
+    """GGN diagonal computation for modules via automatic derivatives."""
+
+    def __init__(self):
+        """Set up the derivatives."""
+        super().__init__(AutomaticDerivatives(), sum_batch=True)
+
+
+class DiagGGNExactLinearAutomatic(DiagGGNBaseModule):
+    """GGN diag. computation for ``torch.nn.Linear`` via automatic derivatives."""
+
+    def __init__(self):
+        """Set up the derivatives."""
+        super().__init__(
+            AutomaticDerivatives(), params=["weight", "bias"], sum_batch=True
+        )
+
+
+class BatchDiagGGNExactAutomatic(DiagGGNBaseModule):
+    """GGN diagonal computation for modules via automatic derivatives."""
+
+    def __init__(self):
+        """Set up the derivatives."""
+        super().__init__(AutomaticDerivatives(), sum_batch=False)
+
+
+class BatchDiagGGNExactLinearAutomatic(DiagGGNBaseModule):
+    """GGN diag. computation for ``torch.nn.Linear`` via automatic derivatives."""
+
+    def __init__(self):
+        """Set up the derivatives."""
+        super().__init__(
+            AutomaticDerivatives(), params=["weight", "bias"], sum_batch=False
+        )
+
+
+class BatchGradLinearAutomatic(BatchGradBase):
+    """Batch gradients for ``torch.nn.Linear`` via automatic derivatives."""
+
+    def __init__(self):
+        """Set up the derivatives."""
+        super().__init__(AutomaticDerivatives(), params=["weight", "bias"])

test/test_automatic_support.py

@@ -0,0 +1,113 @@
+"""Test automatic support of new layers."""
+
+from test.automatic_extensions import (
+    BatchDiagGGNExactAutomatic,
+    BatchDiagGGNExactLinearAutomatic,
+    BatchGradLinearAutomatic,
+    DiagGGNExactAutomatic,
+    DiagGGNExactLinearAutomatic,
+)
+from test.test___init__ import DEVICES, DEVICES_ID
+from test.utils import popattr
+from typing import List, Union
+
+from pytest import mark, raises
+from torch import allclose, device, manual_seed, rand
+from torch.nn import Linear, MSELoss, ReLU, Sequential, Sigmoid
+
+from backpack import backpack, extend, extensions
+
+
+@mark.parametrize("batched", [False, True], ids=["DiagGGNExact", "BatchDiagGGNExact"])
+@mark.parametrize("dev", DEVICES, ids=DEVICES_ID)
+def test_automatic_support_diag_ggn_exact(dev: device, batched: bool):
+    """Test GGN diagonal computation via automatic derivatives.
+
+    Args:
+        dev: The device on which to run the test.
+        batched: Whether to compute the batched or summed GGN diagonal.
+    """
+    manual_seed(0)
+    X, y = rand(10, 5, device=dev), rand(10, 3, device=dev)
+
+    model = extend(Sequential(Linear(5, 4), ReLU(), Linear(4, 3), Sigmoid()).to(dev))
+    loss_func = extend(MSELoss().to(dev))
+    savefield = "diag_ggn_exact_batch" if batched else "diag_ggn_exact"
+
+    # ground truth
+    ext = extensions.BatchDiagGGNExact() if batched else extensions.DiagGGNExact()
+    with backpack(ext):
+        loss = loss_func(model(X), y)
+        loss.backward()
+    manual = [popattr(p, savefield) for p in model.parameters()]
+
+    # same quantity with automatic support
+    ext = extensions.BatchDiagGGNExact() if batched else extensions.DiagGGNExact()
+    new_mappings = {
+        ReLU: (BatchDiagGGNExactAutomatic() if batched else DiagGGNExactAutomatic()),
+        Linear: (
+            BatchDiagGGNExactLinearAutomatic()
+            if batched
+            else DiagGGNExactLinearAutomatic()
+        ),
+        Sigmoid: (BatchDiagGGNExactAutomatic() if batched else DiagGGNExactAutomatic()),
+    }
+    for layer_cls, extension in new_mappings.items():
+        # make sure we need to turn on explicit overwriting
+        with raises(ValueError):
+            ext.set_module_extension(layer_cls, extension)
+        ext.set_module_extension(layer_cls, extension, overwrite=True)
+
+    with backpack(ext):
+        loss = loss_func(model(X), y)
+        loss.backward()
+    automatic = [popattr(p, savefield) for p in model.parameters()]
+
+    assert len(manual) == len(automatic)
+    for m, a in zip(manual, automatic):
+        assert allclose(m, a)
+
+
+SUBSAMPLINGS = [None, [7, 2, 4]]
+SUBSAMPLING_IDS = [f"subsampling={subsampling}" for subsampling in SUBSAMPLINGS]
+
+
+@mark.parametrize("subsampling", SUBSAMPLINGS, ids=SUBSAMPLING_IDS)
+@mark.parametrize("dev", DEVICES, ids=DEVICES_ID)
+def test_automatic_support_batch_grad(dev: device, subsampling: Union[None, List[int]]):
+    """Test per-example gradient computation via automatic derivatives.
+
+    Args:
+        dev: The device on which to run the test.
+        subsampling: Indices of active samples. ``None`` means full batch.
+    """
+    manual_seed(0)
+    X, y = rand(10, 5, device=dev), rand(10, 3, device=dev)
+
+    model = extend(Sequential(Linear(5, 4), ReLU(), Linear(4, 3), Sigmoid()).to(dev))
+    loss_func = extend(MSELoss().to(dev))
+    savefield = "grad_batch"
+
+    # ground truth
+    with backpack(extensions.BatchGrad(subsampling=subsampling)):
+        loss = loss_func(model(X), y)
+        loss.backward()
+    manual = [popattr(p, savefield) for p in model.parameters()]
+
+    # same quantity with automatic support
+    ext = extensions.BatchGrad(subsampling=subsampling)
+    new_mappings = {Linear: BatchGradLinearAutomatic()}
+    for layer_cls, extension in new_mappings.items():
+        # make sure we need to turn on explicit overwriting
+        with raises(ValueError):
+            ext.set_module_extension(layer_cls, extension)
+        ext.set_module_extension(layer_cls, extension, overwrite=True)
+
+    with backpack(ext):
+        loss = loss_func(model(X), y)
+        loss.backward()
+    automatic = [popattr(p, savefield) for p in model.parameters()]
+
+    assert len(manual) == len(automatic)
+    for m, a in zip(manual, automatic):
+        assert allclose(m, a)