workspace.py 2.9 KB

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