3
Eg @ s\ d dl mZmZ d dlmZmZmZmZ G dd deZ G dd deZ
G dd de
Z d S )
)core workspace)ClusterTask TaskGroup
WorkspaceTypec @ s e Zd ZdZdd ZdS )CompiledRunnablez? Wrapper for compiled runnable returned from session.compile() c C s || _ || _d S )N)obj
session_class)selfr r
r A/usr/local/lib64/python3.6/site-packages/caffe2/python/session.py__init__ s zCompiledRunnable.__init__N)__name__
__module__ __qualname____doc__r r r r r
r
s r c @ sz e Zd ZdZi Zdd Zdd ZedddZdd d
Z d d Z
d
d Z dd Z edddZ
dd Zdd Zdd ZdS ) SessionaD
Allows to run Nets, ExecutionSteps, Plans, Tasks and TaskGroups.
A session can potentially run in multiple nodes concurrently.
Example:
from core import Net
from caffe2.python.task import Task, TaskGroup, WorkspaceType
net = Net('test1')
net.Add([net.Const(1), net.Const(2)])
net2 = net.Clone()
step = core.execution_step('step1', [net2])
with TaskGroup(WorkspaceType.GLOBAL) as init_tg:
with Node('node1'):
n1setup = net.Net('n1setup')
n1msg = n1setup.Const('Hello from node 1.')
Task(step=n1setup)
with TaskGroup() as private_tg:
with Node('node1'):
n1 = net.Net('n1')
n1.Print(n1msg, 0)
Task(step=n1)
with Node('node2'):
n2 = net.Net('n2')
n2.Print(n2.Const('Hello from node 2.'), 0)
Task(step=n2)
session = LocalSession()
session.run(net)
session.run(step)
session.run(init_tg)
session.run(private_tg)
Global Workspace:
At the beginning of the session, a global workspace is created and kept
alive for the duration of the session.
Private Workspace:
Tasks can be run either directly on the global workspace, or they can
instantiate a private child workspace that is released after each run.
Blob visibility:
Tasks running in different nodes in parallel will always run under
different workspaces, so it must be assumed that they won't be able to
access each other's blobs. Tasks running on the same node will follow
Workspace hierarchy rules: tasks running on separate private workspaces
will only be able to share blobs defined on a common parent Workspace.
c C s
d| _ d S )NT)_open)r r r r
r N s zSession.__init__c C s | j S )N)r )r r r r
is_openQ s zSession.is_openNc C sx t |tr2| |jks.tdd| j|jjf |S || jkrF| j| S t |tr|r|j r~|j |kstdj||j n||_ |}n|d krt
j }t|d}t |t r|j
| nt |tjr|j
t |d nt |tjr:t|j dkstt|j dkr$|j
t |j d d n|j
t |j d n tjd|}|j
t |d t| j||| d }|| j|< |S )
Nz2Runnable was compiled for different session type. zNeed: %s, got: %sz Require {} but already have {})workspace_type)stepr runnable)r
)
isinstancer r
AssertionErrorr _compiled_cacher r formatZ_workspace_typer GLOBALr addr Z
ExecutionStepPlanlenZStepsZexecution_step_compile_task_group)clsr r setup_net_listtgr compiledr r r
compileT sD
zSession.compilec C s: | j std|dk s td| j| j|||j dS )a Run the given runnable.
Args:
runnable: Object recognized by the Session. Currently, we support
TaskGroup, Task, Plan, ExecutionStep, and Net.
workspace_type: A string defined in the WorkspaceType object.
setup_net_list: A list of Net objects or a list of NetDef protos.
So far this is only used by the DistributedSession, in which we
need to pass a list of special nets to setup the master.
zSession is closed.NzGot a none runnable.)r r
_run_compiledr' r )r r r r$ r r r
run s z Session.runc C s | j r| j d| _d S )NF)r _do_closer )r r r r
close s z
Session.closec C s
t d S )N)NotImplementedError)r outputr r r
fetch_output s zSession.fetch_outputc C s
t d S )N)r, )r
task_groupr r r
r( s zSession._run_compiledc C s |S )Nr )r# r/ r$ r r r
r" s zSession._compile_task_groupc C s d S )Nr )r r r r
r* s zSession._do_closec C s | j std| S )NzSession already closed.)r r )r r r r
__enter__ s zSession.__enter__c C s |d kr| j d S )N)r+ )r Zex_typevalue tracebackr r r
__exit__ s zSession.__exit__)NN)NN)N)r r r r r r r classmethodr' r) r+ r. r( r" r* r0 r3 r r r r
r s 6 0
r c @ s8 e Zd ZdZd ddZed ddZdd Zd d
ZdS )
LocalSessionaQ
Session that runs in a single node.
Tasks are all remapped to run in parallel in the 'local' node.
Currently, LocalSession runs all parallel tasks in the same workspace,
but this behavior may change in the future. Only tasks pointing to the
same logical node are guaranteed to always run in the same workspace.
Nc C s t j| |ptjjj| _d S )N)r r r C Workspacecurrent_ws)r wsr r r
r s
zLocalSession.__init__c
C sD t |j }W d Q R X tjd}|j|j ||j |j fS )NZtask_group_plan)r Zto_taskr r ZAddStepZget_step output_listr )r# r/ r$ taskplanr r r
r" s
z LocalSession._compile_task_groupc C s |\}}}g }x4|j D ](}| jjt| |jtjt| qW |j|| jd |t j
krlt j j
| jn| j}t j| |j| W d Q R X d S )N)Z _fetch_func)namesr9 Z create_blobstrappendr Z
BlobReferenceZ
set_values
_fetch_outputr ZPRIVATEr r6 r7 ZWorkspaceGuardr) )r r& r= r; r outputsnameZtask_wsr r r
r( s
zLocalSession._run_compiledc C s | j jt| j S )N)r9 Zblobsr? fetch)r r- r r r
rA s zLocalSession._fetch_output)N)N) r r r r r r4 r" r( rA r r r r
r5 s
r5 N) Z
caffe2.pythonr r Zcaffe2.python.taskr r r r objectr r r5 r r r r