2
# Module caffe2.python.onnx.workspace
11
from caffe2.python import workspace
13
# Separating out the context manager part so that users won't
14
# (mis-)use Workspace instances as context managers
16
def __init__(self, workspace_id):
17
self.workspace_id = workspace_id
18
# A stack, so that the context manager is reentrant.
19
self.workspace_stack = []
22
self.workspace_stack.append(workspace.CurrentWorkspace())
23
workspace.SwitchWorkspace(self.workspace_id, create_if_missing=True)
25
def __exit__(self, exc_type, exc_value, traceback):
26
w = self.workspace_stack.pop()
27
# Strictly speaking, create_if_missing here is unnecessary, since a user
28
# is not supposed to be allowed to destruct a workspace while we're in
29
# it. However, empirically, it has been observed that during abnormal
30
# shutdown, Caffe2 deletes its default workspace fairly early in the
31
# final calls to destructors. In this case, we may attempt to exit
32
# to a default workspace which no longer exists. create_if_missing=True
33
# will (harmlessly) recreate the workspace before we finally quit.)
34
workspace.SwitchWorkspace(w, create_if_missing=True)
39
An object representing a Caffe2 workspace. It is a context manager,
40
so you can say 'with workspace:' to use the represented workspace
41
as your global workspace. It also supports every method supported
42
by caffe2.python.workspace, but instead of running these operations
43
in the global workspace, it runs them in the workspace represented
44
by this object. When this object goes dead, the workspace (and all
45
nets and blobs within it) are freed.
47
Why do we need this class? Caffe2's workspace model is very "global state"
48
oriented, in that there is always some ambient global workspace you are
49
working in which holds on to all of your networks and blobs. This class
50
makes it possible to work with workspaces more locally, and without
51
forgetting to deallocate everything in the end.
54
# Caffe2 (apparently) doesn't provide any native method of generating
55
# a fresh, unused workspace, so we have to fake it by generating
56
# a unique ID and hoping it's not used already / will not be used
57
# directly in the future.
58
self._ctx = _WorkspaceCtx(str(uuid.uuid4()))
60
def __getattr__(self, attr):
61
def f(*args, **kwargs):
63
return getattr(workspace, attr)(*args, **kwargs)
67
# NB: This is a 'self' call because we need to switch into the workspace
68
# we want to reset before we actually reset it. A direct call to
69
# workspace.ResetWorkspace() will reset the ambient workspace, which
70
# is not want we want.