test_script.py

import tensorflow as tf
import edward as ed
import numpy as np

from edward.models import Normal


def build_toy_dataset(N, w, noise_sd=0.1, data_sd=1):
    D = len(w)
    x = np.random.normal(0, data_sd, size=(N, D))
    y = np.dot(x, w) + np.random.normal(0, noise_sd, size=N)
    return x, y


### Generate the data
# Note that data_sd >> noise_sd
N = 1000
D = 5

w_true = np.random.randn(D)
noise_sd = 0.1
data_sd = 5
X_train, y_train = build_toy_dataset(N, w_true, noise_sd=noise_sd, data_sd=data_sd)


### Define the model
X = tf.placeholder(tf.float32, [N, D])
w = Normal(loc=tf.zeros(D), scale=tf.ones(D))
b = Normal(loc=tf.zeros(1), scale=tf.ones(1))
log_sd = Normal(loc=tf.zeros(1), scale=tf.ones(1))
y = Normal(loc=ed.dot(X, w) + b, scale=tf.exp(log_sd) * tf.ones(N))

qw = Normal(loc=tf.get_variable("qw/loc", [D]),
            scale=tf.nn.softplus(tf.Variable(tf.zeros(D)-10., name='qw/scale')) + 1e-8)
qb = Normal(loc=tf.get_variable("qb/loc", [1]),
            scale=tf.nn.softplus(tf.get_variable("qb/scale", [1])) + 1e-8)
qlog_sd = Normal(loc=tf.get_variable("qlog_sd/loc", [1]),
            scale=tf.nn.softplus(tf.Variable(-10., name='qlog_sd/scale')) + 1e-8)


### Variational Inference
# Many samples and iterations, just to go sure
optimizer = tf.train.AdamOptimizer(0.1)
inference = ed.KLqp({w: qw, b: qb, log_sd: qlog_sd}, data={X: X_train, y: y_train})
inference.run(n_samples=1, n_iter=2000, optimizer=optimizer)


### Print estimates
sess = ed.get_session()

# w is unbiased...
print("w hat: {}".format(sess.run(qw.mean())))
print("w true: {}".format(w_true))

# sd is overestimated..
print("sd hat: {}".format(sess.run(tf.exp(qlog_sd.mean()))))
print("sd true: {}".format(noise_sd))

import os
import json
if os.path.exists('ed_logs/results.json'):
    with open('ed_logs/results.json', 'r') as f:
        results = json.load(f)
    results['sd_hat'].append(float(sess.run(tf.exp(qlog_sd.mean()))[0]))
else:
    results = {'sd_hat': [float(sess.run(tf.exp(qlog_sd.mean()))[0])]}
with open('ed_logs/results.json', 'w') as f:
    json.dump(results, f)