Skip to content

Instantly share code, notes, and snippets.

@viswanathgs
Created July 7, 2016 19:26
Show Gist options
  • Save viswanathgs/46015e51006049e783e33245eb269e06 to your computer and use it in GitHub Desktop.
Save viswanathgs/46015e51006049e783e33245eb269e06 to your computer and use it in GitHub Desktop.
Compare GPU to GPU broadcast with CPU to GPU broadcast
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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment