same board but with gpu

gpu-gradient.py

import tensorflow as tf
import numpy as np
import os
import uuid

TENSORBOARD_PATH = "/tmp/tensorboard-switchpoint"
# tensorboard --logdir=/tmp/tensorboard-switchpoint

x1 = np.random.randn(35) - 1
x2 = np.random.randn(35) * 2 + 5
x_all = np.hstack([x1, x2])
len_x = len(x_all)
time_all = np.arange(1, len_x + 1)

with tf.device("/gpu:0"):
    mu1 = tf.Variable(0, name="mu1", dtype=tf.float32)
    mu2 = tf.Variable(0, name="mu2", dtype=tf.float32)
    sigma1 = tf.Variable(2, name="sigma1", dtype=tf.float32)
    sigma2 = tf.Variable(2, name="sigma2", dtype=tf.float32)
    tau = tf.Variable(15, name="tau", dtype=tf.float32)
    switch = 1. / (1 + tf.exp(tf.pow(time_all - tau, 1)))
    mu = switch * mu1 + (1 - switch) * mu2
    sigma = switch * sigma1 + (1 - switch) * sigma2

    likelihood_arr = tf.log(tf.sqrt(1/(2*np.pi*tf.pow(sigma, 2)))) - tf.pow(x_all - mu, 2)/(2*tf.pow(sigma, 2))
    total_likelihood = tf.reduce_sum(likelihood_arr, name="total_likelihood")

    optimizer = tf.train.AdamOptimizer()
    opt_task = optimizer.minimize(-total_likelihood)
    init = tf.global_variables_initializer()

    tf.summary.scalar("mu1", mu1)
    tf.summary.scalar("mu2", mu2)
    tf.summary.scalar("sigma1", sigma1)
    tf.summary.scalar("sigma2", sigma2)
    tf.summary.scalar("tau", tau)
    tf.summary.scalar("likelihood", total_likelihood)
    merged_summary_op = tf.summary.merge_all()

with tf.Session() as sess:
    sess.run(init)
    uniq_id = os.path.join(TENSORBOARD_PATH, "switchpoint-gpu-" + uuid.uuid1().__str__()[:4])
    summary_writer = tf.summary.FileWriter(uniq_id, graph=tf.get_default_graph())
    for step in range(15000):
        lik, opt, summary = sess.run([total_likelihood, opt_task, merged_summary_op])
        if step % 100 == 0:
            variables = {_.name:_.eval() for _ in [total_likelihood]}
            summary_writer.add_summary(summary, step)
            print("i{}: {}".format(str(step).zfill(5), variables))