/usr/local/lib64/python3.6/site-packages/caffe2/python/__pycache__
Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/__pycache__/task.cpython-36.pyc (22207B)
3
Eg^ @ s d dl mZmZ d dlmZmZ d dlmZ d dlmZ d dl m
Z
dd Z G dd d ej Z
G d
d d ej ZG d d
d
eZdd Zdd ZG dd dejZG dd deZdd ZG dd deZG dd dejZG d d d eZd S ) )corecontext)Fieldfrom_blob_list) defaultdict)copy) viewitemsc C s. | d kr |S |d kr| S t | }|j| |S )N)r update)abc r
>/usr/local/lib64/python3.6/site-packages/caffe2/python/task.py_merge_node_kwargs s
r c @ s8 e Zd ZdZdd Zdd Zdd Zdd Zd
d Zd S )
Clusterz
Context that keeps track of all the node names used.
Users shouldn't have to use them directly, since a Cluster is automatically
generated at the first usage of 'Node'.
c C s g | _ i | _d S )N)_nodes _node_kwargs)selfr
r
r __init__ s zCluster.__init__c C sF t || jkr | jjt | t|j | jjt || jt |< d S )N)strr appendr kwargsr get)r noder
r
r add_node" s
zCluster.add_nodec C s | j S )zQ
Returns the list of unique node names used within this context.
)r )r r
r
r nodes) s z
Cluster.nodesc C s | j S )N)r )r r
r
r node_kwargs/ s zCluster.node_kwargsc C s dj | j | j S )Nz!Cluster(nodes={}, node_kwargs={}))formatr r )r r
r
r __repr__2 s zCluster.__repr__N) __name__
__module__ __qualname____doc__r r r r r r
r
r
r r s r c @ s2 e Zd ZdZd ddZdd Zdd Zd d
Zd S )
Nodea
A Node context is used to indicate that all Tasks instantiated within will
run on the given node name. (Only the name of the node actually counts.)
Example:
with TaskGroup() as tg:
with Node('node1'):
s1 = execution_step(...)
Task(step=s1)
with Node('node2'):
s2 = execution_step(...)
with Node('node1'):
s3 = execution_step(...)
In this example, all three execution steps will run in parallel.
Moreover, s1 and s3 will run on the same node, and can see each
others blobs.
Additionally, a Node can be passed implementation-specific kwargs,
in order to specify properties of the node.
localc K s" t || _|| _tj j| d S )N)r _name_kwargsr currentr )r r r r
r
r r N s
z
Node.__init__c C s | j S )N)r% )r r
r
r __str__S s z Node.__str__c C s dj | j| jS )NzNode(name={}, kwargs={}))r r% r& )r r
r
r r V s z
Node.__repr__c C s | j S )N)r& )r r
r
r r Y s z Node.kwargsN)r$ )r r r! r" r r( r r r
r
r
r r# 7 s
r# c @ s e Zd ZdZdZdZdS )
WorkspaceTypez
Determines whether tasks of a TaskGroup will run directly at the global
workspace, which is kept alive across runs, or whether a new child
workspace will be created for the run and destroyed afterwards.
privateglobalN)r r r! r" PRIVATEGLOBALr
r
r
r r) ] s r) c C s t j| d }t j| d }g }g }g }x>|D ]6}t|drL||j| 7 }q.t|dr.||j| 7 }q.W x|D ]} t| dr| jrqpt| dr| j|krqpt| dr
| j|}
t|
t t
fr||
7 }n8t|
t jt j fr|j |
n|
d k rt
dt|
d | _t| d
rp| j|}
t|
t t
fr8||
7 }n:t|
t jt j frX|j |
n|
d k rrt
dt|
d | _qpW t|j jd kr|jd | t|j jd kr|jd | ||fS ) Nz/initz/exitget_all_attributesget_attributes _setup_used
_setup_targetsetupz Unsupported type for setup: %sTexitr )r Nethasattrr. r/ r0 r1 r2
isinstancelisttuple
ExecutionStepr TypeErrortyper3 lenProtoopinsert) keyZ
steps_or_netstargetinit_netexit_net init_nets exit_netsobjsZ step_or_netobjnetsr
r
r get_setup_netsg sL
rI c C sf | r| r| S g }|r.|j tjd| | |j | t|dkrZ|j tjd| | tj||S )Nz%s:initr z%s:exit)r r execution_stepr< )steprD rE namestepsr
r
r add_setup_steps s
rN c @ s e Zd ZdZdZd ddZdd Zdd Zd
d Zd d
Z dd Z
dd Z d!ddZ d"ddZ
d#ddZd$ddZd d Zd d ZdS )% TaskGroupa
Context that gathers tasks which will run concurrently, potentially on
multiple nodes. All tasks in the same node will share the same workspace
and thus can share blobs, while tasks running in different nodes won't
be able to directly share data.
All tasks of the task group will start concurrently, and the task group
will finish execution when the last task of the group finishes.
Example:
# suppose that s1 ... s5 are execution steps or nets.
with TaskGroup() as tg:
# these tasks go to default node 'local'
Task(step=s1)
Task(step=s2)
with Node('n2'):
Task(step=s3)
with Node('n1'):
Task(step=s4)
with Node('n2'):
Task(step=s5)
# this will run all steps in parallel.
# s1 and s2 will run at default node 'local'
# s3 and s5 will run at node 'n2'
# s4 will run at node 'n1'
session.run(tg)
Z local_setupNc C s@ d | _ g | _d| _d | _g | _i | _g | _|| _d | _g | _ d S )NF)
Z _plan_cache_tasks
_already_usedZ _prev_active
_tasks_to_add _report_nets
_report_steps_workspace_type_tasks_by_node _remote_nets)r workspace_typer
r
r r s zTaskGroup.__init__c C s | j j| d S )N)rW r )r netr
r
r add_remote_net s zTaskGroup.add_remote_netc C s | j S )N)rW )r r
r
r remote_nets s zTaskGroup.remote_netsc C sv | j std| jd ks4|jd ks4| j|jks4t|jd krL| jpHtj|_| jd kr^|j| _|j | jj| d S )Nz-Cannot add Task to an already used TaskGroup.)rQ AssertionErrorrU r) r, _notify_usedrP r )r taskr
r
r add s
z
TaskGroup.addc C s, x| j D ]}| j| qW g | _ d| _| jS )NT)rR r_ rQ rP )r r^ r
r
r tasks s
zTaskGroup.tasksc C s t | jt | j S )N)r< rR rP )r r
r
r num_registered_tasks s z TaskGroup.num_registered_tasksc C s4 g }x*| j | j D ]}|j|kr|j|j qW |S )N)rP rR r r )r usedr^ r
r
r
used_nodes s
zTaskGroup.used_nodes c C s6 t j|}|j| | jjt|p(tj||f dS )a6
Add a "report step" to this TaskGroup. This step will run repeatedly
every `interval_ms` milliseconds for the duration of the TaskGroup
execution on each of the nodes. It is guaranteed that this step
will be run at least once after every Task in the node has finished.
N)r to_execution_stepRunEveryMillisrT r r r# r' )r rK r interval_msr
r
r report_step s
zTaskGroup.report_step c C s` t |ptj|}|dks(|| jks(t|| jkrR|r:|n tjd| |f| j|< | j| d S )z6
DEPRECATED. Use report_step instead.
Nz %s/reporterr )r r# r' rS r\ r r4 )r rY r Zreport_intervalr
r
r
report_net s
zTaskGroup.report_netc C s i }x*| j D ] }|r ||jn|j||j< qW | jd k rX| j\}}||ksTtd|S x.t| jD ] \}\}}| j|||d d qdW i | _tt}x(| j D ] }||j } || j | qW tt}
x$| j
D ]\} } |
|| j | qW t }
xt|D ]\}}|
| }t t j
dd |D | | \}}|}g }tj}x`|D ]X}|j } | j|j tjk | d k rz|j | ||j 7 }|j tjkrFtj}qFW t|dkr|j tjdg t|dkr|d } ntjd | |d
d } t|dks
t|dkrlg }t|dkr2|j tjd | | |j | t|dkr`|j tjd
| | tj||} t|| |d|
|d q W |
|f| _|
S )Nz)Cannot call tasks_by_node multiple times.i )r rg c S s g | ] }|j qS r
)get_step).0tr
r
r
) s z+TaskGroup.tasks_by_node..r empty z%s:bodyT)Zconcurrent_substepsz%s:initz%s:exitgrouped_by_node)r rK outputsrL grouprX )r` r rV r\ r rS rh r r7 r rT rO rI LOCAL_SETUPr) r, rk SetCreateWorkspacerX rr r- r< r rJ Task)r Z
node_remapZnode_mapr^
tasks_by_nodeZ
prev_node_mapr rY intervalZ mapped_nodeZreport_steps_by_nodeZ
original_noderK rq r` report_stepsZ
node_initsZ
node_exitsrM rr Zgrouped_workspace_typer
r
r rw sv
zTaskGroup.tasks_by_nodec s> t tj | j fddj }t|dkr6t S |d S )Nc s S )Nr
)x)r r
r T s z#TaskGroup.to_task..r )r r# r' rw r` r< rv )r r r` r
)r r to_taskR s
zTaskGroup.to_taskc C s | j S )N)rU )r r
r
r rX Y s zTaskGroup.workspace_typec C s dj | j| j | j | j S )Nz6TaskGroup(tasks={}, workspace_type={}, remote_nets={}))r rP rR rX r[ )r r
r
r r \ s
zTaskGroup.__repr__)N)NNrd )NNri )N)N)r r r! r" rt r rZ r[ r_ r` ra rc rh rj rw r| rX r r
r
r
r rO s
I
rO c @ s: e Zd ZdZdd Zd
ddZdd Zd d
Zd d ZdS )
TaskOutputzd
Represents the output of a task. An output can be a blob,
a list of blob, or a record.
c C sT d | _ d| _t|tr&|| _ | j j }t|ttfk| _| jrD|g}|| _d | _ d S )NF)
_schema
_is_scalarr6 r Z field_blobsr; r8 r7 names_values)r r r
r
r r i s
zTaskOutput.__init__Nc C s& t |t | jkst|| _|| _d S )N)r< r r\ r _fetch_func)r valuesr r
r
r setu s zTaskOutput.setc C s@ | j d k std| jr"| j d S | jr6t| j| j S | j S d S )NzOutput value not set yet.r )r r\ r r~ r )r r
r
r r z s
zTaskOutput.getc sN j d k std fdd jD } jr4|d S jrFt j|S |S d S )Nz#Cannot fetch value for this output.c s g | ]} j |qS r
)r )rl v)r r
r rn s z$TaskOutput.fetch..r )r r\ r r r~ r )r Z fetched_valsr
)r r fetch s zTaskOutput.fetchc C s dj | j| jS )Nz TaskOutput(names={}, values={}))r r r )r r
r
r r s zTaskOutput.__repr__)N) r r r! r" r r r r r r
r
r
r r} c s
r} c C s t jddpt }|j| S )a
Adds an output to the current Task, or if no task is active,
create a dummy task that returns the given blob or record
to the client. This will return the value of the blob or record when
the last task of the TaskGroup for a given node finishes.
F)required)rv r'
add_output)Zblob_or_recordZcur_taskr
r
r final_output s r c @ s4 e Zd ZdZd ddZdd Zd ddZd d
ZdS )
TaskOutputListz$ Keeps a list of outputs for a task Nc C s |pg | _ d S )N)rr )r rr r
r
r r s zTaskOutputList.__init__c C s" g }x| j D ]}||j7 }q W |S )z[
Retrive the output names.
TODO(azzolini): make this schema-based.
)rr r )r r or
r
r r s zTaskOutputList.namesc C sV d}x8| j D ].}t|j}|j|||| | ||7 }q W |t|ksRtdd S )Nr z Wrong number of output values.)rr r< r r r\ )r r r offsetr numr
r
r
set_values s
zTaskOutputList.set_valuesc C s dj | jS )NzTaskOutputList(outputs={}))r rr )r r
r
r r s zTaskOutputList.__repr__)N)N)r r r! r" r r r r r
r
r
r r s
r c s e Zd ZdZdZdZdZe Ze dd Z
d"dd Z fd
d Z fd d
Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zd d Zd d Zd d! Z ZS )#rv aA
A Task is composed of an execution step and zero or more outputs.
Tasks are executed in the context of a TaskGroup, which, in turn, can
be run by a Session.
Task outputs are fetched by the session at the end of the run.
The recommended way of creating a task is by using `net_builder.ops`.
Example:
from net_builder import ops
with Node('trainer'), Task(name='my_task', num_instances=2):
with ops.task_init():
globl = ops.Const(0)
with ops.task_instance_init():
local = ops.Const(0)
with ops.loop(100):
ops.Copy(globl, local)
with ops.task_instance_exit():
ops.Add([globl, local], [globl])
with ops.task_exit():
ops.Mul([globl, globl], [globl])
The task above will create 2 instances that will run in parallel.
Each instance will copy `local` to `globl` 100 times, Then Add `local`
to `globl` once. The `Mul` will only execute once, after all the instances
of the task have finished.
Z
task_setupZtask_instance_setuprh c C sd t | d t | }|d kr"tjntdd |jD }|}d}x ||kr^|d7 }d||f }q@W |S )N/c s s | ] }|j V qd S )N)rL )rl rm r
r
r s z&Task._get_next_name..r rp z%s:%d)r rv _global_names_usedr rR )r rs rL basenameZ
names_usedZcur_nameir
r
r _get_next_name s
zTask._get_next_nameNc C s | r t |tjr |j j}|s$d}ttj|dkr6dnt|| _t j|dd| _
t j | j| j
|| _| j
dk r| j
j
j| d| _d| _d| _g | _|dk r| j| |dk r| j| d| _d| _|| _d| _|| _dS )aV
Instantiate a Task and add it to the current TaskGroup and Node.
Args:
step: If provided, this task will run this ExecutionStep.
outputs: If provided, the task will return the provided outputs
to the client at completion time.
node: If provided, force task execution on the given node.
name: Name of the Task.
num_instances: If provided, this task will be cloned num_instances
times at runtime, and all instances will run
concurrently.
r^ NF)r )r6 r r9 r= rL r r# r' r rO rs rv r rR r rQ _step_step_with_setup_outputsset_step add_outputsZ _pipelineZ_is_pipeline_contextrU Z _report_net_num_instances)r rK rr rX rs r rL Z
num_instancesr
r
r r s,
z
Task.__init__c sj t t| j | jd k r&| jjj| | j | jd ks@tdddl m
} |j | j d| _
| j
j | S )Nz(This Task already has an execution step.r ) net_builder)Z _fullname)superrv __enter__rs rR remove_assert_not_usedr r\
caffe2.pythonr Z
NetBuilderrL _net_builder)r r ) __class__r
r r s
zTask.__enter__c sZ t t| j||| | jj||| |d kr8| j| j | jd k rP| jjj| d | _d S )N)r rv __exit__r r rs rR r )r r; value traceback)r r
r r ' s
z
Task.__exit__c C s | j S )N)rU )r r
r
r rX 1 s zTask.workspace_typec C s | j stdd S )Nz1Cannot modify task since it is already been used.)rQ r\ )r r
r
r r 4 s
zTask._assert_not_usedc C s. | j t|tr|nt|}| jj| |S )N)r r6 r} r r )r outputr
r
r r 8 s zTask.add_outputc s8 j t|ttfkr" j|S fdd|D S d S )Nc s g | ]} j |qS r
)r )rl r )r r
r rn D s z$Task.add_outputs..)r r; r7 r8 r )r rr r
)r r r ? s
zTask.add_outputsc C s | j tj|| _d S )N)r r re r )r rK r
r
r r F s z
Task.set_stepc
C sj | j d k r| j S | jd kr0tj| jg | _ | j S dd | jjtjD }x&|D ] }d|_|j j
sN|j d qNW t tj
| jg| | \}}t tj| jg| | \}}t| jdkrtjd| j }| j|jg dtjjdd |j| |s| jntjd | j || jg }t|||| jd
} | jrR| jdkrR| jd tjd | g| jd } t| ||| j| _ | j S )
Nc S s g | ]}t |d s|qS )_report_step_used)r5 )rl sr
r
r rn S s z!Task.get_step..Ti r z %s:outputrp )Zdtyper z%s:bodyz :instancez %s:parallel)Znum_concurrent_instances)r r r rJ rL r. rv REPORT_STEPr r= Z run_every_msrf rI
TASK_SETUPTASK_INSTANCE_SETUPr< r r4 r Z ConstantFillZDataTypeZINT32r rN r ru )
r ry rK Ztask_init_netsZtask_exit_netsZinstance_init_netsZinstance_exit_netsZ
output_netbodyZstep_with_instance_setupr
r
r rk J sD
z
Task.get_stepc C s
t | jS )N)r r )r r
r
r output_listy s zTask.output_listc C s | j S )N)r )r r
r
r rr | s z Task.outputsc C s | j d| _d S )NT)rk rQ )r r
r
r r] s zTask._notify_usedc C s dj | j| j| j S )Nz"Task(name={}, node={}, outputs={}))r rL r rr )r r
r
r r s z
Task.__repr__)NNNNNNN)r r r! r" r r r r r staticmethodr r r r rX r r r r rk r rr r] r
__classcell__r
r
)r r rv s*
+
/rv c @ s2 e Zd ZdZd ddZdd Zdd Zd d
ZdS ) SetupNetsa
Allow to register a list of nets to be run at initialization
and finalization of Tasks or TaskGroups.
For example, let's say you have the following:
init_net = core.Net('init')
my_val = init_net.ConstantFill([], 'my_val', value=0)
net = core.Net('counter')
net.Add([my_val, net.Const(1),], [my_val])
with TaskGroup() as task_group:
with Node('trainer'):
my_task = Task(step=[net])
In order to have `init_net` run once before `net` runs for the
first time, you can do one of the following:
net.add_attribute(Task.TASK_SETUP, SetupNets([init_net]))
or
net.add_attribute(TaskGroup.LOCAL_SETUP, SetupNets([init_net]))
- With Task.TASK_SETUP, init_net will run once at my_task startup.
- With TaskGroup.LOCAL_SETUP, init_net will run once on node 'trainer',
before any task of the task group is run on that node.
The same SetupNets object can be added to multiple nets. It will only
run once per Task/TaskGroup run.
Nc C s || _ || _d S )N)rD rE )r rD rE r
r
r r s zSetupNets.__init__c C s | j S )N)rD )r rB r
r
r r2 s zSetupNets.setupc C s | j S )N)rE )r rC r
r
r r3 s zSetupNets.exitc C s dj | j| jS )Nz%SetupNets(init_nets={}, exit_nets={}))r rD rE )r r
r
r r s zSetupNets.__repr__)NN)r r r! r" r r2 r3 r r
r
r
r r s
r N)r r r Zcaffe2.python.schemar r collectionsr r Z future.utilsr r ZDefaultManagedr r# objectr) rI rN ZManagedrO r} r r rv r r
r
r
r s$ !&
, E/ Q