a
    dȈ                     @   s  d Z ddlZddlZddlZddlZddlZddlmZ ddlm	Z	 ddl
mZmZmZmZ ddlmZ ddlZddlZddlZddlmZ ddlZddlmZ dd	lmZ dd
lmZ ddlmZ ddl m!Z!m"Z"m#Z# ddl$m%Z% e&e'Z(e# Z)G dd deZ*dS )za
Model Checkpointing
===================

Automatically save model checkpoints during training.

    N)deepcopy)	timedelta)AnyDictOptionalSet)proxy)Tensor)get_filesystem)_PATH)
Checkpoint)MisconfigurationException)rank_zero_inforank_zero_warnWarningCache)STEP_OUTPUTc                       sD  e Zd ZdZdZdZdZdZdcee	 ee
 ee
 eee eee
eee ee ee ee d
 fddZee
dddZdde
ddddZddddddZddeeeddddZddddddZddddddZee
ef ddd Zee
ef dd!d"d#Zdee
ef dd$d%d&Zde
dd'd(d)Zded*d+d,Zded*d-d.Zddd/d0Z ee	 ee
 dd1d2d3Z!e
dd4d5d6Z"ee ee ee dd7d8d9Z#eee dd:d;Z$dddee ed<d=d>Z%e&deee
 ee
ef e
ee
d@dAdBZ'dfee
ef ee
 ee e
dCdDdEZ(de	d*dFdGZ)de*e
 d*dHdIZ+e	ddJdKdLZ,dgee
ef dee
 e
dMdNdOZ-dee
ef d*dPdQZ.dee
ef dd$dRdSZ/dee
ef dd$dTdUZ0dee
ef dd$dVdWZ1edee
ef ddXdYdZZ2dhee	 dd[d\d]Z3e	ded^d_d`Z4de
dd'dadbZ5  Z6S )iModelCheckpointa,"  
    Save the model periodically by monitoring a quantity. Every metric logged with
    :meth:`~pytorch_lightning.core.module.log` or :meth:`~pytorch_lightning.core.module.log_dict` in
    LightningModule is a candidate for the monitor key. For more information, see
    :ref:`checkpointing`.

    After training finishes, use :attr:`best_model_path` to retrieve the path to the
    best checkpoint file and :attr:`best_model_score` to retrieve its score.

    Args:
        dirpath: directory to save the model file.

            Example::

                # custom path
                # saves a file like: my/path/epoch=0-step=10.ckpt
                >>> checkpoint_callback = ModelCheckpoint(dirpath='my/path/')

            By default, dirpath is ``None`` and will be set at runtime to the location
            specified by :class:`~pytorch_lightning.trainer.trainer.Trainer`'s
            :paramref:`~pytorch_lightning.trainer.trainer.Trainer.default_root_dir` argument,
            and if the Trainer uses a logger, the path will also contain logger name and version.

        filename: checkpoint filename. Can contain named formatting options to be auto-filled.

            Example::

                # save any arbitrary metrics like `val_loss`, etc. in name
                # saves a file like: my/path/epoch=2-val_loss=0.02-other_metric=0.03.ckpt
                >>> checkpoint_callback = ModelCheckpoint(
                ...     dirpath='my/path',
                ...     filename='{epoch}-{val_loss:.2f}-{other_metric:.2f}'
                ... )

            By default, filename is ``None`` and will be set to ``'{epoch}-{step}'``, where "epoch" and "step" match
            the number of finished epoch and optimizer steps respectively.
        monitor: quantity to monitor. By default it is ``None`` which saves a checkpoint only for the last epoch.
        verbose: verbosity mode. Default: ``False``.
        save_last: When ``True``, saves an exact copy of the checkpoint to a file `last.ckpt` whenever a checkpoint
            file gets saved. This allows accessing the latest checkpoint in a deterministic manner. Default: ``None``.
        save_top_k: if ``save_top_k == k``,
            the best k models according to the quantity monitored will be saved.
            if ``save_top_k == 0``, no models are saved.
            if ``save_top_k == -1``, all models are saved.
            Please note that the monitors are checked every ``every_n_epochs`` epochs.
            if ``save_top_k >= 2`` and the callback is called multiple
            times inside an epoch, the name of the saved file will be
            appended with a version count starting with ``v1``.
        mode: one of {min, max}.
            If ``save_top_k != 0``, the decision to overwrite the current save file is made
            based on either the maximization or the minimization of the monitored quantity.
            For ``'val_acc'``, this should be ``'max'``, for ``'val_loss'`` this should be ``'min'``, etc.
        auto_insert_metric_name: When ``True``, the checkpoints filenames will contain the metric name.
            For example, ``filename='checkpoint_{epoch:02d}-{acc:02.0f}`` with epoch ``1`` and acc ``1.12`` will resolve
            to ``checkpoint_epoch=01-acc=01.ckpt``. Is useful to set it to ``False`` when metric names contain ``/``
            as this will result in extra folders.
            For example, ``filename='epoch={epoch}-step={step}-val_acc={val/acc:.2f}', auto_insert_metric_name=False``
        save_weights_only: if ``True``, then only the model's weights will be
            saved. Otherwise, the optimizer states, lr-scheduler states, etc are added in the checkpoint too.
        every_n_train_steps: Number of training steps between checkpoints.
            If ``every_n_train_steps == None or every_n_train_steps == 0``, we skip saving during training.
            To disable, set ``every_n_train_steps = 0``. This value must be ``None`` or non-negative.
            This must be mutually exclusive with ``train_time_interval`` and ``every_n_epochs``.
        train_time_interval: Checkpoints are monitored at the specified time interval.
            For all practical purposes, this cannot be smaller than the amount
            of time it takes to process a single training batch. This is not
            guaranteed to execute at the exact time specified, but should be close.
            This must be mutually exclusive with ``every_n_train_steps`` and ``every_n_epochs``.
        every_n_epochs: Number of epochs between checkpoints.
            This value must be ``None`` or non-negative.
            To disable saving top-k checkpoints, set ``every_n_epochs = 0``.
            This argument does not impact the saving of ``save_last=True`` checkpoints.
            If all of ``every_n_epochs``, ``every_n_train_steps`` and
            ``train_time_interval`` are ``None``, we save a checkpoint at the end of every epoch
            (equivalent to ``every_n_epochs = 1``).
            If ``every_n_epochs == None`` and either ``every_n_train_steps != None`` or ``train_time_interval != None``,
            saving at the end of each epoch is disabled
            (equivalent to ``every_n_epochs = 0``).
            This must be mutually exclusive with ``every_n_train_steps`` and ``train_time_interval``.
            Setting both ``ModelCheckpoint(..., every_n_epochs=V, save_on_train_epoch_end=False)`` and
            ``Trainer(max_epochs=N, check_val_every_n_epoch=M)``
            will only save checkpoints at epochs 0 < E <= N
            where both values for ``every_n_epochs`` and ``check_val_every_n_epoch`` evenly divide E.
        save_on_train_epoch_end: Whether to run checkpointing at the end of the training epoch.
            If this is ``False``, then the check runs at the end of the validation.

    Note:
        For extra customization, ModelCheckpoint includes the following attributes:

        - ``CHECKPOINT_JOIN_CHAR = "-"``
        - ``CHECKPOINT_NAME_LAST = "last"``
        - ``FILE_EXTENSION = ".ckpt"``
        - ``STARTING_VERSION = 1``

        For example, you can change the default last checkpoint name by doing
        ``checkpoint_callback.CHECKPOINT_NAME_LAST = "{epoch}-last"``

        If you want to checkpoint every N hours, every M train batches, and/or every K val epochs,
        then you should create multiple ``ModelCheckpoint`` callbacks.

        If the checkpoint's ``dirpath`` changed from what it was before while resuming the training,
        only ``best_model_path`` will be reloaded and a warning will be issued.

    Raises:
        MisconfigurationException:
            If ``save_top_k`` is smaller than ``-1``,
            if ``monitor`` is ``None`` and ``save_top_k`` is none of ``None``, ``-1``, and ``0``, or
            if ``mode`` is none of ``"min"`` or ``"max"``.
        ValueError:
            If ``trainer.save_checkpoint`` is ``None``.

    Example::

        >>> from pytorch_lightning import Trainer
        >>> from pytorch_lightning.callbacks import ModelCheckpoint

        # saves checkpoints to 'my/path/' at every epoch
        >>> checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
        >>> trainer = Trainer(callbacks=[checkpoint_callback])

        # save epoch and val_loss in name
        # saves a file like: my/path/sample-mnist-epoch=02-val_loss=0.32.ckpt
        >>> checkpoint_callback = ModelCheckpoint(
        ...     monitor='val_loss',
        ...     dirpath='my/path/',
        ...     filename='sample-mnist-{epoch:02d}-{val_loss:.2f}'
        ... )

        # save epoch and val_loss in name, but specify the formatting yourself (e.g. to avoid problems with Tensorboard
        # or Neptune, due to the presence of characters like '=' or '/')
        # saves a file like: my/path/sample-mnist-epoch02-val_loss0.32.ckpt
        >>> checkpoint_callback = ModelCheckpoint(
        ...     monitor='val/loss',
        ...     dirpath='my/path/',
        ...     filename='sample-mnist-epoch{epoch:02d}-val_loss{val/loss:.2f}',
        ...     auto_insert_metric_name=False
        ... )

        # retrieve the best checkpoint after training
        checkpoint_callback = ModelCheckpoint(dirpath='my/path/')
        trainer = Trainer(callbacks=[checkpoint_callback])
        model = ...
        trainer.fit(model)
        checkpoint_callback.best_model_path

    .. tip:: Saving and restoring multiple checkpoint callbacks at the same time is supported under variation in the
        following arguments:

        *monitor, mode, every_n_train_steps, every_n_epochs, train_time_interval*

        Read more: :ref:`Persisting Callback State <extensions/callbacks_state:save callback state>`
    -lastz.ckpt   NFminT)dirpathfilenamemonitorverbose	save_last
save_top_ksave_weights_onlymodeauto_insert_metric_nameevery_n_train_stepstrain_time_intervalevery_n_epochssave_on_train_epoch_endc                    s   t    || _|| _|| _|| _|| _|	| _|| _d| _	d | _
d | _i | _d| _d | _d| _d| _|  |  | | | || | |
|| |   d S )Nr    )super__init__r   r   r   r   r   r   _save_on_train_epoch_end_last_global_step_saved_last_time_checkedcurrent_scorebest_k_modelskth_best_model_pathbest_model_scorebest_model_pathlast_model_path#_ModelCheckpoint__init_monitor_mode_ModelCheckpoint__init_ckpt_dir_ModelCheckpoint__init_triggers-_ModelCheckpoint__validate_init_configuration)selfr   r   r   r   r   r   r   r   r   r    r!   r"   r#   	__class__ u/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/pytorch_lightning/callbacks/model_checkpoint.pyr&      s,    

zModelCheckpoint.__init__)returnc                 C   s   | j | j| j| j| j| jdS )N)r   r   r    r"   r!   )Z_generate_state_keyr   r   _every_n_train_steps_every_n_epochs_train_time_intervalr4   r7   r7   r8   	state_key   s    zModelCheckpoint.state_keyz
pl.Trainerzpl.LightningModule)trainer	pl_modulestager9   c                 C   s:   |  |}|j|}|| _|jr6|dkr6| | j d S )NZfit)"_ModelCheckpoint__resolve_ckpt_dirstrategy	broadcastr   is_global_zero'_ModelCheckpoint__warn_if_dir_not_empty)r4   r?   r@   rA   r   r7   r7   r8   setup  s
    
zModelCheckpoint.setup)r?   r@   r9   c                 C   s   t  | _d S N)time	monotonicr)   )r4   r?   r@   r7   r7   r8   on_train_start  s    zModelCheckpoint.on_train_start)r?   r@   outputsbatch	batch_idxr9   c                 C   s   |  |rdS | jdk p&|j| j dk}| j}d}t }	|rh| j}
|
du pZ|	|
 | k }|j	|}|rt|rtdS |s~|	| _| 
|}| || | || dS )zTSave checkpoint on train batch end if we meet the criteria for `every_n_train_steps`Nr   r   T)_should_skip_saving_checkpointr:   global_stepr<   rI   rJ   r)   total_secondsrC   rD   _monitor_candidates_save_topk_checkpoint_save_last_checkpoint)r4   r?   r@   rL   rM   rN   Z
skip_batchr!   Z	skip_timenowZprev_time_checkmonitor_candidatesr7   r7   r8   on_train_batch_end  s"    	

z"ModelCheckpoint.on_train_batch_endc                 C   sX   |  |sT| |rT| |}| jdkrH|jd | j dkrH| || | || dS )z3Save a checkpoint at the end of the training epoch.r   r   NrO   _should_save_on_train_epoch_endrR   r;   current_epochrS   rT   r4   r?   r@   rV   r7   r7   r8   on_train_epoch_end+  s
    
z"ModelCheckpoint.on_train_epoch_endc                 C   sX   |  |sT| |sT| |}| jdkrH|jd | j dkrH| || | || dS )z5Save a checkpoint at the end of the validation stage.r   r   NrX   r[   r7   r7   r8   on_validation_end3  s
    
z!ModelCheckpoint.on_validation_endc              
   C   s*   | j | j| j| j| j| j| j| j| jd	S )N	r   r-   r.   r*   r   r+   r,   	kth_valuer/   r^   r=   r7   r7   r8   
state_dict;  s    zModelCheckpoint.state_dict)r`   r9   c                 C   s   | d| j}| j|krd|d | _| d| j| _| d| j| _| d| j| _| d| j| _ntd|d| jd	 |d
 | _	d S )Nr   r-   r,   r_   r+   r/   zThe dirpath has changed from z to z, therefore `best_model_score`, `kth_best_model_path`, `kth_value`, `last_model_path` and `best_k_models` won't be reloaded. Only `best_model_path` will be reloaded.r.   )
getr   r-   r,   r_   r+   r/   warningswarnr.   )r4   r`   Zdirpath_from_ckptr7   r7   r8   load_state_dictH  s    

zModelCheckpoint.load_state_dict)r?   rV   r9   c                 C   s~   | j dkrd S | jd urn| j|vr`d| jdt| d| jd}|jjjjrVt|t	| | 
|| n| || d S )Nr   z`ModelCheckpoint(monitor=z=)` could not find the monitored key in the returned metrics: z. HINT: Did you call `log(z#, value)` in the `LightningModule`?)r   r   listZfit_loopZ
epoch_loopZval_loopZ_has_runr   warning_cacherc   _save_monitor_checkpoint_save_none_monitor_checkpoint)r4   r?   rV   mr7   r7   r8   rS   Z  s    




z%ModelCheckpoint._save_topk_checkpoint)r?   filepathr9   c                 C   s:   | || j |j| _|jr6|jD ]}|t|  q"d S rH   )Zsave_checkpointr   rP   r(   rE   loggersZafter_save_checkpointr   )r4   r?   rj   loggerr7   r7   r8   _save_checkpointm  s
    
z ModelCheckpoint._save_checkpoint)r?   r9   c                 C   s6   ddl m} t|jp4|jj|jkp4|jp4| j|j	kS )Nr   )	TrainerFn)
Z pytorch_lightning.trainer.statesrn   boolZfast_dev_runstatefnZFITTINGZsanity_checkingr(   rP   )r4   r?   rn   r7   r7   r8   rO   w  s    

z.ModelCheckpoint._should_skip_saving_checkpointc                 C   s:   | j d ur| j S |jdkrdS t|jdkr0dS |jdkS )Nr   Fr   Tg      ?)r'   Zcheck_val_every_n_epochsumZnum_val_batchesZval_check_interval)r4   r?   r7   r7   r8   rY     s    

z/ModelCheckpoint._should_save_on_train_epoch_endc                 C   s   | j dk rtd| j  d| jdk r8td| j d| jdk rTtd| j d| jdk}| jdk}| jd u}|| | dkrtd	| j d
| j d| j d| jd u r| j dvrtd| j  d| j dkr| jrtd d S )NzInvalid value for save_top_k=z. Must be >= -1r   z&Invalid value for every_n_train_steps=z. Must be >= 0z!Invalid value for every_n_epochs=r   z.Combination of parameters every_n_train_steps=z, every_n_epochs=z and train_time_interval=z should be mutually exclusive.)rs   r   r   zModelCheckpoint(save_top_k=zM, monitor=None) is not a valid configuration. No quantity for top_k to track.zfModelCheckpoint(save_last=True, save_top_k=-1, monitor=None) will duplicate the last checkpoint saved.)r   r   r:   r;   r<   r   r   r   )r4   Zevery_n_train_steps_triggeredZevery_n_epochs_triggeredZtrain_time_interval_triggeredr7   r7   r8   Z__validate_init_configuration  s8    








z-ModelCheckpoint.__validate_init_configuration)r   r   r9   c                 C   s>   t |r
|nd| _|r.| jjdkr.tj|}|| _|| _d S )Nr$   file)r
   _fsprotocolospathrealpathr   r   )r4   r   r   r7   r7   r8   Z__init_ckpt_dir  s
    zModelCheckpoint.__init_ckpt_dir)r   r9   c                 C   sZ   t tj}|df| dfd}||vrFtdd|  d| || \| _| _d S )Nr   maxr   rz   z`mode` can be z, z	 but got )	torchtensornpZInfr   joinkeysr_   r   )r4   r   Z	torch_infZ	mode_dictr7   r7   r8   Z__init_monitor_mode  s
    z#ModelCheckpoint.__init_monitor_mode)r    r"   r!   r9   c                 C   sR   |d u r,|d u r,|d u r,d}d}t d n|p2d}|p:d}|| _|| _|| _d S )Nr   r   zQBoth every_n_train_steps and every_n_epochs are not set. Setting every_n_epochs=1)logdebugr<   r;   r:   )r4   r    r"   r!   r7   r7   r8   Z__init_triggers  s    	zModelCheckpoint.__init_triggersc                 C   s   | j S rH   )r;   r=   r7   r7   r8   r"     s    zModelCheckpoint.every_n_epochs)r?   currentr9   c                 C   sl   |d u rdS | j dkrdS t| j| j k }|r2dS tjtjd| j }||| j| j }|j	t
|}|S )NFrs   Tr{   )r   lenr+   r|   ltgtr   r,   rC   Zreduce_boolean_decisionro   )r4   r?   r   Zless_than_k_modelsZ
monitor_opZshould_update_best_and_saver7   r7   r8   check_monitor_top_k  s    
z#ModelCheckpoint.check_monitor_top_kr$   )r   metricsprefixr   r9   c                 C   s   |sd| j  d }td|}t|dkr|D ]R}|dd  }|rV|||d | }||d| d}||vr.td||< q.||}|r| j ||g}|S )	Nz{epoch}z{step}z(\{.*?)[:\}]r   r   z={z{0[])	CHECKPOINT_JOIN_CHARrefindallr   replacer|   r}   formatr   )clsr   r   r   r   groupsgroupnamer7   r7   r8   _format_checkpoint_name  s    
z'ModelCheckpoint._format_checkpoint_name)r   r   verr9   c                 C   sb   |p| j }| j||| jd}|dur:| j|d| f}| | j }| jr^tj| j|S |S )a  Generate a filename according to the defined template.

        Example::

            >>> tmpdir = os.path.dirname(__file__)
            >>> ckpt = ModelCheckpoint(dirpath=tmpdir, filename='{epoch}')
            >>> os.path.basename(ckpt.format_checkpoint_name(dict(epoch=0)))
            'epoch=0.ckpt'
            >>> ckpt = ModelCheckpoint(dirpath=tmpdir, filename='{epoch:03d}')
            >>> os.path.basename(ckpt.format_checkpoint_name(dict(epoch=5)))
            'epoch=005.ckpt'
            >>> ckpt = ModelCheckpoint(dirpath=tmpdir, filename='{epoch}-{val_loss:.2f}')
            >>> os.path.basename(ckpt.format_checkpoint_name(dict(epoch=2, val_loss=0.123456)))
            'epoch=2-val_loss=0.12.ckpt'
            >>> os.path.basename(ckpt.format_checkpoint_name(dict(epoch=2, val_loss=0.12), filename='{epoch:d}'))
            'epoch=2.ckpt'
            >>> ckpt = ModelCheckpoint(dirpath=tmpdir,
            ... filename='epoch={epoch}-validation_loss={val_loss:.2f}',
            ... auto_insert_metric_name=False)
            >>> os.path.basename(ckpt.format_checkpoint_name(dict(epoch=2, val_loss=0.123456)))
            'epoch=2-validation_loss=0.12.ckpt'
            >>> ckpt = ModelCheckpoint(dirpath=tmpdir, filename='{missing:d}')
            >>> os.path.basename(ckpt.format_checkpoint_name({}))
            'missing=0.ckpt'
            >>> ckpt = ModelCheckpoint(filename='{step}')
            >>> os.path.basename(ckpt.format_checkpoint_name(dict(step=0)))
            'step=0.ckpt'
        )r   Nv)	r   r   r   r   r   FILE_EXTENSIONr   rw   rx   )r4   r   r   r   Z	ckpt_namer7   r7   r8   format_checkpoint_name  s    
z&ModelCheckpoint.format_checkpoint_namec                 C   s   | j dur| j S t|jdkr|jd jdur<|jd j}n|j}|jd j}|jd j}t|trh|nd| }t	j
|t||d}nt	j
|jd}|S )a  Determines model checkpoint save directory at runtime. Reference attributes from the trainer's logger to
        determine where to save checkpoints. The path for saving weights is set in this priority:

        1.  The ``ModelCheckpoint``'s ``dirpath`` if passed in
        2.  The ``Logger``'s ``log_dir`` if the trainer has loggers
        3.  The ``Trainer``'s ``default_root_dir`` if the trainer has no loggers

        The path gets extended with subdirectory "checkpoints".
        Nr   Zversion_Zcheckpoints)r   r   rk   save_dirZdefault_root_dirr   version
isinstancestrrw   rx   r   )r4   r?   r   r   r   	ckpt_pathr7   r7   r8   Z__resolve_ckpt_dir;  s    

z"ModelCheckpoint.__resolve_ckpt_dirc                    s:     |} j|r4 fdd jj|ddD S t S )Nc                    s.   h | ]&} j tj|d  v rtj|qS )r   )CHECKPOINT_NAME_LASTrw   rx   splitnormpath).0pr=   r7   r8   	<setcomp>\  s   z9ModelCheckpoint._find_last_checkpoints.<locals>.<setcomp>F)Zdetail)rB   ru   existslsset)r4   r?   r   r7   r=   r8   _find_last_checkpointsX  s    

z&ModelCheckpoint._find_last_checkpoints)r   r9   c                 C   s>   | j dkr:| j|r:t| j|dkr:td| d d S )Nr   zCheckpoint directory z exists and is not empty.)r   ru   isdirr   r   r   )r4   r   r7   r7   r8   Z__warn_if_dir_not_emptyc  s    *z'ModelCheckpoint.__warn_if_dir_not_empty)rV   r?   del_filepathr9   c                 C   s@   |  |}| j}| ||r<||kr<| j ||d}|d7 }q|S N)r   r   )r   STARTING_VERSIONfile_exists)r4   rV   r?   r   rj   version_cntr7   r7   r8   &_get_metric_interpolated_filepath_nameg  s    

z6ModelCheckpoint._get_metric_interpolated_filepath_namec                 C   sf   t |j}|d}t|tr&| n
t|j|d< |d}t|trR| n
t|j	|d< |S )Nepochstep)
r   Zcallback_metricsra   r   r	   intr|   r}   rZ   rP   )r4   r?   rV   r   r   r7   r7   r8   rR   s  s    

"
"z#ModelCheckpoint._monitor_candidatesc                 C   s   | j s
d S | || j}| j}| ||rP|| jkrP| j|| j|d}|d7 }q| j| }| _| || |r||kr| || d S r   )r   r   r   r   r   r/   rm   _remove_checkpoint)r4   r?   rV   rj   r   previousr7   r7   r8   rT   }  s    
z%ModelCheckpoint._save_last_checkpointc              	   C   s   | j s
J || j }| ||r>|d us.J | ||| n>| jr||d }|d }td|dd|dd| j d| j  d S )Nr   r   Epoch d, global step : z was not in top )r   ra   r   _update_best_and_saver   r   r   )r4   r?   rV   r   r   r   r7   r7   r8   rg     s    
z(ModelCheckpoint._save_monitor_checkpointc                 C   sL   |  ||}| j| }| _| || | jdkrH|rH||krH| || d S )Nr   )r   r.   rm   r   r   )r4   r?   rV   rj   r   r7   r7   r8   rh     s
    z-ModelCheckpoint._save_none_monitor_checkpoint)r   r?   rV   r9   c           
      C   s  | j dkrt| jd n| j }d }t| j|krJ|dkrJ| j}| j| t|trt|rtj	t
| jdkrrdnd|jd}| |||}|| _|| j|< t| j|kr| jdkrtnt}|| j| jjd| _| j| j | _| jdkrtnt}|| j| jjd| _| j| j | _| jrj|d	 }|d
 }	td|dd|	dd| jd|dd| jdd|d|  | || |d ur||kr| || d S )Nrs   r   r   r   infz-inf)device)keyr   r   r   r   r   r   z	 reached z0.5fz (best z), saving model to z as top )r   r   r+   r,   popr   r	   r|   isnanr}   floatr   r   r   r*   rz   r   ra   r_   r.   r-   r   r   r   rm   r   )
r4   r   r?   rV   kr   rj   Z_opr   r   r7   r7   r8   r     s@    "
"z%ModelCheckpoint._update_best_and_save)rj   r9   c                 C   st   dd | j  D }|du r6| js&J tj| jd}| j|d}t	|| W d   n1 sf0    Y  dS )ztSaves the `best_k_models` dict containing the checkpoint paths with the corresponding scores to a YAML
        file.c                 S   s   i | ]\}}||  qS r7   )item)r   r   r   r7   r7   r8   
<dictcomp>      z+ModelCheckpoint.to_yaml.<locals>.<dictcomp>Nzbest_k_models.yamlw)
r+   itemsr   rw   rx   r   ru   openyamldump)r4   rj   Zbest_kfpr7   r7   r8   to_yaml  s    
zModelCheckpoint.to_yaml)rj   r?   r9   c                 C   s   | j |}|j|S )zChecks if a file exists on rank 0 and broadcasts the result to all other ranks, preventing the internal
        state to diverge between ranks.)ru   r   rC   rD   )r4   rj   r?   r   r7   r7   r8   r     s    zModelCheckpoint.file_existsc                 C   s   |j | dS )z1Calls the strategy to remove the checkpoint file.N)rC   Zremove_checkpoint)r4   r?   rj   r7   r7   r8   r     s    z"ModelCheckpoint._remove_checkpoint)NNNFNr   Fr   TNNNN)N)r$   T)NN)N)N)7__name__
__module____qualname____doc__r   r   r   r   r   r   r   ro   r   r   r&   propertyr>   rG   rK   r   r   rW   r\   r]   r   r`   rd   r	   rS   rm   rO   rY   r3   r1   r0   r2   r"   r   classmethodr   r   rB   r   r   rF   r   rR   rT   rg   rh   r   r   r   r   __classcell__r7   r7   r5   r8   r   0   s                 (	 

"	  
! ( 
	*
r   )+r   loggingrw   r   rI   rb   copyr   datetimer   typingr   r   r   r   weakrefr   numpyr~   r|   r   r	   Zpytorch_lightningplZ#lightning_fabric.utilities.cloud_ior
   Z lightning_fabric.utilities.typesr   Zpytorch_lightning.callbacksr   Z&pytorch_lightning.utilities.exceptionsr   Z%pytorch_lightning.utilities.rank_zeror   r   r   Z!pytorch_lightning.utilities.typesr   	getLoggerr   r   rf   r   r7   r7   r7   r8   <module>   s.   
