If necessary, request deferred restorations of variables.
(network, save_path, map_func, user_map_func,
existing_variables_by_checkpoint_name)
| 928 | |
| 929 | |
| 930 | def _set_restore_on_create(network, save_path, map_func, user_map_func, |
| 931 | existing_variables_by_checkpoint_name): |
| 932 | """If necessary, request deferred restorations of variables.""" |
| 933 | checkpoint_reader = checkpoint_utils.load_checkpoint(save_path) |
| 934 | checkpointed_variables_to_restore = {} |
| 935 | for checkpoint_name, _ in checkpoint_utils.list_variables(save_path): |
| 936 | if checkpoint_name in existing_variables_by_checkpoint_name: |
| 937 | # This variable was already created and restored. |
| 938 | continue |
| 939 | # Save the variable for later restoration in a custom getter. |
| 940 | checkpointed_variables_to_restore[checkpoint_name] = ( |
| 941 | checkpoint_reader.get_tensor(checkpoint_name)) |
| 942 | # Only set a deferred restoration if there are checkpoint variables which |
| 943 | # have not been assigned to existing variables. Note that this loses out on |
| 944 | # some opportunity for error checking, but avoids creating |
| 945 | # _DeferredRestoration objects once a Network has been built (so that |
| 946 | # restoring in a loop does not take increasing amounts of memory). |
| 947 | if checkpointed_variables_to_restore: |
| 948 | if context.executing_eagerly(): |
| 949 | sess = None |
| 950 | else: |
| 951 | sess = ops.get_default_session() |
| 952 | # We need a name for error messages. If we haven't been added to another |
| 953 | # Network yet, we're top-level. |
| 954 | network._finalize_name(False) |
| 955 | network._set_scope() |
| 956 | # Save a record of this restoration for use in the custom getter. |
| 957 | deferred_restoration = _DeferredRestoration( |
| 958 | map_func=map_func, |
| 959 | map_func_is_user=(user_map_func is not None), |
| 960 | checkpointed_variables_to_restore=checkpointed_variables_to_restore, |
| 961 | restored_variables={}, |
| 962 | session=sess, |
| 963 | network_name=network.name, |
| 964 | network_scope_name=network.scope_name) |
| 965 | # Add the deferred registration to non-Network children, and request that |
| 966 | # Networks propagate the request to their children. |
| 967 | _add_deferred_restoration(network, deferred_restoration) |
| 968 | |
| 969 | |
| 970 | @deprecation.deprecated( |
no test coverage detected