## @package control_ops_util # Module caffe2.python.control_ops_util from caffe2.python import core def get_external_blob_names(net, lexical_scope): """ Returns a set of blobs a given net depends on and a set of output blobs that are written by the net Inputs: net - net to return input/output blobs for; lexical_scope - all external blob names visible to the net """ # Use the blobs that are actually read/written to as external inputs/outputs net_proto = net.Proto() net_ssa, _ = core.get_ssa(net_proto) input_names = core.get_undefined_blobs(net_ssa) for input_name in input_names: assert str(input_name) in lexical_scope, \ "Input blob " + input_name + " is undefined" output_names = set() for op in net_proto.op: for output in op.output: if output in lexical_scope: output_names.add(output) return input_names, output_names def add_if_op(if_net, cond_blob, lexical_scope, then_net, else_net=None): """ A helper function to add an If op to the net. Automatically determines whether blobs in the then/else subnets are external (from the outer workspace) or local (visible only inside subnet's workspace) based on lexical scope - set of all outer blob names visible to the 'If' operator. All the blobs in then/else subnets with names matching a name in lexical scope and all the blobs that are first used as the operators' inputs are considered outer blobs - these blobs must exist in the outer workspace, then/else subnets can read their values and new values written into these blobs will be visible outside of the 'If' operator. All other blobs are local - exist only within inner workspaces for then/else. Inputs: if_net - net to add an If op to; cond_blob - scalar bool blob reference, used as If condition; lexical_scope - a set of outer blob names visible to then/else branches; then_net/else_net - nets (core.Net) for then/else branches """ then_input_blob_names, then_output_blob_names = get_external_blob_names( then_net, lexical_scope) else_input_blob_names = set() else_output_blob_names = set() if else_net: else_input_blob_names, else_output_blob_names = get_external_blob_names( else_net, lexical_scope) input_blob_names = then_input_blob_names | else_input_blob_names output_blob_names = then_output_blob_names | else_output_blob_names if_inputs = [cond_blob] if_inputs += [core.BlobReference(name=b, net=None) for b in input_blob_names] if_outputs = [core.BlobReference(name=b, net=None) for b in output_blob_names] do_then_net = core.Net('do_then_net') then_input_blobs = \ [core.BlobReference(name=b, net=None) for b in then_input_blob_names] then_output_blobs = \ [core.BlobReference(name=b, net=None) for b in then_output_blob_names] then_input_output_names_ordered = [ str(b) for b in (then_input_blobs + then_output_blobs)] then_outer_blob_names = list(then_input_blob_names | then_output_blob_names) then_outer_blob_names_idx = [ then_input_output_names_ordered.index(b) for b in then_outer_blob_names] # make sure to use net's name to have unique blob name across multiple subnets do_then_workspace_blob = if_net.NextScopedBlob(if_net.Name() + '/workspace_if_then') then_input_blobs.append(do_then_workspace_blob) then_output_blobs.append(do_then_workspace_blob) # make sure that added workspace pointer blobs are in if inputs/outputs if_inputs.append(do_then_workspace_blob) if_outputs.append(do_then_workspace_blob) do_then_net.Do( then_input_blobs, then_output_blobs, net=then_net.Proto(), inner_blobs=then_outer_blob_names, outer_blobs_idx=then_outer_blob_names_idx) do_then_net.AddExternalOutput(*then_output_blobs) if_args = {} if_args['then_net'] = do_then_net.Proto() do_else_workspace_blob = None if else_net: do_else_net = core.Net('do_else_net') else_input_blobs = \ [core.BlobReference(name=b, net=None) for b in else_input_blob_names] else_output_blobs = \ [core.BlobReference(name=b, net=None) for b in else_output_blob_names] else_input_output_names_ordered = [ str(b) for b in (else_input_blobs + else_output_blobs)] else_outer_blob_names = list(else_input_blob_names | else_output_blob_names) else_outer_blob_names_idx = [ else_input_output_names_ordered.index(b) for b in else_outer_blob_names] do_else_workspace_blob = \ if_net.NextScopedBlob(if_net.Name() + '/workspace_if_else') else_input_blobs.append(do_else_workspace_blob) else_output_blobs.append(do_else_workspace_blob) # make sure that added workspace pointer blobs are in if inputs/outputs if_inputs.append(do_else_workspace_blob) if_outputs.append(do_else_workspace_blob) do_else_net.Do( else_input_blobs, else_output_blobs, net=else_net.Proto(), inner_blobs=else_outer_blob_names, outer_blobs_idx=else_outer_blob_names_idx) do_else_net.AddExternalOutput(*else_output_blobs) if_args['else_net'] = do_else_net.Proto() if_net.CreateScope([], [do_then_workspace_blob]) if do_else_workspace_blob: if_net.CreateScope([], [do_else_workspace_blob]) if_net.If(if_inputs, if_outputs, **if_args) if_net.AddExternalOutput(*if_outputs) def add_while_op( while_net, cond_blob, lexical_scope, loop_body_net, condition_body_net=None): """ A helper function to add a While op to the net. Same rules for determining outer and inner blobs as for the 'If' operator apply for the 'While' operator loop and condition subnets. If specified, condition net is executed in a separate workspace before the first and after each iteration, the last operator must have a single scalar boolean output that is written into the condition blob. Inputs: while_net - net to add a While op to; cond_blob - scalar bool blob reference, used as a stop condition; lexical_scope - a set of outer blob names visible to the loop's body; loop_body_net - net to execute on each iteration; condition_body_net - net to compute condition value """ input_blob_names, output_blob_names = get_external_blob_names( loop_body_net, lexical_scope) # Since it's possible that loop is not going to run even once # we have to add loop's external outputs into inputs input_blob_names |= output_blob_names loop_inputs = [core.BlobReference(name=b, net=None) for b in input_blob_names] loop_outputs = [core.BlobReference(name=b, net=None) for b in output_blob_names] while_inputs = [cond_blob] + loop_inputs while_outputs = [] + loop_outputs do_loop_body_net = core.Net('do_loop_body_net') loop_input_output_names_ordered = [ str(b) for b in (loop_inputs + loop_outputs)] loop_body_outer_blob_names = list(input_blob_names | output_blob_names) loop_body_outer_blob_names_idx = [ loop_input_output_names_ordered.index(b) for b in loop_body_outer_blob_names] do_loop_body_workspace_blob = \ while_net.NextScopedBlob(while_net.Name() + '/workspace_loop_body') loop_inputs.append(do_loop_body_workspace_blob) loop_outputs.append(do_loop_body_workspace_blob) # make sure that added workspace pointer blobs are in While inputs/outputs while_inputs.append(do_loop_body_workspace_blob) while_outputs.append(do_loop_body_workspace_blob) do_loop_body_net.Do( loop_inputs, loop_outputs, net=loop_body_net.Proto(), inner_blobs=loop_body_outer_blob_names, outer_blobs_idx=loop_body_outer_blob_names_idx, copy_external_blobs=True) do_loop_body_net.AddExternalOutput(*loop_outputs) while_args = {} while_args['loop_net'] = do_loop_body_net.Proto() cond_workspace_blob = None if condition_body_net: cond_input_blob_names, cond_output_blob_names = get_external_blob_names( condition_body_net, lexical_scope) # make sure condition blob is written by condition net and is # visible outside of it found_condition_output = False for op in condition_body_net.Proto().op: if str(cond_blob) in op.output: found_condition_output = True break assert found_condition_output, \ "Condition net does not write into condition blob" if str(cond_blob) not in cond_output_blob_names: cond_output_blob_names.add(str(cond_blob)) cond_inputs = [core.BlobReference(name=b, net=None) for b in cond_input_blob_names] assert str(cond_blob) in cond_output_blob_names, \ 'Condition blob expected in condition net output' cond_outputs = [core.BlobReference(name=b, net=None) for b in cond_output_blob_names] condition_net = core.Net('do_loop_condition_net') cond_input_output_names_ordered = [ str(b) for b in (cond_inputs + cond_outputs)] cond_body_outer_blob_names = \ list(cond_input_blob_names | cond_output_blob_names) cond_body_outer_blob_names_idx = [ cond_input_output_names_ordered.index(b) for b in cond_body_outer_blob_names] cond_workspace_blob = \ while_net.NextScopedBlob(while_net.Name() + '/workspace_loop_cond') cond_inputs.append(cond_workspace_blob) cond_outputs.append(cond_workspace_blob) condition_net.Do( cond_inputs, cond_outputs, net=condition_body_net.Proto(), inner_blobs=cond_body_outer_blob_names, outer_blobs_idx=cond_body_outer_blob_names_idx) condition_net.AddExternalOutput(*cond_outputs) while_args['cond_net'] = condition_net.Proto() while_inputs += [b for b in cond_inputs if str(b) not in input_blob_names] while_outputs += [b for b in cond_outputs if str(b) not in output_blob_names] if str(cond_blob) not in lexical_scope: while_net.ConstantFill( [], cond_blob, dtype=core.DataType.BOOL, value=False) while_net.CreateScope([], [do_loop_body_workspace_blob]) if cond_workspace_blob: while_net.CreateScope([], [cond_workspace_blob]) while_net.While(while_inputs, while_outputs, **while_args) while_net.AddExternalOutput(*while_outputs)