mindspore.load_checkpoint

mindspore.load_checkpoint(ckpt_file_name, net=None, strict_load=False, filter_prefix=None, dec_key=None, dec_mode='AES-GCM', specify_prefix=None, choice_func=None, crc_check=False, remove_redundancy=False, format='ckpt')[source]

Load checkpoint info from a specified file.

Note

  • specify_prefix and filter_prefix are in the process of being deprecated, choice_func is recommended instead. specify_prefix and filter_prefix do not affect each other. And using either of those two args will override choice_func at the same time.

  • If none of the parameters are loaded from checkpoint file, it will throw ValueError.

  • When loading a checkpoint that has removed redundancy, the network should be compiled.

  • When net is not None, it will verify whether the remove_redundancy parameter matches the deduplication flag in the loaded safetensors file. If they are different, load the file according to the deduplication flag in the file.

  • The ckpt format stores checkpoint protobuf content as a single message that is read into memory before parsing. If the protobuf content is larger than current available memory, loading fails early to avoid deterministic out-of-memory failure. For large model weights, consider using safetensors or segmented checkpoints.

Parameters:

  • ckpt_file_name (Union[str, io.BytesIO]) – Checkpoint file name or BytesIO object (only support for safetensors format).

  • net (Cell, optional) – The network where the parameters will be loaded. Default: None .

  • strict_load (bool, optional) – Whether to strictly load the parameter into net. If False , it will load parameter into net when parameter name's suffix in checkpoint file is the same as the parameter in the network. When the types are inconsistent, perform type conversion on the parameters of the same type, such as float32 to float16. Default: False .

  • filter_prefix (Union[str, list[str], tuple[str]], optional) – Deprecated(see choice_func). Parameters starting with the filter_prefix will not be loaded. Default: None .

  • dec_key (Union[None, bytes], optional) – Byte type key used for decryption. If the value is None , the decryption is not required. Default: None .

  • dec_mode (str, optional) – This parameter is valid only when dec_key is not set to None . Specifies the decryption mode, currently supports "AES-GCM" , "AES-CBC" and "SM4-CBC" . Default: "AES-GCM" .

  • specify_prefix (Union[str, list[str], tuple[str]], optional) – Deprecated(see choice_func). Parameters starting with the specify_prefix will be loaded. Default: None .

  • choice_func (Union[None, function], optional) – Input value of the function is a Parameter name of type string, and the return value is a bool. If it returns True , the Parameter that matches the custom condition will be loaded. If returns False , the Parameter that matches the custom condition will be removed. Default: None .

  • crc_check (bool, optional) – Whether to perform CRC32 corruption check when loading checkpoint. CRC32 is not a cryptographically secure integrity check and must not be used to authenticate checkpoint files from untrusted sources. For tamper resistance, use an authenticated protection mechanism such as AES-GCM. Default: False .

  • remove_redundancy (bool, optional) – Whether to enable loading of checkpoint saved with redundancy removal. Redundancy removal refers to eliminating redundant data in data parallelism mode. Default: False , means redundant-free loading is not enabled.

  • format (str, optional) – Format of the input file, can be "ckpt" or "safetensors". Default: "ckpt" .

Returns:

Dict, key is parameter name, value is a Parameter or string. When the append_dict parameter of mindspore.save_checkpoint() and the append_info parameter of mindspore.train.CheckpointConfig are used to save the checkpoint, append_dict and append_info are dict types, and their values are strings, then the return value obtained by loading checkpoint is string, and in other cases the return value is Parameter.

Raises:

  • ValueError – Checkpoint file's format is incorrect.

  • ValueError – Parameter's dict is None after load checkpoint file.

  • TypeError – The type of specify_prefix or filter_prefix is incorrect.

Examples

>>> import mindspore as ms
>>>
>>> ckpt_file_name = "./checkpoint/LeNet5-1_32.ckpt"
>>> param_dict = ms.load_checkpoint(ckpt_file_name,
...                                 choice_func=lambda x: x.startswith("conv") and not x.startswith("conv1"))
>>> print(param_dict["conv2.weight"])
Parameter (name=conv2.weight, shape=(16, 6, 5, 5), dtype=Float32, requires_grad=True)
>>> def func(param_name):
...     whether_load = False
...     if param_name.startswith("conv"):
...         whether_load = True
...     if param_name.startswith("conv1"):
...         whether_load = False
...     return whether_load
>>> param_dict1 = ms.load_checkpoint(ckpt_file_name, choice_func=func)
>>> print(param_dict1["conv2.weight"])
Parameter (name=conv2.weight, shape=(16, 6, 5, 5), dtype=Float32, requires_grad=True)
>>> def func(param_name):
...     whether_load = False
...     if param_name.startswith("conv1"):
...         whether_load = True
...     return whether_load
>>> param_dict2 = ms.load_checkpoint(ckpt_file_name, choice_func=func)
>>> print(param_dict2)
{'conv1.weight': Parameter (name=conv1.weight, shape=(6, 1, 5, 5), dtype=Float32, requires_grad=True)}
Tutorial Examples: