import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.python.client import device_lib
from tensorflow.python.ops import variable_scope



def get_available_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']



mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)



gpus = get_available_gpus()
trainers = []
accs = []
xs = []
ys = []



for i_, gpu_id in enumerate(gpus):

    with tf.device(gpu_id):
        # [Build graph in here.]

        with variable_scope.variable_scope(variable_scope.get_variable_scope(), reuse=i_>0):

            x = tf.placeholder(tf.float32, [None, 784])
            W = tf.Variable(tf.zeros([784, 10]))
            b = tf.Variable(tf.zeros([10]))
            y = tf.nn.softmax(tf.matmul(x, W) + b)
            y_ = tf.placeholder(tf.float32, [None, 10])

            cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1]))

            train_step = tf.train.GradientDescentOptimizer(0.05).minimize(cross_entropy)

            correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(y_,1))
            accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

            trainers += [train_step]
            accs += [accuracy]
            xs += [x]
            ys += [y_]


# Start an interactive tensorflow session
sess = tf.Session()

# Initialize all variables associated with this session
sess.run(tf.initialize_all_variables())



gpu_num = 15

with tf.device(gpus[gpu_num]):

    for _ in range(20000):
        batch_xs, batch_ys = mnist.train.next_batch(100)
        sess.run(trainers[gpu_num], feed_dict={xs[gpu_num]: batch_xs, ys[gpu_num]: batch_ys})
        break



gpu_num = 15

print "Accuracy on gpu 15"
print(sess.run(accs[gpu_num], feed_dict={xs[gpu_num]: mnist.test.images, ys[gpu_num]: mnist.test.labels}))
print



gpu_num = 14

print "Accuracy on gpu 14"
print(sess.run(accs[gpu_num], feed_dict={xs[gpu_num]: mnist.test.images, ys[gpu_num]: mnist.test.labels}))