mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 08:02:27 +00:00
19 lines
760 B
Python
19 lines
760 B
Python
import mxnet as mx
|
|
|
|
|
|
def save_checkpoint(prefix, epoch, arg_params, aux_params):
|
|
"""Checkpoint the model data into file.
|
|
:param prefix: Prefix of model name.
|
|
:param epoch: The epoch number of the model.
|
|
:param arg_params: dict of str to NDArray
|
|
Model parameter, dict of name to NDArray of net's weights.
|
|
:param aux_params: dict of str to NDArray
|
|
Model parameter, dict of name to NDArray of net's auxiliary states.
|
|
:return: None
|
|
prefix-epoch.params will be saved for parameters.
|
|
"""
|
|
save_dict = {('arg:%s' % k): v for k, v in arg_params.items()}
|
|
save_dict.update({('aux:%s' % k): v for k, v in aux_params.items()})
|
|
param_name = '%s-%04d.params' % (prefix, epoch)
|
|
mx.nd.save(param_name, save_dict)
|