Created
May 2, 2020 09:54
-
-
Save roblem/8472d29e26db644ebc160645166d131f to your computer and use it in GitHub Desktop.
Ordinary least squares tensorflow probability benchmarks (MHRW)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import sys | |
print("Running in :", sys.executable) | |
import tensorflow as tf | |
print("TF devices: ", tf.config.list_physical_devices()) | |
import tensorflow_probability as tfp | |
from tensorflow_probability import distributions as tfd | |
import numpy as np | |
import pandas as pd | |
import time as time | |
# set tensorflow data type | |
dtype = tf.float32 | |
## | |
## simple OLS Data Generation Process | |
## | |
# True beta | |
N = 50000 | |
K = 500 | |
b = np.random.randn(K) | |
b[0] = b[0] + 3 | |
# True error std deviation | |
sigma_e = 1 | |
x = np.c_[np.ones(N), np.random.randn(N,K-1)] | |
y = x.dot(b) + sigma_e * np.random.randn(N) | |
# estimate parameter vector, errors, sd of errors, and se of parameters | |
bols = np.linalg.inv(x.T.dot(x)).dot(x.T.dot(y)) | |
err = y - x.dot(bols) | |
sigma_ols = np.sqrt(err.dot(err)/(x.shape[0] - x.shape[1])) | |
se = np.sqrt(err.dot(err)/(x.shape[0] - x.shape[1]) * np.diagonal(np.linalg.inv(x.T.dot(x)))) | |
# put results together for easy viewing | |
ols_parms = np.r_[bols, sigma_ols] | |
ols_se = np.r_[se, np.nan] | |
print("\n") | |
indexn = ['b'+str(i) for i in range(K)] | |
indexn.extend(['sigma']) | |
print(pd.DataFrame(np.c_[ols_parms, ols_se],columns=['estimate', 'std err'], | |
index=indexn)) | |
print("\n\n") | |
X = tf.constant(x, dtype=dtype) | |
Y = tf.constant(y, dtype=dtype) | |
N_ = tf.constant(N, dtype=dtype) | |
pi = tf.constant(np.pi, dtype=dtype) | |
nsamples = tf.constant(1000, dtype=tf.int32) | |
nburnin = tf.constant(500, dtype=tf.int32) | |
# initialize | |
init = [tf.constant(np.random.randn(K), dtype=dtype), tf.constant(1., dtype=dtype)] | |
## | |
## Model Log-Likelihood/Posterior | |
## | |
@tf.function(experimental_compile=True) | |
def ols_loglike(beta, sigma): | |
# xb (mu_i for each observation) | |
mu = tf.linalg.matvec(X, beta) | |
# this is normal pdf logged and summed over all observations | |
ll = - (N_/2.)*tf.math.log(2.*pi*sigma**2) -\ | |
(1./(2.*sigma**2.))*tf.math.reduce_sum((Y-mu)**2., axis=-1) | |
return ll | |
# | |
# Evaluate speed of function evals (no tfp required) | |
# | |
with tf.device('/CPU:0'): | |
ll = ols_loglike(init[0], init[1]) | |
startt = time.time() | |
ll = ols_loglike(init[0], init[1]) | |
endt = time.time() | |
print("\n\nLogL calculation in %2.2f MS on CPU"% ((endt - startt)*1000)) | |
print("\n\n") | |
try: | |
with tf.device('/GPU:0'): | |
ll = ols_loglike(init[0], init[1]) | |
startt = time.time() | |
ll = ols_loglike(init[0], init[1]) | |
endt = time.time() | |
print("\n\nLogL calculation in %2.2f MS on GPU"% ((endt - startt)*1000)) | |
print("\n\n") | |
except: | |
print("GPU not available in this python environment") | |
@tf.function(experimental_relax_shapes=True, experimental_compile=True) | |
def mcmc_sampler(init, nsamples=1, nburnin=1): | |
@tf.function | |
def ols_loglike_(beta_, sigma_): | |
# xb (mu_i for each observation) | |
mu = tf.linalg.matvec(X, beta_) | |
# this is normal pdf logged and summed over all observations | |
ll = - (N_/2.)*tf.math.log(2.*pi*sigma_**2) -\ | |
(1./(2.*sigma_**2.))*tf.math.reduce_sum((Y-mu)**2., axis=-1) | |
return ll | |
kernel = tfp.mcmc.RandomWalkMetropolis( | |
target_log_prob_fn=ols_loglike_, | |
new_state_fn=tfp.mcmc.random_walk_normal_fn(scale=.04),seed=42) | |
samples_, stats_ = tfp.mcmc.sample_chain( | |
num_results=nsamples, | |
current_state= init, | |
kernel=kernel, | |
num_burnin_steps=nburnin, | |
parallel_iterations=10) | |
return samples_, stats_ | |
with tf.device('/CPU:0'): | |
samples, stats = mcmc_sampler(init, 1, 1) | |
startt = time.time() | |
samples, stats = mcmc_sampler(init, nsamples, nburnin) | |
endt = time.time() | |
print("\n\nMHRW completed in %2.2f seconds on CPU"% (endt - startt)) | |
print("\n\n") | |
try: | |
with tf.device('/GPU:0'): | |
samples, stats = mcmc_sampler(init, 1, 1) | |
startt = time.time() | |
samples, stats = mcmc_sampler(init, nsamples, nburnin) | |
endt = time.time() | |
print("\n\nMHRE completed in %2.2f seconds on GPU"% (endt - startt)) | |
print("\n\n") | |
except: | |
print("GPU not available in this python environment") | |
trace_sigman = samples[1].numpy() | |
trace_betan = samples[0].numpy() | |
est_nuts = np.r_[trace_betan.mean(axis=0), trace_sigman.mean()] | |
std_nuts = np.r_[trace_betan.std(axis=0), trace_sigman.std()] | |
# assemble and print | |
print(pd.DataFrame(np.c_[est_nuts, std_nuts],columns=['estimate', 'std err'], | |
index=indexn)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment