# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# 测试DVI正态近似的准确性

# +
# %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)


# -

# 首先构建数据集：
# \begin{equation}
#     y = -(x+0.5)\sin(3\pi x) + \eta
# \end{equation}
#
#
# \begin{equation}
#     \eta = 0.45(x + 0.5)^2
# \end{equation}

# +
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})
# -

# 设置模型参数，这里以一个三层、每层5个神经元的网络为例(本项目中的其他网络结果修改“hidden_dims”参数可得到)

hypers = {
    "x_dim": 1,                                "y_dim": 2,
    "hidden_dims": [5,5,5],                  "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)
    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}


# 给出一个比较DVI和MCVI算法结果的函数

from scipy import stats
def show_compare(model_and_metrics, sess):
    plt.figure()
    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])

    # samples_b.shape
    m_min = stats.norm.ppf(
        0.0001, loc=approx[0][0, 0], scale=np.sqrt(approx[1][0, 0, 0]))
    m_max = stats.norm.ppf(
        0.9999, loc=approx[0][0, 0], scale=np.sqrt(approx[1][0, 0, 0]))
    l_min = stats.norm.ppf(
        0.0001, loc=approx[0][0, 1], scale=np.sqrt(approx[1][0, 1, 1]))
    l_max = stats.norm.ppf(
        0.9999, loc=approx[0][0, 1], scale=np.sqrt(approx[1][0, 1, 1]))

    bin_no_m = np.linspace(m_min, m_max, 50)
    bin_no_l = np.linspace(l_min, l_max, 50)

    fig = plt.figure()

    ax1 = fig.add_subplot(221)
    ax1.hist(samples[:, 0, 0], bin_no_m,
             density=True, edgecolor='k', facecolor='#b4c7e7')
    ax1.plot(*gaussian1d(approx[0][0, 0], approx[1][0, 0, 0],
                         m_min, m_max), 'b')
    plt.xlim([m_min, m_max])
    ax1.set_yticks([])
    ax1.set_xlabel('$m$')
    ax1.set_ylabel('$q(m)$')

    ax2 = fig.add_subplot(222)
    ax2.hist(samples[:, 0, 1], bin_no_l,
             density=True, edgecolor='k', facecolor='#b4c7e7', label="MC")

    ax2.plot(*gaussian1d(approx[0][0, 1],
                         approx[1][0, 1, 1],
                         l_min,
                         l_max), 'b', label="ours")
    plt.xlim([l_min, l_max])
    ax2.set_yticks([])
    ax2.set_xlabel('$\ell$')
    ax2.set_ylabel('$q(\ell)$')
    plt.show()
    return None


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


# 定义训练函数

def run(data):
    run_id = u.start_run()

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

    device_id = 1
    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

        show_compare(model_and_metrics, sess)

        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)

        show_compare(model_and_metrics, sess)



# 开始训练，并得到两个DVI和MCVI之间的比较

run(toy_data)


# 将之前的数据集整体增大10倍

# +
def base_model_10(x):
    return 10*(-(x+0.5)*np.sin(3 * np.pi *x))

def noise_model_10(x):
    return 10*(0.45*(x+0.5)**2)

def sample_data_10(x):
    return base_model_10(x) + np.random.normal(0, noise_model_10(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_10(x).reshape(-1)])    
x = np.arange(-1,1,1/100)
toy_data.append([[[_] for _ in x], base_model_10(x)])

pu.toy_results_plot(toy_data, {'mean':base_model_10, 'std':noise_model_10})
plt.ylim([-11,16])
# -

# 同样的方式进行训练，并得到DVI和MCVI之间的比较

run(toy_data)

# 对于m的结果，DVI和MCVI的一致性变强了。



StarcoderdataJupyterScriptsDedupFiltered