<jupyter_start><jupyter_text>Deterministic Variational Inference: Toy exampleCode to accompany [**Fixing Variational Bayes: Deterministic Variational Inference for Bayesian Neural Networks**](https://arxiv.org/abs/1810.03958)<jupyter_code>%matplotlib inline

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from IPython.display import clear_output
import os, json 

import gaussian_variables as gv
import utils as u
import plot_utils as pu
import bayes_layers as bnn
from bayes_models import MLP, PointMLP, AdaptedMLP

np.random.seed(3)<jupyter_output><empty_output><jupyter_text>DatasetFirst we generate a toy dataset according to:\begin{equation}    y = -(x+0.5)\sin(3\pi x) + \eta\end{equation}Where the noise is generated according to:\begin{equation}    \eta = 0.45(x + 0.5)^2\end{equation}<jupyter_code>def base_model(x):
    return -(x+0.5)*np.sin(3 * np.pi *x)

def noise_model(x):
    return 0.45*(x+0.5)**2

def sample_data(x):
    return base_model(x) + np.random.normal(0, noise_model(x))

data_size = {'train': 500, 'valid': 100, 'test': 100}
toy_data = []
for section in ['train', 'valid', 'test']:
    x = (np.random.rand(data_size['train'], 1) - 0.5)
    toy_data.append([x, sample_data(x).reshape(-1)])    
x = np.arange(-1,1,1/100)
toy_data.append([[[_] for _ in x], base_model(x)])

pu.toy_results_plot(toy_data, {'mean':base_model, 'std':noise_model})<jupyter_output><empty_output><jupyter_text>ModelHere we set the hyperparameters and provide a function for constructing a Bayesian Neural Network.<jupyter_code>hypers = {
    "x_dim": 1,                                "y_dim": 2,
    "hidden_dims": [128,128],                  "nonlinearity": "relu",
    "adapter": {
        'in' : {"scale": [[1.0]], "shift": [[0.0]]},
        'out': {"scale": [[1.0, 0.83]], "shift": [[0.0, -3.5]]}
    },
    
    "method": "bayes",                         "style": "heteroskedastic",
    "homo_logvar_scale": 2*np.log(0.2),        "prior_type": ["empirical", "wider_he", "wider_he"],
    
    "n_epochs": 20000,                         "batch_size": 500,
    "learning_rate": 0.001,                    "lambda": 1.0,
    "warmup_updates": {'lambda': 14000.0},     "anneal_updates": {'lambda': 1000.0},
    "optimizer": "adam",                       "gradient_clip": 0.1,

    "data_fraction": 1.0,                      "sections_to_run": ["train"]
}
def make_model(hypers):
    if hypers['method'].lower().strip() == 'bayes':
        MLP_factory = MLP
        prediction = lambda y: tf.reshape(y.mean[:,0], [-1])
        loss = bnn.regression_loss
    else:
        MLP_factory = PointMLP
        prediction = lambda y: tf.reshape(y.mean[:,0], [-1])
        loss = bnn.point_regression_loss
        
    mlp = MLP_factory(hypers['x_dim'], hypers['y_dim'], hypers)
    mlp = AdaptedMLP(mlp)
    mlp.make_placeholders()
    ipt = mlp.placeholders['ipt_mean']
    y = mlp(ipt)
a
    target = tf.placeholder(tf.float32, [None])
    mlp.placeholders['target'] = target
    global_step = tf.Variable(0, trainable=False, name='global_step')
    loss, logprob, all_surprise = loss(y, target, mlp, hypers, global_step)

    accuracy = tf.reduce_mean(tf.abs(target - prediction(y)))
    
    return {
        'model': mlp,
        'metrics': {
            'accuracy': accuracy,   'loss': loss,
            'logprob': logprob,     'all_surprise': all_surprise
            },
        'global_step': global_step}<jupyter_output><empty_output><jupyter_text>Training loopFor the first `warmup_updates` epochs we run with the KL term in the elbo turned off. You will observe that the model starts to overfit the data. The KL term is then ramped back on and the model is trained with the full ELBO to convergence.<jupyter_code>run_id = u.start_run()

data = toy_data

restricted_training_set = u.restrict_dataset_size(data[0], hypers['data_fraction'])
hypers['dataset_size'] = len(restricted_training_set[0])

device_id = 0
device_string = u.get_device_string(device_id)
print(hypers)
with tf.device(device_string):
    if True:
        model_and_metrics = make_model(hypers)
        train_op = u.make_optimizer(model_and_metrics, hypers)
        sess = u.get_session()
        saver = tf.train.Saver()

        all_summaries = []
        best_valid_accuracy = np.inf
    
    for epoch in range(hypers['n_epochs']):
        verbose = (epoch % 20 == 0)
        if verbose:
            print("Epoch %i:        " % epoch, end='')
        epoch_summary, accuracies = u.train_valid_test(
            {
                'train': restricted_training_set,
                'valid': data[1],
                'test': data[2]
            },
            sess, model_and_metrics, train_op, hypers, verbose)

        if epoch % 200 == 0:
            clear_output()
            predictions = u.get_predictions(data[3], sess, model_and_metrics['model'], hypers)
            pu.toy_results_plot(toy_data, {'mean':base_model, 'std':noise_model}, hypers, predictions)
            plt.show()

        # dump log file
        all_summaries.append(epoch_summary)
        #with open(os.path.join("logs", run_id, "results.json"), 'w') as f:
        #    json.dump(all_summaries, f, indent=4, cls=u.NumpyEncoder)<jupyter_output><empty_output><jupyter_text>Comparison with MCWe compare the output activations of the BNN as predicted using our deterministic approximation with the activations predicted by Monte Carlo approximation (see Figure 3 of the paper)<jupyter_code>from matplotlib import rc
rc('text', usetex=True)

n_samp = 20000
x=0.25
ipt = [[[x]] for _ in range(n_samp)]
sample_op = model_and_metrics['model'].run_with_MC(ipt, n_sample=n_samp)
approx_op = model_and_metrics['model'](x)

samples = sess.run(sample_op)
approx = sess.run([approx_op.mean, approx_op.var])

def gaussian1d(mean, var):
    x_axis = np.linspace(-5,5,1000)
    return x_axis, 1.0 / np.sqrt(2.0 * np.pi * var) * np.exp(-1.0 / (2.0 * var) * (x_axis - mean)**2)


samples.shape
plt.figure(figsize=(3,2))
plt.hist(samples[:,0,0], int(np.round(200/(1-0.3))), density=True, edgecolor='k', facecolor='#b4c7e7')
plt.plot(*gaussian1d(approx[0][0,0], approx[1][0,0,0]), 'b')
plt.xlim([-1,-0.3])
plt.yticks([])
plt.xlabel('$m$')
plt.ylabel('$q(m)$')
plt.show()

plt.figure(figsize=(3,2))
plt.hist(samples[:,0,1], int(np.round(200/(3.5-1.9))), density=True, edgecolor='k', facecolor='#b4c7e7', label="MC")
plt.plot(*gaussian1d(approx[0][0,1], approx[1][0,1,1]), 'b', label="ours")
plt.xlim([-3.5,-1.9])
plt.yticks([])
plt.xlabel('$\ell$')
plt.ylabel('$q(\ell)$')
plt.legend()
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1566155248367006, 'fraction_numerical': 0.029925565851435517, 'mean_word_length': 3.42175957018133, 'pattern_counts': {'":': 23, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '287066', 'license': 'MIT', 'n_tokens_mistral': 2501, 'n_tokens_neox': 2350, 'n_words': 496, 'path': 'ToyData.ipynb', 'repo_name': 'yunndlalala/MCS-project', 'size': '6583'}