Include training, testing and uncertainty quantification snippet in README.md

Author: ranganathkrishnan

README.md

@@ -1,5 +1,5 @@
 # Bayesian-Torch: Bayesian neural network layers for uncertainty estimation
-**[Get started](#Installation)** | **[Example usage](#example-usage)** | **[Documentation](doc/bayesian_torch.layers.md)** | **[License](LICENSE)** | **[Citing](#citing)** 
+**[Get started](#installation)** | **[Example usage](#example-usage-training-and-evaluation-of-models)** | **[Documentation](doc/bayesian_torch.layers.md)** | **[License](LICENSE)** | **[Citing](#citing)** 
 
 ### Bayesian layers and utilities to perform stochastic variational inference in PyTorch
 
@@ -38,9 +38,9 @@ The repository has implementations for the following Bayesian layers:
 Please refer to [documentation](doc/bayesian_torch.layers.md#layers) of Bayesian layers for details.
 
 Other features include:
+- [x] [dnn_to_bnn()](https://github.com/IntelLabs/bayesian-torch/blob/main/bayesian_torch/models/dnn_to_bnn.py#L127): An API to convert deterministic deep neural network (dnn) model of any architecture to Bayesian deep neural network (bnn) model, simplifying the model definition i.e. drop-in replacements  of Convolutional, Linear and LSTM layers to corresponding Bayesian layers. This will enable seamless conversion of existing topology of larger models to Bayesian deep neural network models for extending towards uncertainty-aware applications. 
+- [x] [MOPED](https://github.com/IntelLabs/bayesian-torch/blob/main/bayesian_torch/utils/util.py#L72): Specifying weight priors and variational posteriors in Bayesian neural networks with Empirical Bayes [[Krishnan et al. 2020](https://ojs.aaai.org/index.php/AAAI/article/view/5875)]
 - [x] [AvUC](https://github.com/IntelLabs/bayesian-torch/blob/main/bayesian_torch/utils/avuc_loss.py): Accuracy versus Uncertainty Calibration loss [[Krishnan and Tickoo 2020](https://proceedings.neurips.cc/paper/2020/file/d3d9446802a44259755d38e6d163e820-Paper.pdf)]
-- [x] [MOPED](https://github.com/IntelLabs/bayesian-torch/blob/main/bayesian_torch/utils/util.py#L72): Specifying weight priors and variational posteriors with Empirical Bayes [[Krishnan et al. 2020](https://ojs.aaai.org/index.php/AAAI/article/view/5875)]
-- [x] [dnn_to_bnn](https://github.com/IntelLabs/bayesian-torch/blob/main/bayesian_torch/models/dnn_to_bnn.py#L127): An API to convert deterministic deep neural network (dnn) model of any architecture to Bayesian deep neural network (bnn) model, simplifying the model definition i.e. drop-in replacements  of Convolutional, Linear and LSTM layers to corresponding Bayesian layers. This will enable seamless conversion of existing topology of larger models to Bayesian deep neural network models for extending towards uncertainty-aware applications. 
 
 ## Installation
 <!--
@@ -55,6 +55,7 @@ git clone https://github.com/IntelLabs/bayesian-torch
 cd bayesian-torch
 pip install .
 ```
+<!--
 This code has been tested on PyTorch v1.8.1 LTS.
 
 Dependencies:
@@ -64,16 +65,17 @@ Dependencies:
 - conda install -c conda-forge accimage
 - pip install tensorboard
 - pip install scikit-learn
-
+-->
 ## Usage
 There are two ways to build Bayesian deep neural networks using Bayesian-Torch: 
 1. Convert an existing deterministic deep neural network (dnn) model to Bayesian deep neural network (bnn) model with dnn_to_bnn()
 2. Define your custom model using the Bayesian layers ([Flipout](https://github.com/IntelLabs/bayesian-torch/tree/main/bayesian_torch/layers/flipout_layers) or [Reparameterization](https://github.com/IntelLabs/bayesian-torch/tree/main/bayesian_torch/layers/variational_layers))
 
-(1) For instance to build Bayesian-ResNet18 from torchvision deterministic ResNet18 model:
+(1) For instance, building Bayesian-ResNet18 from torchvision deterministic ResNet18 model is as simple as:
 ```
+import torch
 import torchvision
-from bayesian_torch.models.dnn_to_bnn import dnn_to_bnn
+from bayesian_torch.models.dnn_to_bnn import dnn_to_bnn, get_kl_loss
 
 const_bnn_prior_parameters = {
         "prior_mu": 0.0,
@@ -82,13 +84,13 @@ const_bnn_prior_parameters = {
         "posterior_rho_init": -3.0,
         "type": "Reparameterization",  # Flipout or Reparameterization
         "moped_enable": False,  # True to initialize mu/sigma from the pretrained dnn weights
-        "moped_delta": 0.2,
+        "moped_delta": 0.5,
 }
     
 model = torchvision.models.resnet18()
 dnn_to_bnn(model, const_bnn_prior_parameters)
 ```
-To use MOPED method, setting the prior and initializing variational parameters from a pretrained determined model (helps training convergence of larger models):
+To use MOPED method, setting the prior and initializing variational parameters from a pretrained deterministic model (helps training convergence of larger models):
 ```
 const_bnn_prior_parameters = {
         "prior_mu": 0.0,
@@ -97,12 +99,47 @@ const_bnn_prior_parameters = {
         "posterior_rho_init": -3.0,
         "type": "Reparameterization",  # Flipout or Reparameterization
         "moped_enable": True,  # True to initialize mu/sigma from the pretrained dnn weights
-        "moped_delta": 0.2,
+        "moped_delta": 0.5,
 }
     
 model = torchvision.models.resnet18(pretrained=True)
 dnn_to_bnn(model, const_bnn_prior_parameters)
 ```
+Training snippet:
+```
+criterion = torch.nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(model.parameters(), args.learning_rate)
+
+output = model(x_train)
+kl = get_kl_loss(model)
+ce_loss = criterion(output, y_train)
+loss = ce_loss + kl / args.batch_size 
+
+loss.backward()
+optimizer.step()
+```
+Testing snippet:
+```
+model.eval()
+with torch.no_grad():
+    output_mc = []
+    for mc_run in range(args.num_monte_carlo):
+        logits = model(x_test)
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        output_mc.append(probs)
+    output = torch.stack(output_mc)  
+    pred_mean = output.mean(dim=0)
+    y_pred = torch.argmax(pred_mean, axis=-1)
+    test_acc = (y_pred.data.cpu().numpy() == y_test.data.cpu().numpy()).mean()
+```
+Uncertainty Quantification:
+```
+from utils.util import predictive_entropy, mutual_information
+
+predictive_uncertainty = predictive_entropy(output.data.cpu().numpy())
+model_uncertainty = mutual_information(output.data.cpu().numpy())
+```
+
 (2) For building custom models, we have provided [example model implementations](bayesian_torch/models/bayesian) using the Bayesian layers.
 
 ## Example usage (training and evaluation of models)