Compare GPU to GPU broadcast with CPU to GPU broadcast

gpu_io_benchmark.py

import argparse
import tensorflow as tf
import time

def run(source_device, num_gpus=4):
    shape = [100000000]

    source_device = '/%s:0' % source_device
    with tf.device(source_device):
      weight_init = tf.truncated_normal_initializer()
      a = tf.get_variable('source_var', shape=shape, initializer=weight_init)

    broadcasts = []
    for i in range(1, num_gpus):
      with tf.device('/gpu:%d' % i):
        b = tf.get_variable('copy_%d' % i, shape=shape)
        broadcasts.append(tf.assign(b, a))

    session = tf.Session(config=tf.ConfigProto(log_device_placement=True))

    start = time.time() * 1000
    session.run(tf.initialize_all_variables())
    return (time.time() * 1000 - start)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('source_device', choices=['cpu', 'gpu'])
    args = parser.parse_args()

    elapsed = run(args.source_device)
    print("Elapsed time = {} ms".format(elapsed))