Caffe2 - C++ API
A deep learning, cross platform ML framework
distributed.cc
1 #include "caffe2/opt/distributed.h"
2 #include "caffe2/opt/converter.h"
3 
4 namespace caffe2 {
5 
6 using namespace nom::repr;
7 
8 void setDeviceOption(NNGraph::NodeRef n, caffe2::DeviceOption& d) {
9  getOrAddCaffe2Annotation(n);
10  auto op = nn::get<NeuralNetOperator>(n);
11  auto c2Annot = dyn_cast<caffe2::Caffe2Annotation>(op->getMutableAnnotation());
12  CAFFE_ENFORCE(c2Annot, "getOrAddCaffe2Annotation failed!");
13  c2Annot->setDeviceOption(d);
14 }
15 
17  std::map<std::string, caffe2::DeviceOption> blobMap,
18  NNModule* nn) {
19  // Names we've seen in the NNModule
20  std::unordered_set<std::string> seen;
21 
22  auto declareNodes = nn::filter<Declare>(*nn);
23 
24  for (auto& declareNode : declareNodes) {
25  auto inputNode = nn::getOutputs(declareNode).at(0);
26  auto input = nn::get<nom::repr::Tensor>(inputNode);
27 
28  if (!blobMap.count(input->getName())) {
29  continue;
30  }
31 
32  CAFFE_ENFORCE(
33  !seen.count(input->getName()),
34  "Ambiguous name->deviceOption map. Please do this manually.");
35 
36  seen.insert(input->getName());
37  setDeviceOption(declareNode, blobMap[input->getName()]);
38  }
39 
40  auto exportNodes = nn::filter<Export>(*nn);
41 
42  for (auto& exportNode : exportNodes) {
43  auto outputNode = nn::getInputs(exportNode).at(0);
44  auto output = nn::get<nom::repr::Tensor>(outputNode);
45 
46  if (!blobMap.count(output->getName())) {
47  continue;
48  }
49 
50  CAFFE_ENFORCE(
51  !seen.count(output->getName()),
52  "Ambiguous name->deviceOption map. Please do this manually.");
53 
54  seen.insert(output->getName());
55  setDeviceOption(exportNode, blobMap[output->getName()]);
56  }
57 
58  if (seen.size() != blobMap.size()) {
59  std::ostringstream os;
60  for (const auto& kv : blobMap) {
61  if (!(seen.count(kv.first))) {
62  os << "\"" << kv.first << "\" ";
63  }
64  }
65  CAFFE_ENFORCE(
66  seen.size() == blobMap.size(),
67  "Unused names in the blob map: ",
68  os.str());
69  }
70 }
71 
72 void injectDataEdgeIndicators(nom::repr::NNModule* nn) {
73  for (auto& input : nn->inputs) {
74  auto declareNode =
75  nn->dataFlow.createNode(nom::util::make_unique<Declare>());
76  nn->dataFlow.createEdge(declareNode, input);
77  }
78 
79  for (auto& output : nn->outputs) {
80  auto exportNode = nn->dataFlow.createNode(nom::util::make_unique<Export>());
81  nn->dataFlow.createEdge(output, exportNode);
82  }
83 
84  nn->inputs.clear();
85  nn->outputs.clear();
86 }
87 
88 void removeDataEdgeIndicators(nom::repr::NNModule* nn) {
89  auto declareNodes = nn::filter<Declare>(*nn);
90  for (auto& declareNode : declareNodes) {
91  auto input = nn::getOutputs(declareNode).at(0);
92  nn->inputs.insert(input);
93  nn->dataFlow.deleteNode(declareNode);
94  }
95  auto exportNodes = nn::filter<Export>(*nn);
96  for (auto& exportNode : exportNodes) {
97  auto output = nn::getInputs(exportNode).at(0);
98  nn->outputs.insert(output);
99  nn->dataFlow.deleteNode(exportNode);
100  }
101 }
102 
104  caffe2::NetDef& net,
105  std::map<std::string, caffe2::DeviceOption> blobMap) {
106  auto nn = convertToNNModule(net);
107  injectDataEdgeIndicators(&nn);
108  addBlobDeviceOptions(blobMap, &nn);
109  return nn;
110 }
111 
112 } // namespace caffe2
NodeRef createNode(T &&data)
Creates a node and retains ownership of it.
Definition: Graph.h:240
void addBlobDeviceOptions(std::map< std::string, caffe2::DeviceOption > blobMap, NNModule *nn)
Helpers for the convertToNNModule for use if you already have an NNModule.
Definition: distributed.cc:16
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
repr::NNModule convertToNNModule(const caffe2::NetDef &net, bool strict, std::vector< repr::NNGraph::NodeRef > *opNodeVec)
Ingest a caffe2 protobuf model and output an NNModule.
Definition: converter.cc:301
EdgeRef createEdge(NodeRef tail, NodeRef head, U...data)
Creates a directed edge and retains ownership of it.
Definition: Graph.h:415