MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / trace_cpu

Method trace_cpu

tensorflow/python/tpu/tensor_tracer.py:1503–1543  ·  view source on GitHub ↗

Traces the tensors generated by CPU Ops in a TF graph. Args: graph: the graph of Ops executed on the CPU. tensor_fetches: a (list,tuple,or a single object) of tensor fetches returned by model_fn given to session.run. Function must be provided with as least one tensor

(self, graph, tensor_fetches, op_fetches=None)

Source from the content-addressed store, hash-verified

1501 return tensor_fetches
1502
1503 def trace_cpu(self, graph, tensor_fetches, op_fetches=None):
1504 """Traces the tensors generated by CPU Ops in a TF graph.
1505
1506 Args:
1507 graph: the graph of Ops executed on the CPU.
1508 tensor_fetches: a (list,tuple,or a single object) of tensor fetches
1509 returned by model_fn given to session.run. Function must be provided
1510 with as least one tensor to fetch.
1511 op_fetches: A list of op fetches returned by model_fn given to
1512 session.run. op_fetches and tensor_fetches are used to determine the
1513 nodes that will be executed. Can be None.
1514
1515 Returns:
1516 tensor_fetches: an exact copy of tensor_fetches that has additional
1517 dependencies.
1518 Raises:
1519 RuntimeError: If tensor_fetches is None or empty.
1520 """
1521
1522 if graph in TensorTracer._traced_graphs:
1523 logging.warning('Graph is already rewritten with tensor tracer, ignoring '
1524 'multiple calls.')
1525 return tensor_fetches
1526 else:
1527 TensorTracer._traced_graphs.add(graph)
1528
1529 self._tt_config.device_type = _DEVICE_TYPE_CPU
1530 self._tt_config.num_replicas = 1
1531 self._tt_config.num_replicas_per_host = 1
1532 self._tt_config.num_hosts = 1
1533 self._replica_id = 0
1534 if self._parameters.graph_dump_path:
1535 graph_io.write_graph(graph, self._parameters.graph_dump_path,
1536 'graph_before_tt.pbtxt')
1537 with graph.as_default():
1538 tensor_fetches = self._trace_execution(graph, tensor_fetches, op_fetches,
1539 on_tpu=False)
1540 if self._parameters.graph_dump_path:
1541 graph_io.write_graph(graph, self._parameters.graph_dump_path,
1542 'graph_after_tt.pbtxt')
1543 return tensor_fetches

Callers

nothing calls this directly

Calls 3

_trace_executionMethod · 0.95
addMethod · 0.45
as_defaultMethod · 0.45

Tested by

no test coverage detected