1 from __future__
import absolute_import
2 from __future__
import division
3 from __future__
import print_function
4 from __future__
import unicode_literals
8 from caffe2.proto
import caffe2_pb2
10 from subprocess
import Popen, PIPE
15 def __init__(self, net=None, device_map=None):
17 serialized_proto =
None 19 serialized_proto = net.Proto().SerializeToString()
20 elif isinstance(net, caffe2_pb2.NetDef):
21 serialized_proto = net.SerializeToString()
24 if device_map
is not None:
25 serialized_device_map = {}
27 serialized_device_map[k] = device_map[k].SerializeToString()
28 self.
_NNModule = C.NNModuleFromProtobufDistributed(serialized_proto,
29 serialized_device_map)
31 elif serialized_proto:
35 "NNModule can be constructed with core.Net or caffe2_pb2.NetDef types" 42 return self._NNModule.dataFlow()
45 def controlFlow(self):
46 return self._NNModule.getExecutionOrder()
50 return self._NNModule.dataFlow().nodes
54 return self._NNModule.dataFlow().operators
58 return self._NNModule.dataFlow().tensors
60 def createNode(self, val):
61 return self._NNModule.dataFlow().createNode(val)
63 def deleteNode(self, node):
64 return self._NNModule.dataFlow().deleteNode(node)
66 def createEdge(self, a, b):
67 return self._NNModule.dataFlow().createEdge(a, b)
69 def deleteEdge(self, a, b=None):
71 self._NNModule.dataFlow().deleteEdge(a, b)
73 self._NNModule.dataFlow().deleteEdge(a)
75 def replaceNode(self, old_node, new_node):
76 return self._NNModule.dataFlow().replaceNode(old_node, new_node)
78 def replaceProducer(self, tensor, new_producer):
79 C.replaceProducer(tensor, new_producer)
81 def replaceAllUsesWith(self, old_tensor, new_tensor):
82 C.replaceAllUsesWith(old_tensor, new_tensor)
84 def replaceAsConsumer(self, old_consumer, new_consumer):
85 C.replaceAsConsumer(old_consumer, new_consumer)
87 def replaceSubgraph(self, subgraph, new_node, inputs, outputs):
88 self._NNModule.replaceSubgraph(subgraph, new_node, inputs, outputs)
90 def deleteSubgraph(self, subgraph):
91 self._NNModule.deleteSubgraph(subgraph)
93 def createUniqueDataNode(self, prefix="_unique"):
94 return self._NNModule.createUniqueDataNode(prefix)
96 def convertToCaffe2Proto(self, old_proto=None):
98 old_proto = caffe2_pb2.NetDef()
99 output = self._NNModule.convertToCaffe2Proto(old_proto)
100 new_proto = caffe2_pb2.NetDef()
101 new_proto.ParseFromString(output)
104 def match(self, pattern):
105 for n
in self.dataFlow.getMutableNodes():
106 m = C.matchSubgraph(n, pattern)
113 cmd_exists =
lambda x: any(
114 os.access(os.path.join(path, x), os.X_OK)
115 for path
in os.environ[
"PATH"].
split(os.pathsep)
117 if cmd_exists(
"graph-easy"):
118 p = Popen(
"graph-easy", stdin=PIPE)
120 p.stdin.write(s.encode(
"utf-8"))
122 if e.errno == errno.EPIPE
or e.errno == errno.EINVAL:
134 NeuralNetOperator = C.NeuralNetOperator
135 Operator = C.NeuralNetOperator
136 NeuralNetData = C.NeuralNetData
137 Data = C.NeuralNetData
138 NNSubgraph = C.NNSubgraph
139 NNMatchGraph = C.NNMatchGraph
141 Annotation = C.Annotation
Module caffe2.python.layers.split.