def simple_if(condition_val): condition = ng.constant(condition_val, dtype=np.bool) # then_body X_t = ng.parameter([2], np.float32, "X") Y_t = ng.parameter([2], np.float32, "Y") then_mul = ng.multiply(X_t, Y_t) then_body_res_1 = ng.result(then_mul) then_body = GraphBody([X_t, Y_t], [then_body_res_1]) then_body_inputs = [ TensorIteratorInvariantInputDesc(1, 0), TensorIteratorInvariantInputDesc(2, 1) ] then_body_outputs = [TensorIteratorBodyOutputDesc(0, 0)] # else_body X_e = ng.parameter([2], np.float32, "X") Y_e = ng.parameter([2], np.float32, "Y") add_e = ng.add(X_e, Y_e) else_body_res_1 = ng.result(add_e) else_body = GraphBody([X_e, Y_e], [else_body_res_1]) else_body_inputs = [ TensorIteratorInvariantInputDesc(1, 0), TensorIteratorInvariantInputDesc(2, 1) ] else_body_outputs = [TensorIteratorBodyOutputDesc(0, 0)] X = ng.constant([3, 4], dtype=np.float32) Y = ng.constant([2, 1], dtype=np.float32) if_node = ng.if_op(condition, [X, Y], (then_body, else_body), (then_body_inputs, else_body_inputs), (then_body_outputs, else_body_outputs)) relu = ng.relu(if_node) return relu
def create_simple_if_with_two_outputs(condition_val): condition = ng.constant(condition_val, dtype=np.bool) # then_body X_t = ng.parameter([], np.float32, "X") Y_t = ng.parameter([], np.float32, "Y") Z_t = ng.parameter([], np.float32, "Z") add_t = ng.add(X_t, Y_t) mul_t = ng.multiply(Y_t, Z_t) then_body_res_1 = ng.result(add_t) then_body_res_2 = ng.result(mul_t) then_body = GraphBody([X_t, Y_t, Z_t], [then_body_res_1, then_body_res_2]) then_body_inputs = [ TensorIteratorInvariantInputDesc(1, 0), TensorIteratorInvariantInputDesc(2, 1), TensorIteratorInvariantInputDesc(3, 2) ] then_body_outputs = [ TensorIteratorBodyOutputDesc(0, 0), TensorIteratorBodyOutputDesc(1, 1) ] # else_body X_e = ng.parameter([], np.float32, "X") Z_e = ng.parameter([], np.float32, "Z") W_e = ng.parameter([], np.float32, "W") add_e = ng.add(X_e, W_e) pow_e = ng.power(W_e, Z_e) else_body_res_1 = ng.result(add_e) else_body_res_2 = ng.result(pow_e) else_body = GraphBody([X_e, Z_e, W_e], [else_body_res_1, else_body_res_2]) else_body_inputs = [ TensorIteratorInvariantInputDesc(1, 0), TensorIteratorInvariantInputDesc(3, 1), TensorIteratorInvariantInputDesc(4, 2) ] else_body_outputs = [ TensorIteratorBodyOutputDesc(0, 0), TensorIteratorBodyOutputDesc(1, 1) ] X = ng.constant(15.0, dtype=np.float32) Y = ng.constant(-5.0, dtype=np.float32) Z = ng.constant(4.0, dtype=np.float32) W = ng.constant(2.0, dtype=np.float32) if_node = ng.if_op(condition, [X, Y, Z, W], (then_body, else_body), (then_body_inputs, else_body_inputs), (then_body_outputs, else_body_outputs)) return if_node
def loop( trip_count: NodeInput, execution_condition: NodeInput, inputs: List[Node], graph_body: GraphBody, slice_input_desc: List[TensorIteratorSliceInputDesc], merged_input_desc: List[TensorIteratorMergedInputDesc], invariant_input_desc: List[TensorIteratorInvariantInputDesc], body_output_desc: List[TensorIteratorBodyOutputDesc], concat_output_desc: List[TensorIteratorConcatOutputDesc], body_condition_output_idx: int, current_iteration_input_idx: int = -1, name: Optional[str] = None, ) -> Node: """Perform recurrent execution of the network described in the body, iterating through the data. :param trip_count: A scalar or 1D tensor with 1 element specifying maximum number of iterations. :param execution_condition: A scalar or 1D tensor with 1 element specifying whether to execute the first iteration or not. :param inputs: The provided to TensorIterator operator. :param graph_body: The graph representing the body we execute. :param slice_input_desc: The descriptors describing sliced inputs, that is nodes representing tensors we iterate through, processing single data slice in one iteration. :param merged_input_desc: The descriptors describing merged inputs, that is nodes representing variables with initial value at first iteration, which may be changing through iterations. :param invariant_input_desc: The descriptors describing invariant inputs, that is nodes representing variable with persistent value through all iterations. :param body_output_desc: The descriptors describing body outputs from specified iteration. :param concat_output_desc: The descriptors describing specified output values through all the iterations concatenated into one node. :param body_condition_output_idx: Determines the purpose of the corresponding result in the graph_body. This result will determine the dynamic exit condition. If the value of this result is False, then iterations stop. :param current_iteration_input_idx: Determines the purpose of the corresponding parameter in the graph_body. This parameter will be used as an iteration counter. Optional. :return: The new node which performs Loop. """ attributes = { "body": graph_body.serialize(), "input_descriptions": {"slice_input_desc": [desc.serialize() for desc in slice_input_desc], "merged_input_desc": [desc.serialize() for desc in merged_input_desc], "invariant_input_desc": [desc.serialize() for desc in invariant_input_desc]}, "output_descriptions": {"body_output_desc": [desc.serialize() for desc in body_output_desc], "concat_output_desc": [desc.serialize() for desc in concat_output_desc]}, "special_body_ports": {"body_condition_output_idx": body_condition_output_idx, "current_iteration_input_idx": current_iteration_input_idx} } return _get_node_factory_opset5().create("Loop", as_nodes(trip_count, execution_condition, *inputs), attributes)
def create_diff_if_with_two_outputs(condition_val): condition = ng.constant(condition_val, dtype=np.bool) # then_body X_t = ng.parameter([2], np.float32, "X") Y_t = ng.parameter([2], np.float32, "Y") mmul_t = ng.matmul(X_t, Y_t, False, False) mul_t = ng.multiply(Y_t, X_t) then_body_res_1 = ng.result(mmul_t) then_body_res_2 = ng.result(mul_t) then_body = GraphBody([X_t, Y_t], [then_body_res_1, then_body_res_2]) then_body_inputs = [ TensorIteratorInvariantInputDesc(1, 0), TensorIteratorInvariantInputDesc(2, 1) ] then_body_outputs = [ TensorIteratorBodyOutputDesc(0, 0), TensorIteratorBodyOutputDesc(1, 1) ] # else_body X_e = ng.parameter([2], np.float32, "X") Z_e = ng.parameter([], np.float32, "Z") mul_e = ng.multiply(X_e, Z_e) else_body_res_1 = ng.result(Z_e) else_body_res_2 = ng.result(mul_e) else_body = GraphBody([X_e, Z_e], [else_body_res_1, else_body_res_2]) else_body_inputs = [ TensorIteratorInvariantInputDesc(1, 0), TensorIteratorInvariantInputDesc(3, 1) ] else_body_outputs = [ TensorIteratorBodyOutputDesc(0, 0), TensorIteratorBodyOutputDesc(1, 1) ] X = ng.constant([3, 4], dtype=np.float32) Y = ng.constant([2, 1], dtype=np.float32) Z = ng.constant(4.0, dtype=np.float32) if_node = ng.if_op(condition, [X, Y, Z], (then_body, else_body), (then_body_inputs, else_body_inputs), (then_body_outputs, else_body_outputs)) return if_node
def simple_if_without_parameters(condition_val): condition = ng.constant(condition_val, dtype=np.bool) # then_body then_constant = ng.constant(0.7, dtype=np.float) then_body_res_1 = ng.result(then_constant) then_body = GraphBody([], [then_body_res_1]) then_body_inputs = [] then_body_outputs = [TensorIteratorBodyOutputDesc(0, 0)] # else_body else_const = ng.constant(9.0, dtype=np.float) else_body_res_1 = ng.result(else_const) else_body = GraphBody([], [else_body_res_1]) else_body_inputs = [] else_body_outputs = [TensorIteratorBodyOutputDesc(0, 0)] if_node = ng.if_op(condition, [], (then_body, else_body), (then_body_inputs, else_body_inputs), (then_body_outputs, else_body_outputs)) relu = ng.relu(if_node) return relu
def test_tensor_iterator(): from ngraph.utils.tensor_iterator_types import ( GraphBody, TensorIteratorSliceInputDesc, TensorIteratorMergedInputDesc, TensorIteratorInvariantInputDesc, TensorIteratorBodyOutputDesc, TensorIteratorConcatOutputDesc, ) # Body parameters body_timestep = ng.parameter([], np.int32, "timestep") body_data_in = ng.parameter([1, 2, 2], np.float32, "body_in") body_prev_cma = ng.parameter([2, 2], np.float32, "body_prev_cma") body_const_one = ng.parameter([], np.int32, "body_const_one") # CMA = cumulative moving average prev_cum_sum = ng.multiply(ng.convert(body_timestep, "f32"), body_prev_cma) curr_cum_sum = ng.add(prev_cum_sum, ng.squeeze(body_data_in, [0])) elem_cnt = ng.add(body_const_one, body_timestep) curr_cma = ng.divide(curr_cum_sum, ng.convert(elem_cnt, "f32")) cma_hist = ng.unsqueeze(curr_cma, [0]) # TI inputs data = ng.parameter([16, 2, 2], np.float32, "data") # Iterations count zero = ng.constant(0, dtype=np.int32) one = ng.constant(1, dtype=np.int32) initial_cma = ng.constant(np.zeros([2, 2], dtype=np.float32), dtype=np.float32) iter_cnt = ng.range(zero, np.int32(16), np.int32(1)) ti_inputs = [iter_cnt, data, initial_cma, one] graph_body = GraphBody( [body_timestep, body_data_in, body_prev_cma, body_const_one], [curr_cma, cma_hist]) ti_slice_input_desc = [ # timestep # input_idx, body_param_idx, start, stride, part_size, end, axis TensorIteratorSliceInputDesc(0, 0, 0, 1, 1, -1, 0), # data TensorIteratorSliceInputDesc(1, 1, 0, 1, 1, -1, 0), ] ti_merged_input_desc = [ # body prev/curr_cma TensorIteratorMergedInputDesc(2, 2, 0), ] ti_invariant_input_desc = [ # body const one TensorIteratorInvariantInputDesc(3, 3), ] # TI outputs ti_body_output_desc = [ # final average TensorIteratorBodyOutputDesc(0, 0, -1), ] ti_concat_output_desc = [ # history of cma TensorIteratorConcatOutputDesc(1, 1, 0, 1, 1, -1, 0), ] node = ng.tensor_iterator( ti_inputs, graph_body, ti_slice_input_desc, ti_merged_input_desc, ti_invariant_input_desc, ti_body_output_desc, ti_concat_output_desc, ) assert node.get_type_name() == "TensorIterator" assert node.get_output_size() == 2 # final average assert list(node.get_output_shape(0)) == [2, 2] # cma history assert list(node.get_output_shape(1)) == [16, 2, 2]