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)