Traces the tensors generated by TPU Ops in a TF graph. Args: graph: the graph of Ops executed on the TPU. 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,
num_replicas=None,
num_replicas_per_host=None,
num_hosts=None)
| 1442 | processed_t_fetches) |
| 1443 | |
| 1444 | def trace_tpu(self, graph, |
| 1445 | tensor_fetches, |
| 1446 | op_fetches=None, |
| 1447 | num_replicas=None, |
| 1448 | num_replicas_per_host=None, |
| 1449 | num_hosts=None): |
| 1450 | """Traces the tensors generated by TPU Ops in a TF graph. |
| 1451 | |
| 1452 | Args: |
| 1453 | graph: the graph of Ops executed on the TPU. |
| 1454 | tensor_fetches: a (list,tuple,or a single object) of tensor fetches |
| 1455 | returned by model_fn given to session.run. Function must be provided |
| 1456 | with as least one tensor to fetch. |
| 1457 | op_fetches: A list of op fetches returned by model_fn given to |
| 1458 | session.run. op_fetches and tensor_fetches are used to determine the |
| 1459 | nodes that will be executed. Can be None. |
| 1460 | num_replicas: number of replicas used on the TPU. |
| 1461 | num_replicas_per_host: number of replicas per TPU host. |
| 1462 | num_hosts: total number of TPU hosts. |
| 1463 | |
| 1464 | Returns: |
| 1465 | tensor_fetches: an exact copy of tensor_fetches that has additional |
| 1466 | dependencies. |
| 1467 | Raises: |
| 1468 | RuntimeError: If num_replicas_per_host > 8. |
| 1469 | RuntimeError: If tensor_fetches is None or empty. |
| 1470 | """ |
| 1471 | |
| 1472 | if graph in TensorTracer._traced_graphs: |
| 1473 | logging.warning('Graph is already rewritten with tensor tracer, ignoring ' |
| 1474 | 'multiple calls.') |
| 1475 | return tensor_fetches |
| 1476 | else: |
| 1477 | TensorTracer._traced_graphs.add(graph) |
| 1478 | |
| 1479 | self._tt_config.device_type = _DEVICE_TYPE_TPU |
| 1480 | self._tt_config.num_replicas = num_replicas |
| 1481 | self._tt_config.num_replicas_per_host = num_replicas_per_host |
| 1482 | self._tt_config.num_hosts = num_hosts |
| 1483 | if self._tt_config.num_replicas is not None: |
| 1484 | if self._tt_config.num_replicas_per_host is None: |
| 1485 | self._tt_config.num_replicas_per_host = 8 |
| 1486 | if self._tt_config.num_hosts is None: |
| 1487 | self._tt_config.num_hosts = ( |
| 1488 | num_replicas // self._tt_config.num_replicas_per_host + |
| 1489 | (num_replicas % self._tt_config.num_replicas_per_host > 0)) |
| 1490 | |
| 1491 | if self._parameters.graph_dump_path: |
| 1492 | graph_io.write_graph(graph, self._parameters.graph_dump_path, |
| 1493 | 'graph_before_tt.pbtxt') |
| 1494 | with graph.as_default(): |
| 1495 | self._add_replica_id_to_graph() |
| 1496 | tensor_fetches = self._trace_execution(graph, tensor_fetches, op_fetches, |
| 1497 | on_tpu=True) |
| 1498 | if self._parameters.graph_dump_path: |
| 1499 | graph_io.write_graph(graph, self._parameters.graph_dump_path, |
| 1500 | 'graph_after_tt.pbtxt') |
| 1501 | return tensor_fetches |
no test coverage detected