Class LocalSession

caffe2/python/session.py:176–213

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 a

class LocalSession(Session):
    """
    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.
    """
    def __init__(self, ws=None):
        Session.__init__(self)
        self._ws = ws or workspace.C.Workspace.current

    @classmethod
    def _compile_task_group(cls, task_group, setup_net_list=None):
        with Cluster():
            task = task_group.to_task()
        plan = core.Plan('task_group_plan')
        plan.AddStep(task.get_step())
        return (plan, task.output_list(), task.workspace_type())

    def _run_compiled(self, compiled):
        plan, output_list, workspace_type = compiled

        # make sure the output blobs belong to the parent workspace
        outputs = []
        for name in output_list.names():
            self._ws.create_blob(str(name))
            outputs.append(core.BlobReference(str(name)))
        output_list.set_values(outputs, _fetch_func=self._fetch_output)
        task_ws = (
            workspace.C.Workspace(self._ws)
            if workspace_type == WorkspaceType.PRIVATE else self._ws)
        with workspace.WorkspaceGuard(task_ws):
            task_ws.run(plan)

    def _fetch_output(self, output):
        return self._ws.blobs[str(output)].fetch()