Created
March 6, 2019 21:39
-
-
Save xilenteyex/47c0810ec447d87f26cf3ce419f6efb8 to your computer and use it in GitHub Desktop.
toy matmul code with placeholders moved out of the control dependency context
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 tensorflow as tf | |
import json | |
import sys | |
from tensorflow.python.client import timeline | |
from protobuf_to_dict import protobuf_to_dict | |
from google.protobuf.json_format import MessageToJson | |
from tensorflow.core.protobuf import rewriter_config_pb2 | |
from tensorflow.python.framework import graph_io | |
import os | |
dim = 16384 | |
with tf.device('/gpu:0'): | |
X9 = tf.random_uniform([dim, dim], 0, 10) | |
_X9 = tf.placeholder(dtype=tf.float32, shape=[dim, dim]) | |
Z19 = tf.matmul(_X9, _X9) | |
with tf.device('/gpu:1'): | |
Y9 = tf.random_uniform([dim, dim], 0, 10) | |
_Y9 = tf.placeholder(dtype=tf.float32, shape=[dim, dim]) | |
Z29 = tf.matmul(_Y9, _Y9) | |
W9 = tf.matmul(Z19, Z29) | |
n = 8 | |
dim = 32 | |
Z1, Z2, W = [], [], [] | |
X, _X, Y, _Y = [], [], [], [] | |
with tf.device('/gpu:0'): | |
for i in range(n): | |
dim *= 2 | |
X.append(tf.random_uniform([dim, dim], 0, 10, name='X' + str(i))) | |
Y.append(tf.random_uniform([dim, dim], 0, 10, name='Y' + str(i))) | |
_X.append(tf.placeholder(dtype=tf.float32, shape=[dim, dim])) | |
_Y.append(tf.placeholder(dtype=tf.float32, shape=[dim, dim])) | |
with tf.control_dependencies([Z19]): | |
Z1.append(tf.matmul(_X[i], _X[i])) | |
Z2.append(tf.matmul(_Y[i], _Y[i])) | |
W.append(tf.matmul(Z1[i], Z2[i])) | |
config_proto = tf.ConfigProto(graph_options=tf.GraphOptions(build_cost_model=1)) | |
config_proto.intra_op_parallelism_threads = 1 | |
config_proto.inter_op_parallelism_threads = 1 | |
config_proto.graph_options.optimizer_options.opt_level = -1 | |
config_proto.graph_options.rewrite_options.constant_folding = (rewriter_config_pb2.RewriterConfig.OFF) | |
config_proto.graph_options.rewrite_options.arithmetic_optimization = (rewriter_config_pb2.RewriterConfig.OFF) | |
sess = tf.Session(config=config_proto) | |
sess.run(tf.global_variables_initializer()) | |
X_, Y_, X9_, Y9_ = sess.run([X, Y, X9, Y9]) | |
X_Y_ = X_ + [X9_] + Y_ + [Y9_] | |
_X_Y = _X + [_X9] + _Y + [_Y9] | |
run_metadata = tf.RunMetadata() | |
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, output_partition_graphs=True) | |
W_ = sess.run(W + [W9], | |
{_i: i_ for _i, i_ in zip(_X_Y, X_Y_)}, | |
options=run_options, | |
run_metadata=run_metadata) | |
jsonObj = MessageToJson(run_metadata) | |
with open('metadata_matmul.json', 'w') as outfile: | |
json.dump(jsonObj, outfile) | |
trace = timeline.Timeline(step_stats=run_metadata.step_stats) | |
trace_file = open('timeline_matmul.ctf.json', 'w') | |
trace_file.write(trace.generate_chrome_trace_format()) | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment