Decorator for asserting that no new Tensors persist after a test. Mainly useful for checking that code using the Python C API has correctly manipulated reference counts. Clears the caches that it knows about, runs the garbage collector, then checks that there are no Tensor or Tensor-like o
(f)
| 627 | |
| 628 | |
| 629 | def assert_no_new_tensors(f): |
| 630 | """Decorator for asserting that no new Tensors persist after a test. |
| 631 | |
| 632 | Mainly useful for checking that code using the Python C API has correctly |
| 633 | manipulated reference counts. |
| 634 | |
| 635 | Clears the caches that it knows about, runs the garbage collector, then checks |
| 636 | that there are no Tensor or Tensor-like objects still around. This includes |
| 637 | Tensors to which something still has a reference (e.g. from missing |
| 638 | Py_DECREFs) and uncollectable cycles (i.e. Python reference cycles where one |
| 639 | of the objects has __del__ defined). |
| 640 | |
| 641 | Args: |
| 642 | f: The test case to run. |
| 643 | |
| 644 | Returns: |
| 645 | The decorated test case. |
| 646 | """ |
| 647 | |
| 648 | def decorator(self, **kwargs): |
| 649 | """Finds existing Tensors, runs the test, checks for new Tensors.""" |
| 650 | |
| 651 | def _is_tensorflow_object(obj): |
| 652 | try: |
| 653 | return isinstance(obj, |
| 654 | (ops.Tensor, variables.Variable, |
| 655 | tensor_shape.Dimension, tensor_shape.TensorShape)) |
| 656 | except ReferenceError: |
| 657 | # If the object no longer exists, we don't care about it. |
| 658 | return False |
| 659 | |
| 660 | tensors_before = set( |
| 661 | id(obj) for obj in gc.get_objects() if _is_tensorflow_object(obj)) |
| 662 | outside_executed_eagerly = context.executing_eagerly() |
| 663 | # Run the test in a new graph so that collections get cleared when it's |
| 664 | # done, but inherit the graph key so optimizers behave. |
| 665 | outside_graph_key = ops.get_default_graph()._graph_key |
| 666 | with ops.Graph().as_default(): |
| 667 | ops.get_default_graph()._graph_key = outside_graph_key |
| 668 | if outside_executed_eagerly: |
| 669 | with context.eager_mode(): |
| 670 | result = f(self, **kwargs) |
| 671 | else: |
| 672 | result = f(self, **kwargs) |
| 673 | # Make an effort to clear caches, which would otherwise look like leaked |
| 674 | # Tensors. |
| 675 | context.context()._clear_caches() # pylint: disable=protected-access |
| 676 | gc.collect() |
| 677 | tensors_after = [ |
| 678 | obj for obj in gc.get_objects() |
| 679 | if _is_tensorflow_object(obj) and id(obj) not in tensors_before |
| 680 | ] |
| 681 | if tensors_after: |
| 682 | raise AssertionError(("%d Tensors not deallocated after test: %s" % ( |
| 683 | len(tensors_after), |
| 684 | str(tensors_after), |
| 685 | ))) |
| 686 | return result |