Verifies Python loops for TF-specific limits.
| 773 | |
| 774 | |
| 775 | class _PythonLoopChecker(object): |
| 776 | """Verifies Python loops for TF-specific limits.""" |
| 777 | |
| 778 | def __init__(self): |
| 779 | self.iterations = 0 |
| 780 | self.check_inefficient_unroll = WARN_INEFFICIENT_UNROLL |
| 781 | |
| 782 | # Triggered when we decided to test the op counts. |
| 783 | self.check_op_count_after_iteration = False |
| 784 | |
| 785 | def _get_ops(self): |
| 786 | return ops.get_default_graph().get_operations() |
| 787 | |
| 788 | def _check_unroll_limits(self): |
| 789 | if LIMIT_PYTHON_ITERATIONS and self.iterations > PYTHON_MAX_ITERATIONS: |
| 790 | raise ValueError('iteration limit exceeded') |
| 791 | |
| 792 | def _stop_checking_inefficient_unroll(self): |
| 793 | self.check_inefficient_unroll = False |
| 794 | self.ops_before_iteration = None |
| 795 | |
| 796 | def _verify_ineffcient_unroll(self): |
| 797 | """Checks for possibly-inefficient creation of ops in a Python loop.""" |
| 798 | assert self.ops_before_iteration is not None |
| 799 | ops_after_iteration = self._get_ops() |
| 800 | new_ops = tuple( |
| 801 | op for op in ops_after_iteration if op not in self.ops_before_iteration) |
| 802 | |
| 803 | if len(new_ops) < INEFFICIENT_UNROLL_MIN_OPS: |
| 804 | return False |
| 805 | |
| 806 | # TODO(mdan): Add location information. |
| 807 | ag_logging.warn( |
| 808 | 'TensorFlow ops are being created in a Python loop with large number' |
| 809 | ' of iterations. This can lead to slow startup. Did you mean to use a' |
| 810 | ' TensorFlow loop? For example, `while True:` is a Python loop, and' |
| 811 | ' `while tf.constant(True):` is a TensorFlow loop. The following' |
| 812 | ' ops were created after iteration %s: %s', self.iterations, new_ops) |
| 813 | return True |
| 814 | |
| 815 | def before_iteration(self): |
| 816 | """Called before each iteration in a Python loop.""" |
| 817 | if (self.check_inefficient_unroll and |
| 818 | self.iterations > INEFFICIENT_UNROLL_MIN_ITERATIONS): |
| 819 | self.ops_before_iteration = self._get_ops() |
| 820 | self.check_op_count_after_iteration = True |
| 821 | |
| 822 | def after_iteration(self): |
| 823 | """Called after each iteration in a Python loop.""" |
| 824 | self.iterations += 1 |
| 825 | |
| 826 | self._check_unroll_limits() |
| 827 | |
| 828 | if self.check_inefficient_unroll and self.check_op_count_after_iteration: |
| 829 | did_warn = self._verify_ineffcient_unroll() |
| 830 | if did_warn: |
| 831 | self._stop_checking_inefficient_unroll() # Only warn once. |
| 832 | elif self.iterations > INEFFICIENT_UNROLL_MIN_ITERATIONS + 3: |