def compute_cost_graph(self, devices=None):
"""Computes a CostGraphDef protobuf based on this graph.
Defined in tensorflow/core/framework/cost_graph.proto.
Args:
devices: optional [string], the names of devices to consider. If
specified, any tensor on a device not listed is given a size of zero.
Any device-less tensor (e.g. Mesh TensorFlow tensor) is not affected.
Returns:
a CostGraphDef protobuf with a Node for every operation in the graph, each
of which is populated with size/dtype information for its inputs and
outputs (which match the input/output order of the operation).
"""
cost_graph_def = cost_graph_pb2.CostGraphDef()
for i, operation_name in enumerate(self.get_all_operation_names()):
node = cost_graph_def.node.add(
name=operation_name,
device=self.get_operation_device(operation_name),
id=i)
for input_name in self.get_operation_input_names(operation_name):
id1, id2 = self._tensor_name_to_ids[input_name]
node.input_info.add(preceding_node=id1, preceding_port=id2)
for output_name in self.get_operation_output_names(operation_name):
tensor_device = self.get_tensor_device(output_name)
# devices = [] is not the same as None, and tensor_device = '' is also
# not the same as None.
if devices is None or tensor_device is None or tensor_device in devices:
node.output_info.add(
size=self.get_tensor_size(output_name),
alias_input_port=-1,
dtype=self.get_tensor_dtype(
output_name).as_datatype_enum,
shape=self.get_tensor_shape(output_name).as_proto(),
)
else:
node.output_info.add(
size=0,
alias_input_port=-1,
dtype=self.get_tensor_dtype(
output_name).as_datatype_enum,
)
# NOTE(joshuawang): Unfortunately, the CostGraphDef protobuf has final
# operations, not tensors. As a result, we have to declare any operation
# that outputs a final tensor as final, which may expand the final set
# of tensors to keep in memory. This issue also arises in the scheduler
# code we will interface with.
if self.is_tensor_final(output_name):
node.is_final = True
return cost_graph_def
def StripCostGraphDef(self, cost_graph, to_strip):
"""Removes fields from a CostGraphDef protobuf.
Helper method to reduce the initialization of CostGraphDef(s).
Args:
cost_graph: a CostGraphDef to strip.
to_strip: a string, either "SIZES" or "DEVICES".
Returns:
a new CostGraphDef with either size information or device information
stripped, as appropriate.
"""
new_cost_graph = cost_graph_pb2.CostGraphDef()
new_cost_graph.CopyFrom(cost_graph)
for node in new_cost_graph.node:
if to_strip == "SIZES":
for output_info in node.output_info:
output_info.size = 0
output_info.ClearField("shape")
if to_strip == "DEVICES":
node.ClearField("device")
return new_cost_graph
def setUp(self):
super(GraphInterfaceTest, self).setUp()
self._cost_graph = cost_graph_pb2.CostGraphDef(node=[
cost_graph_pb2.CostGraphDef.Node(
name="X",
device="/device:CPU:0",
id=0,
output_info=[
cost_graph_pb2.CostGraphDef.Node.OutputInfo(
size=48,
alias_input_port=-1,
dtype=types_pb2.DT_INT32,
shape=tensor_shape_pb2.TensorShapeProto(dim=[
tensor_shape_pb2.TensorShapeProto.Dim(size=3),
tensor_shape_pb2.TensorShapeProto.Dim(size=4),
])),
],
),
cost_graph_pb2.CostGraphDef.Node(
name="Y",
device="/device:CPU:0",
id=1,
output_info=[
cost_graph_pb2.CostGraphDef.Node.OutputInfo(
size=80,
alias_input_port=-1,
dtype=types_pb2.DT_INT32,
shape=tensor_shape_pb2.TensorShapeProto(dim=[
tensor_shape_pb2.TensorShapeProto.Dim(size=4),
tensor_shape_pb2.TensorShapeProto.Dim(size=5),
])),
],
),
cost_graph_pb2.CostGraphDef.Node(
name="Z1",
device="/device:CPU:0",
id=2,
input_info=[
cost_graph_pb2.CostGraphDef.Node.InputInfo(
preceding_node=0,
preceding_port=0,
),
cost_graph_pb2.CostGraphDef.Node.InputInfo(
preceding_node=1,
preceding_port=0,
),
],
output_info=[
cost_graph_pb2.CostGraphDef.Node.OutputInfo(
size=60,
alias_input_port=-1,
dtype=types_pb2.DT_INT32,
shape=tensor_shape_pb2.TensorShapeProto(dim=[
tensor_shape_pb2.TensorShapeProto.Dim(size=3),
tensor_shape_pb2.TensorShapeProto.Dim(size=5),
])),
],
is_final=True,
),
cost_graph_pb2.CostGraphDef.Node(
name="Z2",
device="/device:CPU:0",
id=3,
input_info=[
cost_graph_pb2.CostGraphDef.Node.InputInfo(
preceding_node=0,
preceding_port=0,
),
cost_graph_pb2.CostGraphDef.Node.InputInfo(
preceding_node=1,
preceding_port=0,
),
],
output_info=[
cost_graph_pb2.CostGraphDef.Node.OutputInfo(
size=60,
alias_input_port=-1,
dtype=types_pb2.DT_INT32,
shape=tensor_shape_pb2.TensorShapeProto(dim=[
tensor_shape_pb2.TensorShapeProto.Dim(size=3),
tensor_shape_pb2.TensorShapeProto.Dim(size=5),
])),
],
),
])
self._sizeless_cost_graph = self.StripCostGraphDef(
self._cost_graph, "SIZES")
self._deviceless_cost_graph = self.StripCostGraphDef(
self._cost_graph, "DEVICES")
self._cost_graph_string = self._cost_graph.SerializeToString()
self._sizeless_cost_graph_string = (
self._sizeless_cost_graph.SerializeToString())
self._deviceless_cost_graph_string = (
self._deviceless_cost_graph.SerializeToString())
