diff --git a/README.md b/README.md index 25956cf7..ecee838c 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ # BackPACK 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. diff --git a/backpack/core/derivatives/automatic.py b/backpack/core/derivatives/automatic.py new file mode 100644 index 00000000..d44941f5 --- /dev/null +++ b/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 diff --git a/backpack/utils/subsampling.py b/backpack/utils/subsampling.py index 75cb6e73..04d09868 100644 --- a/backpack/utils/subsampling.py +++ b/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: diff --git a/fully_documented.txt b/fully_documented.txt index ebf4bd3b..ea974eb6 100644 --- a/fully_documented.txt +++ b/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 diff --git a/test/automatic_extensions.py b/test/automatic_extensions.py new file mode 100644 index 00000000..c8b73553 --- /dev/null +++ b/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"]) diff --git a/test/test_automatic_support.py b/test/test_automatic_support.py new file mode 100644 index 00000000..a1ce7471 --- /dev/null +++ b/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) diff --git a/test/utils/__init__.py b/test/utils/__init__.py index 40711349..c6b0c7f2 100644 --- a/test/utils/__init__.py +++ b/test/utils/__init__.py @@ -1,6 +1,6 @@ """Helper functions for tests.""" -from typing import List +from typing import Any, List def chunk_sizes(total_size: int, num_chunks: int) -> List[int]: @@ -25,3 +25,18 @@ def chunk_sizes(total_size: int, num_chunks: int) -> List[int]: sizes.append(rest) return sizes + + +def popattr(obj: Any, name: str) -> Any: + """Pop an attribute from an object. + + Args: + obj: The object from which to pop the attribute. + name: The name of the attribute to pop. + + Returns: + The attribute's value. + """ + value = getattr(obj, name) + delattr(obj, name) + return value