def execute_tpu(self, graph_fn, inputs):
"""Constructs the graph, executes it on TPU and returns the result.
Args:
graph_fn: a callable that constructs the tensorflow graph to test. The
arguments of this function should correspond to `inputs`.
inputs: a list of numpy arrays to feed input to the computation graph.
Returns:
A list of numpy arrays or a scalar returned from executing the tensorflow
graph.
"""
with self.test_session(graph=tf.Graph()) as sess:
placeholders = [
tf.placeholder_with_default(v, v.shape) for v in inputs
]
tpu_computation = tpu.rewrite(graph_fn, placeholders)
sess.run(tpu.initialize_system())
sess.run([
tf.global_variables_initializer(),
tf.tables_initializer(),
tf.local_variables_initializer()
])
materialized_results = sess.run(tpu_computation,
feed_dict=dict(
zip(placeholders, inputs)))
sess.run(tpu.shutdown_system())
if (hasattr(materialized_results, '__len__')
and len(materialized_results) == 1
and (isinstance(materialized_results, list)
or isinstance(materialized_results, tuple))):
materialized_results = materialized_results[0]
return materialized_results
def q1():
global l_returnflag_group_size
global l_linestatus_group_size
returnflag_groups = np.unique(l_returnflag)
linestatus_groups = np.unique(l_linestatus)
l_returnflag_group_size = len(returnflag_groups)
l_linestatus_group_size = len(linestatus_groups)
inputs = [
tf.convert_to_tensor(l_shipdate, np.float32),
tf.convert_to_tensor(l_returnflag, np.float32),
tf.convert_to_tensor(l_linestatus, np.float32),
tf.convert_to_tensor(l_quantity, np.float32),
tf.convert_to_tensor(l_extendedprice, np.float32),
tf.convert_to_tensor(l_discount, np.float32),
tf.convert_to_tensor(l_tax, np.float32),
tf.convert_to_tensor(returnflag_groups, np.float32),
tf.convert_to_tensor(linestatus_groups, np.float32)
]
tpu_computation = tpu.rewrite(q1_computation, inputs)
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
with tf.Session(tpu_grpc_url) as sess:
sess.run(tpu.initialize_system())
sess.run(tf.global_variables_initializer())
for i in range(0, 5):
res = sess.run(tpu_computation)
sess.run(tpu.shutdown_system())
print(res)
return res
def execute_tpu(self, graph_fn, inputs):
"""Constructs the graph, executes it on TPU and returns the result.
Args:
graph_fn: a callable that constructs the tensorflow graph to test. The
arguments of this function should correspond to `inputs`.
inputs: a list of numpy arrays to feed input to the computation graph.
Returns:
A list of numpy arrays or a scalar returned from executing the tensorflow
graph.
"""
with self.test_session(graph=tf.Graph()) as sess:
placeholders = [tf.placeholder_with_default(v, v.shape) for v in inputs]
tpu_computation = tpu.rewrite(graph_fn, placeholders)
sess.run(tpu.initialize_system())
sess.run([tf.global_variables_initializer(), tf.tables_initializer(),
tf.local_variables_initializer()])
materialized_results = sess.run(tpu_computation,
feed_dict=dict(zip(placeholders, inputs)))
sess.run(tpu.shutdown_system())
if (hasattr(materialized_results, '__len__') and
len(materialized_results) == 1 and
(isinstance(materialized_results, list) or
isinstance(materialized_results, tuple))):
materialized_results = materialized_results[0]
return materialized_results
def run(size):
a_ = []
b_ = []
c_ = []
for i in range(size):
a_.append((i * 1.0 + 4.0) * 2.5)
b_.append((i * 1.0 + 5.0) * 2.5)
c_.append((i * 1.0 + 6.0) * 0.1)
inputs = [tf.constant(a_), tf.constant(b_), tf.constant(c_)]
tpu_computation = tpu.rewrite(expression, inputs)
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
with tf.Session(tpu_grpc_url) as sess:
sess.run(tpu.initialize_system())
t1 = time()
sess.run(tf.global_variables_initializer())
sess.run(tpu_computation)
t2 = time()
print(str(size) + " : " + str(t2 - t1))
sess.run(tpu.shutdown_system())
print('Done !')
def filter_sum():
inputs = [tf.convert_to_tensor(l_quantity, np.float32)]
tpu_computation = tpu.rewrite(filter_sum_computation, inputs)
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
with tf.Session(tpu_grpc_url) as sess:
sess.run(tpu.initialize_system())
sess.run(tf.global_variables_initializer())
for i in range(0, 5):
res = sess.run(tpu_computation)
sess.run(tpu.shutdown_system())
print(res)
return res
def apply_comp(inputs):
tpu_computation = tpu.rewrite(apply, inputs)
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
with tf.Session(tpu_grpc_url) as sess:
sess.run(tpu.initialize_system())
sess.run(tf.global_variables_initializer())
t1 = time()
sess.run(tpu_computation)
t2 = time()
sess.run(tpu.shutdown_system())
print(t2 - t1)
def _evaluate(self, variables, observation, **extra_feed):
sess = self.sess
feed_dict = {self.X: adjust_shape(self.X, observation)}
for inpt_name, data in extra_feed.items():
if inpt_name in self.__dict__.keys():
inpt = self.__dict__[inpt_name]
if isinstance(inpt,
tf.Tensor) and inpt._op.type == 'Placeholder':
feed_dict[inpt] = adjust_shape(inpt, data)
if USING_TPUS:
return sess.run(tpu.rewrite(variables, feed_dict))
else:
return sess.run(variables, feed_dict)
def group_by():
unique_groups = np.unique(l_returnflag)
inputs = [
tf.convert_to_tensor(l_quantity, np.float32),
tf.convert_to_tensor(l_returnflag, np.float32),
tf.convert_to_tensor(unique_groups, np.float32)
]
tpu_computation = tpu.rewrite(group_by_computation, inputs)
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
with tf.Session(tpu_grpc_url) as sess:
sess.run(tpu.initialize_system())
sess.run(tf.global_variables_initializer())
for i in range(0, 5):
res = sess.run(tpu_computation)
sess.run(tpu.shutdown_system())
print(res)
def _run_tpu_computation(self):
"""Attempt to run computation graph directly on TPU."""
def _computation_fn(alpha, x, y):
return alpha * x + y
alpha = tf.Variable(3.0, name='alpha')
x = tf.Variable(tf.ones([3, 3], tf.float32), name='x')
y = tf.Variable(tf.ones([3, 3], tf.float32), name='y')
result = contrib_tpu.rewrite(_computation_fn, [alpha, x, y])
with tf.Session('grpc://{0}:8470'.format(self.tpu_ip)) as sess:
sess.run(contrib_tpu.initialize_system())
sess.run(tf.global_variables_initializer())
logging.info(sess.run(result))
sess.run(tpu.shutdown_system())
logging.info('Output should be a 3x3 matrix with all 4s.')
self.tpu_computation = 'Passed'
logging.info('Successfully ran a computation on the TPU')
def test_large_input(self):
if test_case.FLAGS.tpu_test:
input_size = 1408
min_level = 2
max_level = 6
batch_size = 2
num_boxes = 512
num_filters = 256
output_size = [7, 7]
with self.test_session() as sess:
features = []
for level in range(min_level, max_level + 1):
feat_size = int(input_size / 2**level)
features.append(
tf.constant(np.reshape(
np.arange(batch_size * feat_size * feat_size *
num_filters,
dtype=np.float32),
[batch_size, feat_size, feat_size, num_filters]),
dtype=tf.bfloat16))
boxes = np.array([
[[0, 0, 256, 256]] * num_boxes,
],
dtype=np.float32) / input_size
boxes = np.tile(boxes, [batch_size, 1, 1])
tf_boxes = tf.constant(boxes)
tf_levels = tf.random_uniform([batch_size, num_boxes],
maxval=5,
dtype=tf.int32)
def crop_and_resize_fn():
return spatial_ops.multilevel_roi_align(
features, tf_boxes, tf_levels, output_size)
tpu_crop_and_resize_fn = contrib_tpu.rewrite(
crop_and_resize_fn)
sess.run(contrib_tpu.initialize_system())
sess.run(tf.global_variables_initializer())
roi_features = sess.run(tpu_crop_and_resize_fn)
self.assertEqual(roi_features[0].shape,
(batch_size, num_boxes, output_size[0],
output_size[1], num_filters))
sess.run(contrib_tpu.shutdown_system())
def timer(inputs):
reps = 2
times = []
for i in range(reps):
t1 = time()
tpu_computation = tpu.rewrite(blackscholes, inputs)
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
with tf.Session(tpu_grpc_url) as sess:
sess.run(tpu.initialize_system())
sess.run(tf.global_variables_initializer())
sess.run(tpu_computation)
sess.run(tpu.shutdown_system())
t2 = time()
print(str(i) + "_ : " + str(t2 - t1))
times.append(t2 - t1)
print(sum(times) / reps)
def run():
S0 = 100.
K0 = 100.
r0 = 0.05
T0 = 1.0
v0 = 0.2
S = tf.constant(S0)
K = tf.constant(K0)
r = tf.constant(r0)
T = tf.constant(T0)
v = tf.constant(v0)
inputs = [S, K, r, T, v]
tpu_computation = tpu.rewrite(blackscholes, inputs)
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
timer(tpu_computation, tpu_grpc_url)
print('Done !')
def execute_tpu_tf1(self, compute_fn, inputs, graph=None):
"""Executes compute_fn on TPU with Tensorflow 1.X.
Args:
compute_fn: a function containing Tensorflow computation that takes a list
of input numpy tensors, performs computation and returns output numpy
tensors.
inputs: a list of numpy arrays to feed input to the `compute_fn`.
graph: (optional) If not None, provided `graph` is used for computation
instead of a brand new tf.Graph().
Returns:
A list of numpy arrays or a single numpy array.
"""
with self.session(graph=(graph or tf.Graph())) as sess:
placeholders = [
tf.placeholder_with_default(v, v.shape) for v in inputs
]
def wrap_graph_fn(*args, **kwargs):
results = compute_fn(*args, **kwargs)
if (not (isinstance(results, dict)
or isinstance(results, tf.Tensor))
and hasattr(results, '__iter__')):
results = list(results)
return results
tpu_computation = contrib_tpu.rewrite(wrap_graph_fn, placeholders)
sess.run(contrib_tpu.initialize_system())
sess.run([
tf.global_variables_initializer(),
tf.tables_initializer(),
tf.local_variables_initializer()
])
materialized_results = sess.run(tpu_computation,
feed_dict=dict(
zip(placeholders, inputs)))
sess.run(contrib_tpu.shutdown_system())
return self.maybe_extract_single_output(materialized_results)
def train(obs, actions, rewards, dones, mus, states, masks, steps):
cur_lr = lr.value_steps(steps)
td_map = {
train_model.X: obs,
polyak_model.X: obs,
A: actions,
R: rewards,
D: dones,
MU: mus,
LR: cur_lr
}
if states is not None:
td_map[train_model.S] = states
td_map[train_model.M] = masks
td_map[polyak_model.S] = states
td_map[polyak_model.M] = masks
if USING_TPUS:
return names_ops, sess.run(tpu.rewrite(
run_ops, td_map))[1:] # strip off _train
else:
return names_ops, sess.run(run_ops,
td_map)[1:] # strip off _train
for t in range(config.steps):
roll = pool.map(generate_rollout, [t] * config.processes)
gradients = [np.mean(tup, axis=0) for tup in zip(*roll)]
gnorms = model.update(gradients)
if t % 10 == 0:
logger.log_gradients(gnorms)
logger.log_weights()
saver.save(sess,
save_path=config.file_prefix,
global_step=step,
write_meta_graph=False)
print(f"Epoch took: {time.time() - dur:.2f}")
dur = time.time()
if __name__ == '__main__':
if config.use_tpu:
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
tpu_computation = tpu.rewrite(train, tpu_grpc_url)
with tf.Session(tpu_grpc_url) as sess:
sess.run(tpu.initialize_system())
sess.run(tpu_computation)
sess.run(tpu.shutdown_system())
else:
train()
def axy_computation(a, x, y):
return a * x + y
output_shape = [80, 80]
inputs = [
3.0,
tf.random_uniform(output_shape, dtype=tf.float32),
tf.random_uniform(output_shape, dtype=tf.float32),
]
if use_tpu:
print('Setting up TPU')
tpu_grpc_url = TPUClusterResolver(tpu=[tpu_name]).get_master()
computation = tpu.rewrite(axy_computation, inputs)
else:
print(
'TPU IS NOT ENABLED (pass a TPU name or grpc://ip:port as the TPU_NAME envvar)'
)
computation = tf.py_func(axy_computation, inputs, tf.float32)
tpu_grpc_url = None
with tf.Session(tpu_grpc_url) as sess:
if use_tpu:
print('Running TPU initializer')
sess.run(tpu.initialize_system())
sess.run(tf.global_variables_initializer())
print('Running computation {}'.format(computation))
output = sess.run(computation)
print(output)
def handler(queue, kvm_fd, mm):
global object_dict
global object_id
global callback_stack
global initialized
if not initialized:
callback_stack = []
object_dict = dict()
object_id = 1
# TODO: forward logging or disable it in test
tf.logging.set_verbosity(tf.logging.INFO)
initialized = True
print("handler is initialized")
while True:
task = None
task = queue.get(block=True)
while task is None:
try:
task = queue.get(block=True, timeout=5)
except Queue.Empty:
task = None
if callback_stack:
if time.time() > callback_stack[-1]["deadline"]:
print("callback failed deadline")
return STATUS_CALLBACK_TIMEOUT
vm_id = task.vm_id
if vm_id == STOP_HANDLER:
break
param = TF_PY_PARAM.from_buffer(mm, task.data_ptr)
callback_param = TF_PY_PARAM.from_buffer(
mm, task.data_ptr + param.base.callback_param_offset)
print(
"retrieve [vm#%d] tensorflow task=%d cmd=%d, obj=%d, dstore=%lx, done=%d"
% (task.vm_id, task.node_id, param.base.cmd_id,
param.base.object_id, param.base.dstore_size, param.base.done))
print(
"retrieve [vm#%d] callback node cmd=%d, obj=%d, dstore=%lx, done=%d"
% (task.vm_id, callback_param.base.cmd_id,
callback_param.base.object_id, callback_param.base.dstore_size,
callback_param.base.done))
cmd_id = param.base.cmd_id
try:
if cmd_id == TF_PY_NW_CALLBACK_DONE:
param.base.done = STATUS_TASK_DONE
ret = fcntl.ioctl(kvm_fd, IOCTL_KVM_NOTIFY_TASK_FINISHED,
task.node_id)
if ret < 0:
print("notify task completion failed: %d\n" % ret)
if callback_stack and \
callback_stack[-1]["callback_id"] == param.base.object_id:
print("callback is finished")
return STATUS_CALLBACK_DONE
else:
print("callback is error")
return STATUS_CALLBACK_ERROR
if cmd_id == TF_PY_SESSION_INIT:
print("SessionInit!!!")
param1 = parse_param(vm_id, mm, param, param.param1)
print(param1)
sess = tf.Session(param1)
# assign object_id
object_dict[object_id] = sess
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_SESSION_ENTER:
sess = object_dict[param.base.object_id]
ctx_sess = sess.__enter__()
if sess is ctx_sess:
pass
else: # unlikely
print("unlikely to search for sess")
param.base.object_id = next(
obj_id for obj_id, obj in object_dict.items()
if obj is ctx_sess)
elif cmd_id == TF_PY_SESSION_EXIT:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
sess = object_dict[param.base.object_id]
sess.__exit__(param1, param2, param3)
elif cmd_id == TF_PY_SESSION_DEL:
sess = object_dict[param.base.object_id]
sess.__del__()
# deprecated
elif cmd_id == TF_PY_SESSION_RUN:
sess = object_dict[param.base.object_id]
param1 = parse_param(vm_id, mm, param, param.param1)
if type(param1) == NwObject:
print("get NwObject=%d" % param1.object_id())
param1 = object_dict[param1.object_id()]
print(param1)
ret_val = sess.run(param1)
print(ret_val)
writeback_result(vm_id, mm, param, param.ret_val1, ret_val)
elif cmd_id == TF_PY_TPU_CLUSTER_RESOLVER_INIT:
print("resloverInit!!!")
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
if param1 is None:
param1 = None
if param2 is None:
param2 = None
if param3 is None:
param3 = None
print("TPUClusterResolver", param1, param2, param3)
tpu_grpc = tf.contrib.cluster_resolver.TPUClusterResolver(
tpu=param1, zone=param2, project=param3)
# assign object_id
object_dict[object_id] = tpu_grpc
param.base.object_id = object_id
print("assign obj_id=%d" % object_id)
object_id += 1
# deprecated
elif cmd_id == TF_PY_TPU_CLUSTER_RESOLVER_MASTER:
# FIXED: use __getattr__
print("master!!")
tpu_grpc = object_dict[param.base.object_id]
# FIXED: may have parameters
tpu_grpc_url = tpu_grpc.master()
# serialize return value
writeback_result(vm_id, mm, param, param.ret_val1,
tpu_grpc_url)
elif cmd_id == TF_PY_TPU_INITIALIZE_SYSTEM:
# TODO: may have parameters
ts = tpu.initialize_system()
object_dict[object_id] = ts
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_TPU_SHUTDOWN_SYSTEM:
# TODO: may have parameters
ts = tpu.shutdown_system()
object_dict[object_id] = ts
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_GLOBAL_VARIABLES_INITIALIZER:
# TODO: may have parameters
ts = tf.global_variables_initializer()
object_dict[object_id] = ts
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_ONES:
print("param1 size=%ld,offset=%ld" %
(param.param1.size, param.param1.offset))
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
if param2 is None:
param2 = dtypes.float32
print(param2)
var = tf.ones(param1, param2)
object_dict[object_id] = var
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_RANDOM_UNIFORM:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param5 = parse_param(vm_id, mm, param, param.param5)
param6 = parse_param(vm_id, mm, param, param.param6)
if param2 is None:
param2 = 0
if param4 is None:
param4 = dtypes.float32
print(param1, param2, param3, param4)
var = tf.random_uniform(param1, param2, param3, param4, param5,
param6)
object_dict[object_id] = var
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_TRANSPOSE:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param1 = object_dict[param1.object_id()]
if param3 is None:
param3 = "transpose"
if param4 is None:
param4 = False
print("transpose", param1, param2, param3, param4)
var = tf.transpose(param1, param2, param3, param4)
object_dict[object_id] = var
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_CAST:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param1 = object_dict[param1.object_id()]
print("cast", param1, param2, param3)
var = tf.cast(param1, param2, param3)
object_dict[object_id] = var
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_EXPAND_DIMS:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param1 = object_dict[param1.object_id()]
print("expand_dims", param1, param2, param3, param4)
var = tf.expand_dims(param1, param2, param3, param4)
object_dict[object_id] = var
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_CONCAT:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param1 = object_dict[param1.object_id()]
if param3 is None:
param3 = "concat"
print("concat", param1, param2, param3)
var = tf.concat(param1, param2, param3)
object_dict[object_id] = var
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_EQUAL:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param1 = object_dict[param1.object_id()]
print("equal", param1, param2, param3)
if isinstance(param2, NwObject):
param2 = object_dict[param2.object_id()]
result = tf.equal(param1, param2, param3)
print(result)
object_dict[object_id] = result
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_FIXED_LEN_FEATURE:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
feature = tf.FixedLenFeature(param1, param2, param3)
print(feature)
object_dict[object_id] = feature
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_VAR_LEN_FEATURE:
param1 = parse_param(vm_id, mm, param, param.param1)
feature = tf.VarLenFeature(param1)
print(feature)
object_dict[object_id] = feature
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_PARSE_SINGLE_EXAMPLE:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
print(param1, param2)
# expand embedded NwObject
if isinstance(param1, NwObject):
param1 = object_dict[param1.object_id()]
dict_walker(param2)
print("after translation", param1, param2)
result = tf.parse_single_example(param1, param2, param3,
param4)
print(result)
dict_mapper(result)
print(result)
writeback_result(vm_id, mm, param, param.ret_val1, result)
elif cmd_id == TF_PY_CONTROL_FLOW_OPS_SWITCH:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param1 = object_dict[param1.object_id()]
param2 = object_dict[param2.object_id()]
print("switch", param1, param2, param3, param4)
result = control_flow_ops.switch(param1, param2, param3,
param4)
print(result)
mapped_tuple = tuple_mapper(result, [0, 1])
print(mapped_tuple)
writeback_result(vm_id, mm, param, param.ret_val1,
mapped_tuple)
elif cmd_id == TF_PY_CONTROL_FLOW_OPS_MERGE:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param1 = object_dict[param1.object_id()]
print("merge", param1, param2)
list_walker(param1)
print("merge-new", param1, param2)
result = control_flow_ops.merge(param1, param2)
print(result)
mapped_tuple = tuple_mapper(result, [0])
print(mapped_tuple)
writeback_result(vm_id, mm, param, param.ret_val1,
mapped_tuple)
elif cmd_id == TF_PY_TPU_REWRITE:
# TODO: may have parameters
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
# default parameter
if param2 is None:
param2 = None
# expand embedded NwObject
list_walker(param2)
func = tpu.rewrite(param1, param2)
object_dict[object_id] = func
param.base.object_id = object_id
print("rewrite object_id=%d" % object_id)
object_id += 1
elif cmd_id == TF_PY_TPU_RUN_CONFIG:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param5 = parse_param(vm_id, mm, param, param.param5)
# default parameter
if param1 is None:
param1 = None
if param2 is None:
param2 = None
if param3 is None:
param3 = None
if param4 is None:
param4 = None
# expand embedded NwObject
param4 = object_dict[param4.object_id()]
print(param4, param5)
func = tpu.RunConfig(param1, param2, param3, param4, **param5)
object_dict[object_id] = func
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_TPU_TPU_ESTIMATOR:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param5 = parse_param(vm_id, mm, param, param.param5)
param6 = parse_param(vm_id, mm, param, param.param6)
param7 = parse_param(vm_id, mm, param, param.param7)
param8 = parse_param(vm_id, mm, param, param.param8)
param9 = parse_param(vm_id, mm, param, param.param9)
param10 = parse_param(vm_id, mm, param, param.param10)
param11 = parse_param(vm_id, mm, param, param.param11)
param12 = parse_param(vm_id, mm, param, param.param12)
# default parameter
if param1 is None:
param1 = None
if param2 is None:
param2 = None
if param3 is None:
param3 = None
if param4 is None:
param4 = None
if param5 is None:
param5 = True
if param6 is None:
param6 = None
if param7 is None:
param7 = None
if param8 is None:
param8 = None
if param9 is None:
param9 = None
if param10 is None:
param10 = True
if param11 is None:
param11 = True
if param12 is None:
param12 = None
# expand embedded NwObject
param3 = object_dict[param3.object_id()]
print(param3)
func = tpu.TPUEstimator(param1, param2, param3, param4, param5,
param6, param7, param8, param9,
param10, param11, param12)
object_dict[object_id] = func
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_IMAGE_RESIZE_IMAGES:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param5 = parse_param(vm_id, mm, param, param.param5)
# default parameter
if param3 is None:
param3 = ResizeMethod.BILINEAR
if param4 is None:
param4 = False
if param5 is None:
param5 = False
# expand embedded NwObject
param1 = object_dict[param1.object_id()]
print(param1)
img = tf.image.resize_images(param1, param2, param3, param4,
param5)
# TODO: it may return a float
object_dict[object_id] = img
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_SLICE:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
# expand embedded NwObject
print(param1, param2, param3)
param1 = object_dict[param1.object_id()]
ret = tf.slice(param1, param2, param3, param4)
object_dict[object_id] = ret
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_SHAPE:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
if param3 is None:
param3 = dtypes.int32
# expand embedded NwObject
print(param1, param2, param3)
param1 = object_dict[param1.object_id()]
ret = tf.shape(param1, param2, param3)
object_dict[object_id] = ret
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_IMAGE_SAMPLE_DISTORTED_BOUNDING_BOX:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param5 = parse_param(vm_id, mm, param, param.param5)
param6 = parse_param(vm_id, mm, param, param.param6)
param7 = parse_param(vm_id, mm, param, param.param7)
param8 = parse_param(vm_id, mm, param, param.param8)
param9 = parse_param(vm_id, mm, param, param.param9)
param10 = parse_param(vm_id, mm, param, param.param10)
# default parameter
if param5 is None:
param5 = 0.1
print("sample_distorted_bounding_box", param1, param2)
result = tf.image.sample_distorted_bounding_box(
param1, param2, param3, param4, param5, param6, param7,
param8, param9, param10)
print(result)
mapped_tuple = tuple_mapper(result, [0, 1, 2])
print(mapped_tuple)
writeback_result(vm_id, mm, param, param.ret_val1,
mapped_tuple)
elif cmd_id == TF_PY_IMAGE_DRAW_BOUNDING_BOXES:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
# expand embedded NwObject
print(param1, param2, param3)
param1 = object_dict[param1.object_id()]
param2 = object_dict[param2.object_id()]
ret = tf.image.draw_bounding_boxes(param1, param2, param3)
object_dict[object_id] = ret
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_IMAGE_DECODE_JPEG:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
param5 = parse_param(vm_id, mm, param, param.param5)
param6 = parse_param(vm_id, mm, param, param.param6)
param7 = parse_param(vm_id, mm, param, param.param7)
param8 = parse_param(vm_id, mm, param, param.param8)
if param2 is None:
param2 = 0
if param3 is None:
param3 = 1
if param4 is None:
param4 = True
if param5 is None:
param5 = False
if param6 is None:
param6 = 1
if param7 is None:
param7 = ""
param1 = object_dict[param1.object_id()]
img = tf.image.decode_jpeg(param1, param2, param3, param4,
param5, param6, param7, param8)
object_dict[object_id] = img
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_IMAGE_CONVERT_IMAGE_DTYPE:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
param4 = parse_param(vm_id, mm, param, param.param4)
# expand embedded NwObject
print(param1, param2, param3)
param1 = object_dict[param1.object_id()]
if param3 is None:
param3 = False
ret = tf.image.convert_image_dtype(param1, param2, param3,
param4)
object_dict[object_id] = ret
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_DATA_DATASET_LIST_FILES:
param1 = parse_param(vm_id, mm, param, param.param1)
param2 = parse_param(vm_id, mm, param, param.param2)
param3 = parse_param(vm_id, mm, param, param.param3)
print(param1, param2, param3)
if isinstance(param1, NwObject):
param1 = object_dict[oaram1.object_id()]
ret = tf.data.Dataset.list_files(param1, param2, param3)
object_dict[object_id] = ret
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_NW_OBJECT:
print("nw_object!! id = %d" % param.base.object_id)
obj = object_dict[param.base.object_id]
name = parse_param(vm_id, mm, param, param.param1)
args = parse_param(vm_id, mm, param, param.param2)
kwargs = parse_param(vm_id, mm, param, param.param3)
print("NwObject", obj, name, args, kwargs)
# expand embedded NwObject
args = list(args)
list_walker(args)
args = tuple(args)
dict_walker(kwargs)
print("after translation", obj, name, args, kwargs)
# run
result = getattr(obj, name)(*(args or []), **(kwargs or {}))
param.base.object_id = -1
param.ret_val1.size = 0
print("analyze type", type(result), result)
# TODO: go through tuple, dict or list
if isinstance(result, tuple):
result = tuple_mapper(result, range(len(result)))
if isinstance(result, dict):
dict_mapper(result)
if isinstance(result, list):
list_mapper(result)
# serialize return value
if is_unpickleable_type(result) or \
pickle.pickles(result) is False:
object_dict[object_id] = result
param.base.object_id = object_id
object_id += 1
elif result is not None:
writeback_result(vm_id, mm, param, param.ret_val1, result)
elif cmd_id == TF_PY_NW_METHOD:
# Reuse as callback
#ins = parse_param(vm_id, mm, param, param.param1)
#name = parse_param(vm_id, mm, param, param.param2)
#print(ins, name)
#method = getattr(ins, name)
#print(method)
#object_dict[object_id] = method
cw = callback_constructor(object_id, callback_param, param, mm,
vm_id, queue, kvm_fd)
object_dict[object_id] = cw
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_NW_CALLBACK_TEST:
nw_func = parse_param(vm_id, mm, param, param.param1)
print(nw_func, nw_func.object_id())
func = object_dict[nw_func.object_id()]
print("callback func", func)
x = parse_param(vm_id, mm, param, param.param2)
y = parse_param(vm_id, mm, param, param.param3)
result = func(x, y)
print(result)
writeback_result(vm_id, mm, param, param.ret_val1, result)
else:
print("unsupported Tensorflow API")
except Exception, error:
param.base.done = STATUS_TASK_ERROR
#mm.flush(task.data_ptr, sizeof(PARAM_BASE))
print "fault: ", str(error)
exc_type, exc_obj, exc_tb = sys.exc_info()
fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
print(exc_type, fname, exc_tb.tb_lineno)
traceback.print_stack()
print("finished [vm#%d] TF task %d cmd %d" %
(task.vm_id, task.node_id, param.base.cmd_id))
param.base.done = STATUS_TASK_DONE
#mm.flush(task.data_ptr, sizeof(PARAM_BASE))
#mm.flush(INVOKER_FIFO_SIZE + VGPU_DSTORE_SIZE * (vm_id - 1) +
# param.base.dstore_offset + param.ret_val1.offset,
# param.ret_val1.size)
# notify hypervisor
ret = fcntl.ioctl(kvm_fd, IOCTL_KVM_NOTIFY_TASK_FINISHED, task.node_id)
if ret < 0:
print("notify task completion failed: %d\n" % ret)
slice3 = tf.slice(fijk, [i + d3, j, 0], [1, 1, size])
slice4 = tf.slice(fijk, [i - d4, j, 0], [1, 1, size])
fdo = c1 * slice1 + c2 * slice2 + c3 * slice3 + c4 * slice4
return fdo
if __name__ == "__main__":
dim1 = [0., 1., 2., 3., 4.]
dim2 = [5., 6., 7., 8., 9.]
dim3 = [10., 11., 12., 13., 14.]
dim4 = [15., 16., 17., 18., 19.]
fijk = tf.constant([[dim1, dim2, dim3, dim4], [dim2, dim3, dim4, dim1],
[dim3, dim4, dim1, dim2], [dim4, dim1, dim2, dim3]])
i = tf.constant(1)
j = tf.constant(1)
dk = tf.constant(0)
inputs = [fijk, i, j, dk]
tpu_computation = tpu.rewrite(apply_, inputs)
tpu_grpc_url = TPUClusterResolver(
tpu=[os.environ['TPU_NAME']]).get_master()
run(tpu_computation, tpu_grpc_url)
print('Done !')
def handler(cmd_queue, chan):
global object_dict
global object_id
global callback_stack
global initialized
if not initialized:
callback_stack = []
object_dict = dict()
object_id = 1
# TODO: forward logging or disable it in test
tf.logging.set_verbosity(tf.logging.INFO)
initialized = True
print("handler is initialized")
while True:
cmd = cmd_queue.get(block=True)
cmd_id = cmd.__get_cmd_id()
print("new command id %d" % cmd_id)
try:
if cmd_id == TF_PY_NW_CALLBACK_DONE:
param.base.done = STATUS_TASK_DONE
ret = fcntl.ioctl(kvm_fd, IOCTL_KVM_NOTIFY_TASK_FINISHED, task.node_id)
if ret < 0:
print("notify task completion failed: %d\n" % ret);
if callback_stack and \
callback_stack[-1]["callback_id"] == param.base.object_id:
print("callback is finished")
return STATUS_CALLBACK_DONE
else:
print("callback is error")
return STATUS_CALLBACK_ERROR
if cmd_id == TF_PY_SESSION_INIT:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
sess = tf.Session(param0, param1, param2)
# assign object_id
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = sess
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
#elif cmd_id == TF_PY_SESSION_ENTER:
# sess = object_dict[param.base.object_id]
# ctx_sess = sess.__enter__()
# if sess is ctx_sess:
# pass
# else: # unlikely
# print("unlikely to search for sess")
# param.base.object_id = next(obj_id for obj_id, obj in
# object_dict.items() if obj is ctx_sess)
#elif cmd_id == TF_PY_SESSION_EXIT:
# param1 = parse_param(vm_id, mm, param, param.param1)
# param2 = parse_param(vm_id, mm, param, param.param2)
# param3 = parse_param(vm_id, mm, param, param.param3)
# sess = object_dict[param.base.object_id]
# sess.__exit__(param1, param2, param3)
#elif cmd_id == TF_PY_SESSION_DEL:
# sess = object_dict[param.base.object_id]
# sess.__del__()
# deprecated
#elif cmd_id == TF_PY_SESSION_RUN:
# sess = object_dict[param.base.object_id]
# param1 = parse_param(vm_id, mm, param, param.param1)
# if type(param1) == NwObject:
# print("get NwObject=%d" % param1.object_id())
# param1 = object_dict[param1.object_id()]
# print(param1)
# ret_val = sess.run(param1)
# print(ret_val)
# writeback_result(vm_id, mm, param, param.ret_val1, ret_val);
elif cmd_id == TF_PY_TPU_CLUSTER_RESOLVER_INIT:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
print("TPUClusterResolver", param0, param1, param2)
tpu_grpc = tf.contrib.cluster_resolver.TPUClusterResolver(
tpu=param0, zone=param1, project=param2)
# assign object_id
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = tpu_grpc
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
# deprecated
elif cmd_id == TF_PY_TPU_CLUSTER_RESOLVER_MASTER:
tpu_grpc = object_dict[cmd.__get_object_id()]
# FIXED: may have parameters
tpu_grpc_url = tpu_grpc.master()
# serialize return value
dump_ret, len_ret = pickle_arg(tpu_grpc_url)
total_buffer_size = chan.buffer_size(len_ret)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), total_buffer_size)
ret_cmd.__set_cmd_base(TENSORFLOW_PY_API, TF_PY_TPU_CLUSTER_RESOLVER_MASTER_RET)
offset_ret = ret_cmd.attach_buffer(dump_ret, len_ret)
ret_cmd.__set_tf_args([(len_ret, offset_ret)])
ret_cmd.send()
elif cmd_id == TF_PY_TPU_INITIALIZE_SYSTEM:
# TODO: may have parameters
ts = tpu.initialize_system()
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ts
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_TPU_SHUTDOWN_SYSTEM:
# TODO: may have parameters
ts = tpu.shutdown_system()
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ts
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_GLOBAL_VARIABLES_INITIALIZER:
# TODO: may have parameters
ts = tf.global_variables_initializer()
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ts
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_ONES:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
print(param0)
if param1 is None:
param1 = dtypes.float32
print(param1)
var = tf.ones(param0, param1)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = var
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_RANDOM_UNIFORM:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param4 = unpickle_arg(cmd, 4)
param5 = unpickle_arg(cmd, 5)
if param1 is None:
param1 = 0
if param3 is None:
param3 = dtypes.float32
print(param0, param1, param2, param3)
var = tf.random_uniform(param0, param1, param2, param3, param4, param5)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = var
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_TRANSPOSE:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param0 = object_dict[param0.object_id()]
if param2 is None:
param2 = "transpose"
if param3 is None:
param3 = False
print("transpose", param0, param1, param2, param3)
var = tf.transpose(param0, param1, param2, param3)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = var
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_CAST:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param0 = object_dict[param0.object_id()]
print("cast", param0, param1, param2)
var = tf.cast(param0, param1, param2)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = var
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_EXPAND_DIMS:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param0 = object_dict[param0.object_id()]
print("expand_dims", param0, param1, param2, param3)
var = tf.expand_dims(param0, param1, param2, param3)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = var
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_CONCAT:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param0 = object_dict[param0.object_id()]
if param2 is None:
param2 = "concat"
print("concat", param0, param1, param2)
var = tf.concat(param0, param1, param2)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = var
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_EQUAL:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param0 = object_dict[param0.object_id()]
print("equal", param0, param1, param2)
if isinstance(param1, NwObject):
param1 = object_dict[param1.object_id()]
result = tf.equal(param0, param1, param2)
print(result)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = result
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_FIXED_LEN_FEATURE:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
feature = tf.FixedLenFeature(param0, param1, param2)
print(feature)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = feature
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_VAR_LEN_FEATURE:
param0 = unpickle_arg(cmd, 0)
feature = tf.VarLenFeature(param0)
print(feature)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = feature
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_PARSE_SINGLE_EXAMPLE:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
print(param1, param2)
# expand embedded NwObject
if isinstance(param0, NwObject):
param0 = object_dict[param0.object_id()]
dict_walker(param1)
print("after translation", param0, param1)
result = tf.parse_single_example(param0, param1, param2, param3)
print(result)
dict_mapper(result)
print(result)
dump_ret, len_ret = pickle_arg(result)
total_buffer_size = chan.buffer_size(len_ret)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), total_buffer_size)
ret_cmd.__set_cmd_base(TENSORFLOW_PY_API, TF_PY_PARSE_SINGLE_EXAMPLE_RET)
offset_ret = ret_cmd.attach_buffer(dump_ret, len_ret)
ret_cmd.__set_tf_args([(len_ret, offset_ret)])
ret_cmd.send()
elif cmd_id == TF_PY_CONTROL_FLOW_OPS_SWITCH:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param0 = object_dict[param0.object_id()]
param1 = object_dict[param1.object_id()]
print("switch", param0, param1, param2, param3)
result = control_flow_ops.switch(param0, param1, param2, param3)
print(result)
mapped_tuple = tuple_mapper(result, [0, 1])
print(mapped_tuple)
dump_ret, len_ret = pickle_arg(mapped_tuple)
total_buffer_size = chan.buffer_size(len_ret)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), total_buffer_size)
ret_cmd.__set_cmd_base(TENSORFLOW_PY_API, TF_PY_CONTROL_FLOW_OPS_SWITCH_RET)
offset_ret = ret_cmd.attach_buffer(dump_ret, len_ret)
ret_cmd.__set_tf_args([(len_ret, offset_ret)])
ret_cmd.send()
elif cmd_id == TF_PY_CONTROL_FLOW_OPS_MERGE:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param0 = object_dict[param0.object_id()]
print("merge", param0, param1)
list_walker(param0)
print("merge-new", param0, param1)
result = control_flow_ops.merge(param0, param1)
print(result)
mapped_tuple = tuple_mapper(result, [0])
print(mapped_tuple)
dump_ret, len_ret = pickle_arg(mapped_tuple)
total_buffer_size = chan.buffer_size(len_ret)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), total_buffer_size)
ret_cmd.__set_cmd_base(TENSORFLOW_PY_API, TF_PY_CONTROL_FLOW_OPS_MERGE_RET)
offset_ret = ret_cmd.attach_buffer(dump_ret, len_ret)
ret_cmd.__set_tf_args([(len_ret, offset_ret)])
ret_cmd.send()
elif cmd_id == TF_PY_TPU_REWRITE:
# TODO: may have more parameters
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
# default parameter
if param1 is None:
param1 = None
# expand embedded NwObject
list_walker(param1)
func = tpu.rewrite(param0, param1)
print("Rewrite result:", func, " object id =", object_id)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = func
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_TPU_RUN_CONFIG:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param4 = unpickle_arg(cmd, 4)
# default parameter
if param0 is None:
param0 = None
if param1 is None:
param1 = None
if param2 is None:
param2 = None
if param3 is None:
param3 = None
# expand embedded NwObject
param3 = object_dict[param3.object_id()]
print(param3, param4)
func = tpu.RunConfig(param0, param1, param2, param3, **param4)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = func
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_TPU_TPU_ESTIMATOR:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param4 = unpickle_arg(cmd, 4)
param5 = unpickle_arg(cmd, 5)
param6 = unpickle_arg(cmd, 6)
param7 = unpickle_arg(cmd, 7)
param8 = unpickle_arg(cmd, 8)
param9 = unpickle_arg(cmd, 9)
param10 = unpickle_arg(cmd, 10)
param11 = unpickle_arg(cmd, 11)
# default parameter
if param0 is None:
param0 = None
if param1 is None:
param1 = None
if param2 is None:
param2 = None
if param3 is None:
param3 = None
if param4 is None:
param4 = True
if param5 is None:
param5 = None
if param6 is None:
param6 = None
if param7 is None:
param7 = None
if param8 is None:
param8 = None
if param9 is None:
param9 = True
if param10 is None:
param10 = True
if param11 is None:
param11 = None
# expand embedded NwObject
param2 = object_dict[param2.object_id()]
print(param2)
func = tpu.TPUEstimator(param0, param1, param2, param3, param4,
param5, param6, param7, param8, param9,
param10, param11)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = func
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_IMAGE_RESIZE_IMAGES:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param4 = unpickle_arg(cmd, 4)
# default parameter
if param2 is None:
param2 =ResizeMethod.BILINEAR
if param3 is None:
param3 = False
if param4 is None:
param4 = False
# expand embedded NwObject
param0 = object_dict[param0.object_id()]
print(param0)
img = tf.image.resize_images(param0, param1, param2, param3, param4)
# TODO: it may return a float
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = img
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_SLICE:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
# expand embedded NwObject
print(param0, param1, param2, param3)
param0 = object_dict[param0.object_id()]
ret = tf.slice(param0, param1, param2, param3)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ret
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_SHAPE:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
if param2 is None:
param2 = dtypes.int32
# expand embedded NwObject
print(param0, param1, param2)
param0 = object_dict[param0.object_id()]
ret = tf.shape(param0, param1, param2)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ret
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_IMAGE_SAMPLE_DISTORTED_BOUNDING_BOX:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param4 = unpickle_arg(cmd, 4)
param5 = unpickle_arg(cmd, 5)
param6 = unpickle_arg(cmd, 6)
param7 = unpickle_arg(cmd, 7)
param8 = unpickle_arg(cmd, 8)
param9 = unpickle_arg(cmd, 9)
# default parameter
if param4 is None:
param4 = 0.1
print("sample_distorted_bounding_box", param0, param1)
result = tf.image.sample_distorted_bounding_box(
param0, param1, param2, param3, param4, param5, param6,
param7, param8, param9)
print(result)
mapped_tuple = tuple_mapper(result, [0, 1, 2])
print(mapped_tuple)
dump_ret, len_ret = pickle_arg(mapped_tuple)
total_buffer_size = chan.buffer_size(len_ret)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), total_buffer_size)
ret_cmd.__set_cmd_base(TENSORFLOW_PY_API, TF_PY_CONTROL_FLOW_OPS_MERGE_RET)
offset_ret = ret_cmd.attach_buffer(dump_ret, len_ret)
ret_cmd.__set_tf_args([(len_ret, offset_ret)])
ret_cmd.send()
elif cmd_id == TF_PY_IMAGE_DRAW_BOUNDING_BOXES:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
# expand embedded NwObject
print(param0, param1, param2)
param0 = object_dict[param0.object_id()]
param1 = object_dict[param1.object_id()]
ret = tf.image.draw_bounding_boxes(param0, param1, param2)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ret
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_IMAGE_DECODE_JPEG:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
param4 = unpickle_arg(cmd, 4)
param5 = unpickle_arg(cmd, 5)
param6 = unpickle_arg(cmd, 6)
param7 = unpickle_arg(cmd, 7)
if param1 is None:
param1 = 0
if param2 is None:
param2 = 1
if param3 is None:
param3 = True
if param4 is None:
param4 = False
if param5 is None:
param5 = 1
if param6 is None:
param6 = ""
param0 = object_dict[param0.object_id()]
ret = tf.image.decode_jpeg(param0, param1, param2, param3,
param4, param5, param6, param7)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ret
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_IMAGE_CONVERT_IMAGE_DTYPE:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
param3 = unpickle_arg(cmd, 3)
# expand embedded NwObject
print(param0, param1, param2, param3)
param0 = object_dict[param0.object_id()]
if param2 is None:
param2 = False
ret = tf.image.convert_image_dtype(param0, param1, param2, param3)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ret
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_DATA_DATASET_LIST_FILES:
param0 = unpickle_arg(cmd, 0)
param1 = unpickle_arg(cmd, 1)
param2 = unpickle_arg(cmd, 2)
print(param0, param1, param2)
if isinstance(param0, NwObject):
param0 = object_dict[param0.object_id()]
ret = tf.data.Dataset.list_files(param0, param1, param2)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = ret
ret_cmd.__set_object_id(object_id)
object_id += 1
ret_cmd.send()
elif cmd_id == TF_PY_NW_OBJECT:
obj = object_dict[cmd.__get_object_id()]
name = unpickle_arg(cmd, 0)
args = unpickle_arg(cmd, 1)
kwargs = unpickle_arg(cmd, 2)
print("NwObject", obj, name, args, kwargs)
# expand embedded NwObject
args = list(args)
list_walker(args)
args = tuple(args)
dict_walker(kwargs)
print("after translation", obj, name, args, kwargs)
# run
result = getattr(obj, name)(*(args or []), **(kwargs or {}))
print("analyze type", type(result), result)
# TODO: go through tuple, dict or list
if isinstance(result, tuple):
result = tuple_mapper(result, range(len(result)))
if isinstance(result, dict):
dict_mapper(result)
if isinstance(result, list):
list_mapper(result)
ret_cmd = None
# serialize return value
# if isinstance(result, list):
# Check whether a nested list pickles
# https://github.com/uqfoundation/dill/issues/307
# pickleable = pickle.pickles(reduce(operator.add, result))
# else:
# pickleable = pickle.pickles(result)
pickleable = True
if is_unpickleable_type(result) or not pickleable:
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), 0)
object_dict[object_id] = result
ret_cmd.__set_object_id(object_id)
object_id += 1
else:
dump_ret, len_ret = pickle_arg(result)
total_buffer_size = chan.buffer_size(len_ret)
ret_cmd = chan.new_command(chan.__get_sizeof_tf_cmd(), total_buffer_size)
ret_cmd.__set_cmd_base(TENSORFLOW_PY_API, TF_PY_NW_OBJECT_RET)
offset_ret = ret_cmd.attach_buffer(dump_ret, len_ret)
ret_cmd.__set_object_id(0)
ret_cmd.__set_tf_args([(len_ret, offset_ret)])
ret_cmd.send()
elif cmd_id == TF_PY_NW_METHOD:
# Reuse as callback
#ins = parse_param(vm_id, mm, param, param.param1)
#name = parse_param(vm_id, mm, param, param.param2)
#print(ins, name)
#method = getattr(ins, name)
#print(method)
#object_dict[object_id] = method
cw = callback_constructor(object_id, callback_param,
param, mm, vm_id, cmd_queue, kvm_fd)
object_dict[object_id] = cw
param.base.object_id = object_id
object_id += 1
elif cmd_id == TF_PY_NW_CALLBACK_TEST:
nw_func = parse_param(vm_id, mm, param, param.param1)
print(nw_func, nw_func.object_id())
func = object_dict[nw_func.object_id()]
print("callback func", func)
x = parse_param(vm_id, mm, param, param.param2)
y = parse_param(vm_id, mm, param, param.param3)
result = func(x, y)
print(result)
writeback_result(vm_id, mm, param, param.ret_val1, result);
else:
print("unsupported Tensorflow API %d" % cmd_id)
except Exception as error:
print("fault: ", str(error))
exc_type, exc_obj, exc_tb = sys.exc_info()
fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
print(exc_type, fname, exc_tb.tb_lineno)
traceback.print_stack()
cmd.free_command()
print("finished cmd %d" % cmd_id)
