Metadata-Version: 2.4
Name: fastbnns
Version: 0.4.0
Summary: Fast training and inference for Bayesian neural networks.
Author-email: "David J. Schodt" <dschodt@lanl.gov>
License-Expression: BSD-3-Clause
Project-URL: Homepage, https://github.com/lanl/FastBNNs
Project-URL: Repository, https://github.com/lanl/FastBNNs
Project-URL: Issues, https://github.com/lanl/FastBNNs/issues
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE.md
Requires-Dist: torch>=2.6.0
Provides-Extra: examples
Requires-Dist: lightning; extra == "examples"
Requires-Dist: matplotlib; extra == "examples"
Requires-Dist: scipy; extra == "examples"
Provides-Extra: laplace
Requires-Dist: curvlinops-for-pytorch<3; python_version < "3.13" and extra == "laplace"
Requires-Dist: laplace-torch; python_version < "3.13" and extra == "laplace"
Provides-Extra: comparisons
Requires-Dist: curvlinops-for-pytorch<3; python_version < "3.13" and extra == "comparisons"
Requires-Dist: laplace-torch; python_version < "3.13" and extra == "comparisons"
Requires-Dist: bayesian-torch; extra == "comparisons"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: sphinx; extra == "dev"
Requires-Dist: myst-parser; extra == "dev"
Requires-Dist: furo; extra == "dev"
Dynamic: license-file

# FastBNNs

## Introduction
FastBNNs implements fast and flexible Bayesian inference of neural networks based on propagation of statistical moments [1] and the unscented transform [2] in PyTorch.
FastBNNs enables one-line conversion of many PyTorch-based models to Bayesian counterparts, whereby learnable model parameters are treated as Normal random variables.

## Installation
FastBNNs is primarily built around PyTorch.
To install FastBNNs, first install PyTorch

```
pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu129
```

then install FastBNNs as 

```
pip install fastbnns
```


Alternatively, a [requirements.txt](requirements.txt) file is provided to support installation of FastBNNs.
To install in a virtual environment, run

```
python -m venv .venv
source .venv/bin/activate  # bash
# .venv/Scripts/Activate.ps1  # Windows PowerShell
pip install -r requirements.txt
```

## Basic usage
A neural network `nn: torch.nn.Module` can be converted to a Bayesian neural network using the
`bnn.base.BNN` wrapper class:

```
import torch

from fastbnns.bnn.base import BNN
from fastbnns.models.mlp import MLP

hidden_features = 32
n_hidden_layers = 1
in_features = 1
out_features = 1
nn = MLP(
    in_features=in_features,
    out_features=out_features,
    n_hidden_layers=n_hidden_layers,
    hidden_features=hidden_features,
    activation=torch.nn.LeakyReLU,
)
bnn = BNN(nn=nn, convert_in_place=False)
```

Forward calls through `bnn` can be made identically to `nn`:

```
data = torch.randn((1, in_features), dtype=torch.float32)
out_nn = nn(data)
out_bnn = bnn(data)
```

In this usage, a single sample of each network parameter is made in `bnn` before computing the forward computation identically to `nn`.
As such, multiple network samples can be made to characterize the output distribution as

```
n_samples = 100
out_bnn_mc = torch.stack([bnn(data) for _ in range(n_samples)])
out_bnn_mc_mean = out_bnn_mc.mean(dim=0)
out_bnn_mc_var = out_bnn_mc.var(dim=0)
```

Alternatively, to leverage the fast inference methods (i.e., non-sampling-based), the network input can be wrapped in the custom type `bnn.types.MuVar`:

```
out_bnn_fast = out_bnn(bnn.types.MuVar(data))
out_bnn_fast_mean = out_bnn_fast.mu
out_bnn_fast_var = out_bnn_fast.var
```


## Bayesian treatment
The wrapped model `bnn: torch.nn.Module` is still an instance of `torch.nn.Module` and can be trained using standard PyTorch or PyTorch Lightning strategies.
However, a Bayesian treatment of `bnn` training requires use of a custom loss function, such as the evidence lower bound (ELBO) used in Bayes-by-backprop [3].
Examples of training the Bayesian MLP from [Basic Usage](#basic-usage) using the ELBO loss are provided in [PyTorch](examples/mlp.py) and [PyTorch Lightning](examples/mlp_lightning.py).


## 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].

[3] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight Uncertainty
in Neural Networks, May 2015. arXiv:1505.05424 [cs, stat]

## Copyright

LANL O4956

&copy; 2025. Triad National Security, LLC. All rights reserved.

This program was produced under U.S. Government contract 89233218CNA000001 for Los Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC for the U.S. Department of Energy/National Nuclear Security Administration. All rights in the program are reserved by Triad National Security, LLC, and the U.S. Department of Energy/National Nuclear Security Administration. The Government is granted for itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license in this material to reproduce, prepare. derivative works, distribute copies to the public, perform publicly and display publicly, and to permit others to do so.
