Save variables from the Network to a checkpoint. Args: network: A Network object to save. save_path: Either a checkpoint prefix or the name of a directory to save the checkpoint in (in which case the checkpoint will be named based on the Network name). global_step: The glo
(network, save_path, global_step=None,
map_func=None)
| 766 | "Please inherit from tf.keras.Model instead of tfe.Network, and use " |
| 767 | "tf.keras.Model.save_weights.")) |
| 768 | def save_network_checkpoint(network, save_path, global_step=None, |
| 769 | map_func=None): |
| 770 | """Save variables from the Network to a checkpoint. |
| 771 | |
| 772 | Args: |
| 773 | network: A Network object to save. |
| 774 | save_path: Either a checkpoint prefix or the name of a directory to save the |
| 775 | checkpoint in (in which case the checkpoint will be named based on the |
| 776 | Network name). |
| 777 | global_step: The global step to use when naming the checkpoint. If None |
| 778 | (default), we will first try to get the default global step. If that fails |
| 779 | because no default global step exists, then the checkpoint is created |
| 780 | without a global step suffix. |
| 781 | map_func: A function mapping fully qualified variable names (e.g. |
| 782 | 'my_network_1/dense_1/kernel') to names in the checkpoint. By default (if |
| 783 | `map_func=None`), the variable prefix for the network being restored |
| 784 | (`Network.scope_name + '/'`, e.g. 'my_network_1/') is stripped and all |
| 785 | other variable names (shared with other Networks) are left unchanged. |
| 786 | |
| 787 | Returns: |
| 788 | The checkpoint prefix for the saved checkpoint, which may be passed to |
| 789 | `Network.restore`. |
| 790 | Raises: |
| 791 | ValueError: If the Network has not yet been called, or if map_func results |
| 792 | in a name collision. |
| 793 | """ |
| 794 | if not network.built: |
| 795 | raise ValueError( |
| 796 | "Attempt to save the Network before it was first called. This means " |
| 797 | "variables have not yet been created, so there is nothing to save.") |
| 798 | network._set_scope() # scope_name should be available to map_funcs |
| 799 | if global_step is None: |
| 800 | global_step = training_util.get_global_step() |
| 801 | if os.path.isdir(save_path): |
| 802 | # If we were passed a directory, default to naming based on the Network |
| 803 | # name. |
| 804 | save_path = os.path.join(save_path, network.name.replace("/", "_")) |
| 805 | user_map_func = map_func |
| 806 | if map_func is None: |
| 807 | map_func = _make_prefix_stripping_map_fn(network.scope_name) |
| 808 | variable_map = {} |
| 809 | for variable in network.variables: |
| 810 | mapped_name = map_func(variable._shared_name) |
| 811 | if variable_map.setdefault(mapped_name, variable) is not variable: |
| 812 | if user_map_func is None: |
| 813 | # Instead of erroring out, we could just re-try and silently use the |
| 814 | # full variable names in the checkpoint. This could be odd for deeply |
| 815 | # nested sub-Networks (since the full prefix from the nesting would |
| 816 | # get added), so for now we'll let the user deal with this case. |
| 817 | raise ValueError( |
| 818 | _default_naming_conflict_error_message( |
| 819 | mapped_name=mapped_name, |
| 820 | first_variable=variable_map[mapped_name], |
| 821 | second_variable=variable, |
| 822 | network_name=network.name, |
| 823 | network_scope_name=network.scope_name)) |
| 824 | else: |
| 825 | # The user passed their own problematic map_func. |
nothing calls this directly
no test coverage detected