Rewrites `computation` for inference on a TPU system. Other than 'rewriting' the computation to run on a TPU, if using variables in your computation, it moves the ReadVariableOps outside the TPU computation, and adds GuaranteeConst ops just after the ReadVariableOps. This mechan
(computation,
inputs=None,
infeed_queue=None,
device_assignment=None,
name=None)
| 1611 | |
| 1612 | |
| 1613 | def rewrite_for_inference(computation, |
| 1614 | inputs=None, |
| 1615 | infeed_queue=None, |
| 1616 | device_assignment=None, |
| 1617 | name=None): |
| 1618 | """Rewrites `computation` for inference on a TPU system. |
| 1619 | |
| 1620 | Other than 'rewriting' the computation to run on a TPU, if using variables |
| 1621 | in your computation, it moves the ReadVariableOps outside the TPU |
| 1622 | computation, and adds GuaranteeConst ops just after the ReadVariableOps. |
| 1623 | This mechanism works only if you are using tf.compat.v1.get_variable() to |
| 1624 | create and access variables in your tpu computation. You can validate |
| 1625 | whether this worked, by calling validate_inference_rewrite_for_variables() |
| 1626 | method immediately after this method to check whether GuaranteeConstOps |
| 1627 | where added to the graph. |
| 1628 | |
| 1629 | Args: |
| 1630 | computation: A Python function that builds a computation to apply to the |
| 1631 | input. If the function takes n inputs, 'inputs' should be a list of n |
| 1632 | tensors. If the function returns m outputs, rewrite will return a list of |
| 1633 | m tensors. |
| 1634 | inputs: A list of input tensors or `None` (equivalent to an empty list). |
| 1635 | infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple |
| 1636 | of arguments as inputs to `computation`. |
| 1637 | device_assignment: if not `None`, a `DeviceAssignment` describing the |
| 1638 | mapping between logical cores in the computation with physical cores in |
| 1639 | the TPU topology. May be omitted for a single-core computation, in which |
| 1640 | case the core attached to task 0, TPU device 0 is used. |
| 1641 | name: The name of the operator. |
| 1642 | Returns: |
| 1643 | A list of output tensors. |
| 1644 | """ |
| 1645 | |
| 1646 | def guarantee_const_getter(getter, name, *args, **kwargs): |
| 1647 | with ops.control_dependencies(None): |
| 1648 | return array_ops.guarantee_const( |
| 1649 | getter(name, *args, **kwargs), name=name + "/GuaranteeConst") |
| 1650 | |
| 1651 | def wrapped_computation(*args, **kwargs): |
| 1652 | """Execute computation under `_TPUInferenceContext`.""" |
| 1653 | context = _TPUInferenceContext( |
| 1654 | name=ops.get_default_graph().unique_name("rewrite_for_inference")) |
| 1655 | try: |
| 1656 | context.Enter() |
| 1657 | |
| 1658 | vscope = variable_scope.get_variable_scope() |
| 1659 | prev_custom_getter = vscope.custom_getter |
| 1660 | prev_caching_device = vscope.caching_device |
| 1661 | vscope.set_custom_getter(guarantee_const_getter) |
| 1662 | vscope.set_caching_device(lambda op: op.device) |
| 1663 | |
| 1664 | result = computation(*args, **kwargs) |
| 1665 | |
| 1666 | vscope.set_custom_getter(prev_custom_getter) |
| 1667 | vscope.set_caching_device(prev_caching_device) |
| 1668 | finally: |
| 1669 | context.Exit() |
| 1670 | return result |