.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/scientific_machine_learning/bbb_trainer/bbbtrainer_quadratic.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code or to run this example in your browser via Binder .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_scientific_machine_learning_bbb_trainer_bbbtrainer_quadratic.py: Training a Bayesian neural network ============================================================= In this example, we train a Bayesian neural network to learn the function :math:`f(x)=x^2` .. GENERATED FROM PYTHON SOURCE LINES 10-11 First, we have to import the necessary modules. .. GENERATED FROM PYTHON SOURCE LINES 14-24 .. code-block:: default # Default imports import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import matplotlib.pyplot as plt import UQpy.scientific_machine_learning as sml torch.manual_seed(123) .. GENERATED FROM PYTHON SOURCE LINES 25-27 We define the network architecture using the ``nn.Sequential`` object and instantiate the ``BayesianNeuralNetwork``. .. GENERATED FROM PYTHON SOURCE LINES 30-41 .. code-block:: default width = 8 network = nn.Sequential( sml.BayesianLinear(1, width), nn.ReLU(), sml.BayesianLinear(width, width), nn.ReLU(), sml.BayesianLinear(width, 1), ) model = sml.FeedForwardNeuralNetwork(network) .. GENERATED FROM PYTHON SOURCE LINES 42-48 With the neural network defined, we turn our attention to the training data. We want to learn the function :math:`f(x)=x^2` and define the training data using the pytorch Dataset and Dataloader. For more information on defining the training data, see the pytorch documentation at https://pytorch.org/tutorials/beginner/basics/data_tutorial.html .. GENERATED FROM PYTHON SOURCE LINES 51-66 .. code-block:: default class QuadraticDataset(Dataset): def __init__(self, n_samples=200): self.n_samples = n_samples self.x = torch.linspace(-5, 5, n_samples, dtype=torch.float).reshape(-1, 1) self.y = self.x**2 def __len__(self): return self.n_samples def __getitem__(self, item): return self.x[item], self.y[item] .. GENERATED FROM PYTHON SOURCE LINES 67-68 Before we continue with training the network, let's get the initial prediction of the neural network on the data. .. GENERATED FROM PYTHON SOURCE LINES 71-74 .. code-block:: default initial_prediction = model(QuadraticDataset().x) .. GENERATED FROM PYTHON SOURCE LINES 75-78 So far we have the neural network and training data. The ``BBBTrainer`` combines the two along with a pytorch optimization algorithm to learn the network parameters. We instantiate the ``BBBTrainer``, train the network, then print the initial and final loss alongside a model summary. .. GENERATED FROM PYTHON SOURCE LINES 81-93 .. code-block:: default optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) train_data = DataLoader(QuadraticDataset(), batch_size=40, shuffle=True) trainer = sml.BBBTrainer(model, optimizer) print("Starting Training...", end="") trainer.run(train_data=train_data, epochs=500, beta=1e-5, num_samples=10) print("done") print("Initial loss:", trainer.history["train_loss"][0]) print("Final loss:", trainer.history["train_loss"][-1]) model.summary() .. GENERATED FROM PYTHON SOURCE LINES 94-95 We compare the initial and final predictions and plot the loss history using matplotlib. .. GENERATED FROM PYTHON SOURCE LINES 98-136 .. code-block:: default x = QuadraticDataset().x y = QuadraticDataset().y model.train(False) model.sample(False) final_prediction = model(x) fig, ax = plt.subplots() ax.plot( x.detach().numpy(), initial_prediction.detach().numpy(), label="Initial Prediction", color="tab:blue", ) ax.plot( x.detach().numpy(), final_prediction.detach().numpy(), label="Final Prediction", color="tab:orange", ) ax.plot( x.detach().numpy(), y.detach().numpy(), label="Exact", color="black", linestyle="dashed", ) ax.set_title("Initial and Final NN Predictions") ax.set(xlabel="x", ylabel="f(x)") ax.legend() train_loss = trainer.history["train_loss"].detach().numpy() fig, ax = plt.subplots() ax.semilogy(train_loss) ax.set_title("Bayes By Backpropagation Training Loss") ax.set(xlabel="Epoch", ylabel="Loss") plt.show() .. GENERATED FROM PYTHON SOURCE LINES 137-144 The Bayesian neural network is a probabilistic model. Each of its parameters, in this case weights and biases, are governed by Gaussian distributions. We can get a deterministic output from the BNN by setting ``model.sample(False)``, which sets each parameter to the mean of its distribution. We can obtain error bars on model's output by sampling the parameters from their governing distribution. This is done by setting ``model.sample(True)`` and computing the forward model evaluation many times, then computing the sample variance .. GENERATED FROM PYTHON SOURCE LINES 146-175 .. code-block:: default model.sample(False) print("BNN is deterministic:", model.is_deterministic()) mean = model(x) model.sample(True) print("BNN is deterministic:", model.is_deterministic()) n = 10_000 samples = torch.zeros(len(x), n) for i in range(n): samples[:, i] = model(x).squeeze() variance = torch.var(samples, dim=1) standard_deviation = torch.sqrt(variance) x_plot = x.squeeze().detach().numpy() mu = mean.squeeze().detach().numpy() sigma = standard_deviation.squeeze().detach().numpy() fig, ax = plt.subplots() ax.plot(x_plot, mu, label="$\mu$") ax.plot(x_plot, y.detach().numpy(), label="Exact", color="black", linestyle="dashed") ax.fill_between( x_plot, mu - (3 * sigma), mu + (3 * sigma), label="$\mu \pm 3\sigma$,", alpha=0.3 ) ax.set_title("Bayesian Neural Network $\mu \pm 3\sigma$") ax.set(xlabel="x", ylabel="f(x)") ax.legend() plt.show() .. rst-class:: sphx-glr-timing **Total running time of the script:** ( 0 minutes 0.000 seconds) .. _sphx_glr_download_auto_examples_scientific_machine_learning_bbb_trainer_bbbtrainer_quadratic.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: binder-badge .. image:: images/binder_badge_logo.svg :target: https://mybinder.org/v2/gh/SURGroup/UQpy/master?urlpath=lab/tree/notebooks/auto_examples/scientific_machine_learning/bbb_trainer/bbbtrainer_quadratic.ipynb :alt: Launch binder :width: 150 px .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: bbbtrainer_quadratic.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: bbbtrainer_quadratic.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_