pytorch

1
#include "caffe2/core/workspace.h"
2
#include "caffe2/python/pybind_workspace.h"
3

4
namespace caffe2 {
5
namespace python {
6

7
// NOLINTNEXTLINE(modernize-use-equals-default)
8
BlobFetcherBase::~BlobFetcherBase() {}
9

10
C10_DEFINE_TYPED_REGISTRY(
11
    BlobFetcherRegistry,
12
    TypeIdentifier,
13
    BlobFetcherBase,
14
    std::unique_ptr);
15

16
// gWorkspace is the pointer to the current workspace. The ownership is kept
17
// by the gWorkspaces map.
18
static Workspace* gWorkspace = nullptr;
19
static std::string gCurrentWorkspaceName;
20
// gWorkspaces allows us to define and switch between multiple workspaces in
21
// Python.
22
static std::map<std::string, std::unique_ptr<Workspace>> gWorkspaces;
23

24
Workspace* GetCurrentWorkspace() {
25
  return gWorkspace;
26
}
27

28
void SetCurrentWorkspace(Workspace* workspace) {
29
  gWorkspace = workspace;
30
}
31

32
Workspace* NewWorkspace() {
33
  std::unique_ptr<Workspace> new_workspace(new Workspace());
34
  gWorkspace = new_workspace.get();
35
  return gWorkspace;
36
}
37

38
Workspace* GetWorkspaceByName(const std::string& name) {
39
  if (gWorkspaces.count(name)) {
40
    return gWorkspaces[name].get();
41
  }
42
  return nullptr;
43
}
44

45
std::string GetCurrentWorkspaceName() {
46
  return gCurrentWorkspaceName;
47
}
48
void InsertWorkspace(const std::string& name, std::unique_ptr<Workspace> ws) {
49
  gWorkspaces.insert(std::make_pair(name, std::move(ws)));
50
}
51

52
void SwitchWorkspaceInternal(const std::string& name, bool create_if_missing) {
53
  if (gWorkspaces.count(name)) {
54
    gCurrentWorkspaceName = name;
55
    gWorkspace = gWorkspaces[name].get();
56
    return;
57
  }
58

59
  CAFFE_ENFORCE(create_if_missing);
60
  std::unique_ptr<Workspace> new_workspace(new Workspace());
61
  gWorkspace = new_workspace.get();
62
  gWorkspaces.insert(std::make_pair(name, std::move(new_workspace)));
63
  gCurrentWorkspaceName = name;
64
}
65

66
void ResetWorkspace(Workspace* workspace) {
67
  gWorkspaces[gCurrentWorkspaceName].reset(workspace);
68
  gWorkspace = gWorkspaces[gCurrentWorkspaceName].get();
69
}
70

71
void GetWorkspaceNames(std::vector<std::string>& names) {
72
  for (const auto& kv : gWorkspaces) {
73
    // NOLINTNEXTLINE(performance-inefficient-vector-operation)
74
    names.emplace_back(kv.first);
75
  }
76
}
77

78
void ClearWorkspaces() {
79
  gWorkspaces.clear();
80
}
81
} // namespace python
82
} // namespace caffe2
83