/usr/local/lib64/python3.6/site-packages/caffe2/python
Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/control.py (19309B)
## @package control
# Module caffe2.python.control
"""
Implement functions for controlling execution of nets and steps, including
Do
DoParallel
For-loop
While-loop
Do-While-loop
Switch
If
"""
from caffe2.python import core
from future.utils import viewitems
# Used to generate names of the steps created by the control functions.
# It is actually the internal index of these steps.
_current_idx = 1
_used_step_names = set()
def _get_next_step_name(control_name, base_name):
global _current_idx, _used_step_names
concat_name = '%s/%s' % (base_name, control_name)
next_name = concat_name
while next_name in _used_step_names:
next_name = '%s_%d' % (concat_name, _current_idx)
_current_idx += 1
_used_step_names.add(next_name)
return next_name
def _MakeList(input):
""" input is a tuple.
Example:
(a, b, c) --> [a, b, c]
(a) --> [a]
([a, b, c]) --> [a, b, c]
"""
if len(input) == 0:
raise ValueError(
'input cannot be empty.')
elif len(input) == 1:
output = input[0]
if not isinstance(output, list):
output = [output]
else:
output = list(input)
return output
def _IsNets(nets_or_steps):
if isinstance(nets_or_steps, list):
return all(isinstance(n, core.Net) for n in nets_or_steps)
else:
return isinstance(nets_or_steps, core.Net)
def _PrependNets(nets_or_steps, *nets):
nets_or_steps = _MakeList((nets_or_steps,))
nets = _MakeList(nets)
if _IsNets(nets_or_steps):
return nets + nets_or_steps
else:
return [Do('prepend', nets)] + nets_or_steps
def _AppendNets(nets_or_steps, *nets):
nets_or_steps = _MakeList((nets_or_steps,))
nets = _MakeList(nets)
if _IsNets(nets_or_steps):
return nets_or_steps + nets
else:
return nets_or_steps + [Do('append', nets)]
def GetConditionBlobFromNet(condition_net):
"""
The condition blob is the last external_output that must
be a single bool
"""
assert len(condition_net.Proto().external_output) > 0, (
"Condition net %s must has at least one external output" %
condition_net.Proto.name)
# we need to use a blob reference here instead of a string
# otherwise, it will add another name_scope to the input later
# when we create new ops (such as OR of two inputs)
return core.BlobReference(condition_net.Proto().external_output[-1])
def BoolNet(*blobs_with_bool_value):
"""A net assigning constant bool values to blobs. It is mainly used for
initializing condition blobs, for example, in multi-task learning, we
need to access reader_done blobs before reader_net run. In that case,
the reader_done blobs must be initialized.
Args:
blobs_with_bool_value: one or more (blob, bool_value) pairs. The net will
assign each bool_value to the corresponding blob.
returns
bool_net: A net assigning constant bool values to blobs.
Examples:
- BoolNet((blob_1, bool_value_1), ..., (blob_n, bool_value_n))
- BoolNet([(blob_1, net1), ..., (blob_n, bool_value_n)])
- BoolNet((cond_1, bool_value_1))
"""
blobs_with_bool_value = _MakeList(blobs_with_bool_value)
bool_net = core.Net('bool_net')
for blob, bool_value in blobs_with_bool_value:
out_blob = bool_net.ConstantFill(
[],
[blob],
shape=[],
value=bool_value,
dtype=core.DataType.BOOL)
bool_net.AddExternalOutput(out_blob)
return bool_net
def NotNet(condition_blob_or_net):
"""Not of a condition blob or net
Args:
condition_blob_or_net can be either blob or net. If condition_blob_or_net
is Net, the condition is its last external_output
that must be a single bool.
returns
not_net: the net NOT the input
out_blob: the output blob of the not_net
"""
if isinstance(condition_blob_or_net, core.Net):
condition_blob = GetConditionBlobFromNet(condition_blob_or_net)
else:
condition_blob = condition_blob_or_net
not_net = core.Net('not_net')
out_blob = not_net.Not(condition_blob)
not_net.AddExternalOutput(out_blob)
return not_net, out_blob
def _CopyConditionBlobNet(condition_blob):
"""Make a condition net that copies the condition_blob
Args:
condition_blob is a single bool.
returns
not_net: the net NOT the input
out_blob: the output blob of the not_net
"""
condition_net = core.Net('copy_condition_blob_net')
out_blob = condition_net.Copy(condition_blob)
condition_net.AddExternalOutput(out_blob)
return condition_net, out_blob
def MergeConditionNets(name, condition_nets, relation):
"""
Merge multi condition nets into a single condition nets.
Args:
name: name of the new condition net.
condition_nets: a list of condition nets. The last external_output
of each condition net must be single bool value.
relation: can be 'And' or 'Or'.
Returns:
- A new condition net. Its last external output is relation of all
condition_nets.
"""
if not isinstance(condition_nets, list):
return condition_nets
if len(condition_nets) <= 1:
return condition_nets[0] if condition_nets else None
merged_net = core.Net(name)
for i in range(len(condition_nets)):
net_proto = condition_nets[i].Proto()
assert net_proto.device_option == merged_net.Proto().device_option
assert net_proto.type == merged_net.Proto().type
merged_net.Proto().op.extend(net_proto.op)
merged_net.Proto().external_input.extend(net_proto.external_input)
# discard external outputs as we're combining them together
curr_cond = GetConditionBlobFromNet(condition_nets[i])
if i == 0:
last_cond = curr_cond
else:
last_cond = merged_net.__getattr__(relation)([last_cond, curr_cond])
# merge attributes
for k, v in viewitems(condition_nets[i]._attr_dict):
merged_net._attr_dict[k] += v
merged_net.AddExternalOutput(last_cond)
return merged_net
def CombineConditions(name, condition_nets, relation):
"""
Combine conditions of multi nets into a single condition nets. Unlike
MergeConditionNets, the actual body of condition_nets is not copied into
the combine condition net.
One example is about multi readers. Each reader net has a reader_done
condition. When we want to check whether all readers are done, we can
use this function to build a new net.
Args:
name: name of the new condition net.
condition_nets: a list of condition nets. The last external_output
of each condition net must be single bool value.
relation: can be 'And' or 'Or'.
Returns:
- A new condition net. Its last external output is relation of all
condition_nets.
"""
if not condition_nets:
return None
if not isinstance(condition_nets, list):
raise ValueError('condition_nets must be a list of nets.')
if len(condition_nets) == 1:
condition_blob = GetConditionBlobFromNet(condition_nets[0])
condition_net, _ = _CopyConditionBlobNet(condition_blob)
return condition_net
combined_net = core.Net(name)
for i in range(len(condition_nets)):
curr_cond = GetConditionBlobFromNet(condition_nets[i])
if i == 0:
last_cond = curr_cond
else:
last_cond = combined_net.__getattr__(relation)(
[last_cond, curr_cond])
combined_net.AddExternalOutput(last_cond)
return combined_net
def Do(name, *nets_or_steps):
"""
Execute the sequence of nets or steps once.
Examples:
- Do('myDo', net1, net2, ..., net_n)
- Do('myDo', list_of_nets)
- Do('myDo', step1, step2, ..., step_n)
- Do('myDo', list_of_steps)
"""
nets_or_steps = _MakeList(nets_or_steps)
if (len(nets_or_steps) == 1 and isinstance(
nets_or_steps[0], core.ExecutionStep)):
return nets_or_steps[0]
else:
return core.scoped_execution_step(
_get_next_step_name('Do', name), nets_or_steps)
def DoParallel(name, *nets_or_steps):
"""
Execute the nets or steps in parallel, waiting for all of them to finish
Examples:
- DoParallel('pDo', net1, net2, ..., net_n)
- DoParallel('pDo', list_of_nets)
- DoParallel('pDo', step1, step2, ..., step_n)
- DoParallel('pDo', list_of_steps)
"""
nets_or_steps = _MakeList(nets_or_steps)
if (len(nets_or_steps) == 1 and isinstance(
nets_or_steps[0], core.ExecutionStep)):
return nets_or_steps[0]
else:
return core.scoped_execution_step(
_get_next_step_name('DoParallel', name),
nets_or_steps,
concurrent_substeps=True)
def _RunOnceIf(name, condition_blob_or_net, nets_or_steps):
"""
Execute nets_or_steps once if condition_blob_or_net evaluates as true.
If condition_blob_or_net is Net, the condition is its last external_output
that must be a single bool. And this net will be executed before
nets_or_steps so as to get the condition.
"""
condition_not_net, stop_blob = NotNet(condition_blob_or_net)
if isinstance(condition_blob_or_net, core.Net):
nets_or_steps = _PrependNets(
nets_or_steps, condition_blob_or_net, condition_not_net)
else:
nets_or_steps = _PrependNets(nets_or_steps, condition_not_net)
def if_step(control_name):
return core.scoped_execution_step(
_get_next_step_name(control_name, name),
nets_or_steps,
should_stop_blob=stop_blob,
only_once=True,
)
if _IsNets(nets_or_steps):
bool_net = BoolNet((stop_blob, False))
return Do(name + '/_RunOnceIf',
bool_net, if_step('_RunOnceIf-inner'))
else:
return if_step('_RunOnceIf')
def _RunOnceIfNot(name, condition_blob_or_net, nets_or_steps):
"""
Similar to _RunOnceIf() but Execute nets_or_steps once if
condition_blob_or_net evaluates as false.
"""
if isinstance(condition_blob_or_net, core.Net):
condition_blob = GetConditionBlobFromNet(condition_blob_or_net)
nets_or_steps = _PrependNets(nets_or_steps, condition_blob_or_net)
else:
copy_net, condition_blob = _CopyConditionBlobNet(condition_blob_or_net)
nets_or_steps = _PrependNets(nets_or_steps, copy_net)
return core.scoped_execution_step(
_get_next_step_name('_RunOnceIfNot', name),
nets_or_steps,
should_stop_blob=condition_blob,
only_once=True,
)
def For(name, nets_or_steps, iter_num):
"""
Execute nets_or_steps iter_num times.
Args:
nets_or_steps: a ExecutionStep or a Net or a list of ExecutionSteps or
a list nets.
iter_num: the number times to execute the nets_or_steps.
Returns:
A ExecutionStep instance.
"""
init_net = core.Net('init-net')
iter_cnt = init_net.CreateCounter([], init_count=iter_num)
iter_net = core.Net('For-iter')
iter_done = iter_net.CountDown([iter_cnt])
for_step = core.scoped_execution_step(
_get_next_step_name('For-inner', name),
_PrependNets(nets_or_steps, iter_net),
should_stop_blob=iter_done)
return Do(name + '/For',
Do(name + '/For-init-net', init_net),
for_step)
def While(name, condition_blob_or_net, nets_or_steps):
"""
Execute nets_or_steps when condition_blob_or_net returns true.
Args:
condition_blob_or_net: If it is an instance of Net, its last
external_output must be a single bool.
nets_or_steps: a ExecutionStep or a Net or a list of ExecutionSteps or
a list nets.
Returns:
A ExecutionStep instance.
"""
condition_not_net, stop_blob = NotNet(condition_blob_or_net)
if isinstance(condition_blob_or_net, core.Net):
nets_or_steps = _PrependNets(
nets_or_steps, condition_blob_or_net, condition_not_net)
else:
nets_or_steps = _PrependNets(nets_or_steps, condition_not_net)
def while_step(control_name):
return core.scoped_execution_step(
_get_next_step_name(control_name, name),
nets_or_steps,
should_stop_blob=stop_blob,
)
if _IsNets(nets_or_steps):
# In this case, while_step has sub-nets:
# [condition_blob_or_net, condition_not_net, nets_or_steps]
# If stop_blob is pre-set to True (this may happen when While() is
# called twice), the loop will exit after executing
# condition_blob_or_net. So we use BootNet to set stop_blob to
# False.
bool_net = BoolNet((stop_blob, False))
return Do(name + '/While', bool_net, while_step('While-inner'))
else:
return while_step('While')
def Until(name, condition_blob_or_net, nets_or_steps):
"""
Similar to While() but execute nets_or_steps when
condition_blob_or_net returns false
"""
if isinstance(condition_blob_or_net, core.Net):
stop_blob = GetConditionBlobFromNet(condition_blob_or_net)
nets_or_steps = _PrependNets(nets_or_steps, condition_blob_or_net)
else:
stop_blob = core.BlobReference(str(condition_blob_or_net))
return core.scoped_execution_step(
_get_next_step_name('Until', name),
nets_or_steps,
should_stop_blob=stop_blob)
def DoWhile(name, condition_blob_or_net, nets_or_steps):
"""
Execute nets_or_steps when condition_blob_or_net returns true. It will
execute nets_or_steps before evaluating condition_blob_or_net.
Args:
condition_blob_or_net: if it is an instance of Net, tts last external_output
must be a single bool.
nets_or_steps: a ExecutionStep or a Net or a list of ExecutionSteps or
a list nets.
Returns:
A ExecutionStep instance.
"""
condition_not_net, stop_blob = NotNet(condition_blob_or_net)
if isinstance(condition_blob_or_net, core.Net):
nets_or_steps = _AppendNets(
nets_or_steps, condition_blob_or_net, condition_not_net)
else:
nets_or_steps = _AppendNets(nets_or_steps, condition_not_net)
# If stop_blob is pre-set to True (this may happen when DoWhile() is
# called twice), the loop will exit after executing the first net/step
# in nets_or_steps. This is not what we want. So we use BootNet to
# set stop_blob to False.
bool_net = BoolNet((stop_blob, False))
return Do(name + '/DoWhile', bool_net, core.scoped_execution_step(
_get_next_step_name('DoWhile-inner', name),
nets_or_steps,
should_stop_blob=stop_blob,
))
def DoUntil(name, condition_blob_or_net, nets_or_steps):
"""
Similar to DoWhile() but execute nets_or_steps when
condition_blob_or_net returns false. It will execute
nets_or_steps before evaluating condition_blob_or_net.
Special case: if condition_blob_or_net is a blob and is pre-set to
true, then only the first net/step of nets_or_steps will be executed and
loop is exited. So you need to be careful about the initial value the
condition blob when using DoUntil(), esp when DoUntil() is called twice.
"""
if not isinstance(condition_blob_or_net, core.Net):
stop_blob = core.BlobReference(condition_blob_or_net)
return core.scoped_execution_step(
_get_next_step_name('DoUntil', name),
nets_or_steps,
should_stop_blob=stop_blob)
nets_or_steps = _AppendNets(nets_or_steps, condition_blob_or_net)
stop_blob = GetConditionBlobFromNet(condition_blob_or_net)
# If stop_blob is pre-set to True (this may happen when DoWhile() is
# called twice), the loop will exit after executing the first net/step
# in nets_or_steps. This is not what we want. So we use BootNet to
# set stop_blob to False.
bool_net = BoolNet((stop_blob, False))
return Do(name + '/DoUntil', bool_net, core.scoped_execution_step(
_get_next_step_name('DoUntil-inner', name),
nets_or_steps,
should_stop_blob=stop_blob,
))
def Switch(name, *conditions):
"""
Execute the steps for which the condition is true.
Each condition is a tuple (condition_blob_or_net, nets_or_steps).
Note:
1. Multi steps can be executed if their conditions are true.
2. The conditions_blob_or_net (if it is Net) of all steps will be
executed once.
Examples:
- Switch('name', (cond_1, net_1), (cond_2, net_2), ..., (cond_n, net_n))
- Switch('name', [(cond_1, net1), (cond_2, net_2), ..., (cond_n, net_n)])
- Switch('name', (cond_1, net_1))
"""
conditions = _MakeList(conditions)
return core.scoped_execution_step(
_get_next_step_name('Switch', name),
[_RunOnceIf(name + '/Switch', cond, step) for cond, step in conditions])
def SwitchNot(name, *conditions):
"""
Similar to Switch() but execute the steps for which the condition is False.
"""
conditions = _MakeList(conditions)
return core.scoped_execution_step(
_get_next_step_name('SwitchNot', name),
[_RunOnceIfNot(name + '/SwitchNot', cond, step)
for cond, step in conditions])
def If(name, condition_blob_or_net,
true_nets_or_steps, false_nets_or_steps=None):
"""
condition_blob_or_net is first evaluated or executed. If the condition is
true, true_nets_or_steps is then executed, otherwise, false_nets_or_steps
is executed.
If condition_blob_or_net is Net, the condition is its last external_output
that must be a single bool. And this Net will be executred before both
true/false_nets_or_steps so as to get the condition.
"""
if not false_nets_or_steps:
return _RunOnceIf(name + '/If',
condition_blob_or_net, true_nets_or_steps)
if isinstance(condition_blob_or_net, core.Net):
condition_blob = GetConditionBlobFromNet(condition_blob_or_net)
else:
condition_blob = condition_blob_or_net
return Do(
name + '/If',
_RunOnceIf(name + '/If-true',
condition_blob_or_net, true_nets_or_steps),
_RunOnceIfNot(name + '/If-false', condition_blob, false_nets_or_steps)
)
def IfNot(name, condition_blob_or_net,
true_nets_or_steps, false_nets_or_steps=None):
"""
If condition_blob_or_net returns false, executes true_nets_or_steps,
otherwise executes false_nets_or_steps
"""
if not false_nets_or_steps:
return _RunOnceIfNot(name + '/IfNot',
condition_blob_or_net, true_nets_or_steps)
if isinstance(condition_blob_or_net, core.Net):
condition_blob = GetConditionBlobFromNet(condition_blob_or_net)
else:
condition_blob = condition_blob_or_net
return Do(
name + '/IfNot',
_RunOnceIfNot(name + '/IfNot-true',
condition_blob_or_net, true_nets_or_steps),
_RunOnceIf(name + '/IfNot-false', condition_blob, false_nets_or_steps)
)