A custom getter which processes deferred restorations.
(getter,
name,
shape=None,
dtype=None,
initializer=None,
*args,
**kwargs)
| 652 | deferred_restorations = [] |
| 653 | |
| 654 | def _custom_getter(getter, |
| 655 | name, |
| 656 | shape=None, |
| 657 | dtype=None, |
| 658 | initializer=None, |
| 659 | *args, |
| 660 | **kwargs): |
| 661 | """A custom getter which processes deferred restorations.""" |
| 662 | # Iterate over restorations, newest first (newer restorations will take |
| 663 | # precedence over older restorations, just like with immediate restorations |
| 664 | # into existing variables). |
| 665 | delayed_restoration = None |
| 666 | found_value = False |
| 667 | value_to_restore = None |
| 668 | for delayed_restoration in reversed(deferred_restorations): |
| 669 | checkpoint_name = delayed_restoration.map_func(name) |
| 670 | if (checkpoint_name in |
| 671 | delayed_restoration.checkpointed_variables_to_restore): |
| 672 | found_value = True |
| 673 | value_to_restore = ( |
| 674 | delayed_restoration |
| 675 | .checkpointed_variables_to_restore[checkpoint_name]) |
| 676 | if found_value: |
| 677 | break |
| 678 | # value_to_restore may be False because this variable is not in any |
| 679 | # checkpoint we are restoring, or None because we have explicitly set it to |
| 680 | # None when it was previously fetched. In either case, we don't need to |
| 681 | # set an initializer. |
| 682 | if found_value and value_to_restore is not None: |
| 683 | initializer = value_to_restore |
| 684 | shape = None |
| 685 | variable = getter( |
| 686 | name, |
| 687 | shape=shape, |
| 688 | dtype=dtype, |
| 689 | initializer=initializer, |
| 690 | *args, |
| 691 | **kwargs) |
| 692 | if found_value and value_to_restore is not None: |
| 693 | # Mark as already restored from this checkpoint. |
| 694 | delayed_restoration.checkpointed_variables_to_restore[ |
| 695 | checkpoint_name] = None |
| 696 | if not context.executing_eagerly(): |
| 697 | delayed_restoration.session.run(variable.initializer) |
| 698 | if found_value: |
| 699 | # Error checking should run even if we've already restored a value. |
| 700 | if delayed_restoration.restored_variables.setdefault( |
| 701 | checkpoint_name, variable) is not variable: |
| 702 | # Naming conflict. We've tried to initialize two variables with the |
| 703 | # same value from the checkpoint. |
| 704 | if delayed_restoration.map_func_is_user: |
| 705 | raise ValueError( |
| 706 | _restore_custom_map_func_error_message( |
| 707 | mapped_name=checkpoint_name, |
| 708 | first_variable=delayed_restoration |
| 709 | .restored_variables[checkpoint_name], |
| 710 | second_variable=variable, |
| 711 | network_name=delayed_restoration.network_name, |
nothing calls this directly
no test coverage detected