FastBNNs developer guide¶

Introduction¶

FastBNNs was designed to enable automatic conversion of PyTorch-based neural networks (NNs) to Bayesian neural networks (BNNs), coupled with fast, approximate variational inference using the algorithms presented in [1] and [2]. To accomplish this, FastBNNs searches for leaf modules of a PyTorch neural network nn, wraps them in a module-dependent wrapper, and pairs each wrapper with an appropriate inference class. FastBNNs additionally implements a custom data type MuVar found in types.py. The MuVar type wraps two PyTorch Tensors mu and var, corresponding to the mean and variance of a Normal distribution, respectively. At inference, FastBNNs leverages a combination of the selected inference classes and custom-implemented MuVar operations to propagate network inputs and their uncertainties through the Bayesian neural network.

In the following sections, we provide additional details about this process that developers may find useful. In particular, we expand on the mechanics of the automated NN-to-BNN conversion process, leaf module wrappers, and inference classes/the MuVar type, as well as provide useful tips for developers that need to expand or modify the behavior of key components of FastBNNs.

Throughout development, we encourage developers to regularly run the tests provided in FastBNNs/tests using pytest and to extend or update these tests when appropriate.

Automated NN-to-BNN conversion¶

A PyTorch-based NN nn can be converted to a BNN using

from fastbnns.bnn.base import BNN
bnn = BNN(nn=nn, convert_in_place=False)

Internally, the base class BNN will call the helper function convert_to_bnn_() found in wrappers.py. convert_to_bnn_() attempts to find all unique leaf modules of nn (i.e., nn.named_modules() that do not have any child modules). Using a combination of default arguments as well as user-defined args and kwargs passed through BNN, convert_to_bnn_() will wrap each leaf module with a FastBNNs wrapper found in wrappers.py. At instantiation, each wrapper will additionally select an inference class using the helper function select_default_propagator() found in wrappers.py.

`convert_to_bnn_()`¶

The convert_to_bnn_() utility was designed to reroute forward passes through nn through bnn by hijacking the forward pass nn.forward() of the base nn. The intention is to capture standard operations made through a forward pass of a NN, e.g., the actions of each distinct module in a torch.nn.Sequential() module. To do this, module wrappers store a copy of the original leaf module (used for its forward() call) and its learnable parameters (used to define the mean of the parameter distributions). If appropriate (e.g., for a BayesianModule wrapper in wrappers.py with argument learn_var=True wrapping a module with learnable parameters), the wrapper will create another set of parameters associated with the variances of learnable parameters. In most cases, convert_to_bnn_() itself should not be modified. Instead, users should leverage the arguments of convert_to_bnn_() passed through the BNN base class to modify which wrappers and inference classes are used by each wrapper. We suggest developers carefully review convert_to_bnn_() to understand the selection process of wrappers and corresponding inference classes.

Module wrappers¶

A basic set of torch.nn.Module wrappers have been defined in wrappers.py. The wrapper BayesianModule is the primary wrapper for arbitrary leaf modules of nn. To write a custom wrapper for arbitrary leaf modules, we recommend inheriting from the base class BayesianModuleBase to ensure compatibility with the rest of the FastBNNs package. An additional wrapper BroadcastModule was added to reduce overhead when wrapping modules whose action applies independently to an input distribution (e.g., the torch.nn.Identity() module, which should leave an input distribution unchanged). A set of modules that are automatically wrapped by BroadcastModule is maintained in the variable BROADCAST in wrappers.py. Developers are encouraged to expand this list if appropriate modules from torch.nn are identified that should be present in this list.

Inference and the MuVar type¶

FastBNNs essentially hijacks the forward pass of the underlying nn to avoid custom mapping of the computational graph. This is achieved through the combined action of custom inference classes inference.py and the datatype MuVar found in types.py.

Inference classes are intended to reroute module inputs to the underlying modules forward call as appropriate for each inference algorithm. For example, the custom inference algorithm Linear in inference.py will use analytic moment propagation rules to compute the mean and variance of torch.nn.Linear when called on an input random variable. More general algorithms found in inference.py include UnscentedTransform, JointUnscentedTransform, and MonteCarlo, which similarly route inputs through forward calls to underlying torch.nn.Modules and estimate the output mean and variance from the results. If the user does not specify a specific inference algorithm for each module when using BNN() or convert_to_bnn_(), a default is selected to based on the module name and whether or not it has learnable parameters (see select_default_propagator() in wrappers.py understand the selection process). Analytic propagation algorithms for layers such as torch.nn.Linear, torch.nn.Conv(1,2,3)d, torch.nn.ConvTranspose(1,2,3)d, and torch.nn.AvgPool(1,2,3)d have already been implemented. Custom analytic propagation algorithms for modules in torch.nn should be added to inference.py with the algorithm defined as a MomentPropagator class whose name matches the name of the torch.nn.Module. Analytic propagation algorithms for torch.nn.ReLU and torch.nn.LeakyReLU are included as reference under the aliases ReLUa and LeakyReLUa. The suffix a was added to intentionally prevent their selection by select_default_propagator(), as we have found the UnscentedTransform to be faster with minimal loss in accuracy (unpublished result).

Since many PyTorch NN workflows include additional operations not encompassed in a subclass of torch.nn.Module, FastBNNs implements a custom type MuVar to carry the mean and variance of a distribution through such operations. The MuVar type implements custom handlers in types.py which apply analytic propagation rules (e.g., for the addition of two random variables) or fallbacks like the unscented transform (e.g., for arbitrary torch function calls whose computational graph is unknown). If operations that work for torch.Tensor instances do not work (i.e., throw an error) for the MuVar type, developers should review the sets SIMPLE_TORCH_FUNCS and TENSOR_METHODS to determine if the operation can be added to these sets. If instead an operation works for torch.Tensor and works with MuVar without throwing an error, yet the output differs from the result expected by the user, developers should consider implementing a custom handler for that operation and register it with the MUVAR_HANDLERS in types.py.

To review the FastBNNs inference workflow with a simple example, consider a NN

class model(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.module_list = torch.nn.ModuleList([
            torch.nn.Linear(1, 2),
            torch.nn.LeakyReLU(),
            torch.nn.Linear(2, 1)
        ])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.module_list:
            x = layer(x)
        
        return x + x**2

We then convert to a BNN and run a test pass through the model:

from fastbnns.bnn.base import BNN
from fastbnns.bnn.types import MuVar

nn = model()
bnn = BNN(nn=nn, convert_in_place=False)
out = bnn(MuVar(torch.randn(1, 1)))

At conversion, the torch.nn.Linear layers will be wrapped in a BayesianModule and assigned an inference class Linear, while the torch.nn.LeakyReLU layer will be wrapped in a BayesianModule with inference class UnscentedTransform. When calling bnn on a MuVar wrapped input, the action x=layer(x) will be rerouted through the inference module of each layer to produce an output MuVar corresponding to moment propagation through that layer. In contrast, the final set of operations x + x**2, which are not contained within a torch.nn.Module, will rely on handlers for add and pow defined in the MuVar type.

References¶

[1] David J. Schodt, Ryan Brown, Michael Merritt, Samuel Park, Delsin Menolascino, and Mark A. Peot. A framework for variational inference of lightweight bayesian neural networks with heteroscedastic uncertainties. 2024. arXiv:2402.14532 [cs].

[2] David J. Schodt. Few-sample Variational Inference of Bayesian Neural Networks with Arbitrary Nonlinearities. 2024. arXiv:2405.02063 [cs].