a
    dT                     @   s  d Z dgZddlZddlZddlmZ ddlmZmZm	Z	m
Z
mZmZmZ ddlmZ ddlmZ ddlZddlmZmZmZ dd	lmZ dd
lmZmZ ddlmZ ddlm Z  edZ!e!rddl"m#Z" ddl$m%Z% nd\Z"Z%e&e'Z(dZ)G dd deZ*dS )z
Neptune Logger
--------------
NeptuneLogger    N)	Namespace)AnyDict	GeneratorListOptionalSetUnion)RequirementCache)Tensor)_add_prefix_convert_params_sanitize_callable_params)
Checkpoint)Loggerrank_zero_experiment)ModelSummary)rank_zero_onlyzneptune-client)new)Run)NNz*source_code/integrations/pytorch-lightningc                	       sF  e Zd ZdZdZdZdZdddddddee ee ee ed	 ee	 ee
d
 fddZddddZeedddZeedddZeee ee ee ed	 eddddZeee
f dddZeee
f ddddZeeedddZeeeddd Zeeeee
f ef dd!d"d#ZedAeeeeef f ee  dd$d%d&Z!eedd' fd(d)Z"eee dd*d+Z#edBd-e dd.d/d0Z$ee%dd1d2d3Z&eee%ed4d5d6Z'e(eee
f ee)e d7d8d9Z*e(dCeee
f ee+d:d;d<Z,eee dd=d>Z-eee dd?d@Z.  Z/S )Dr   aX  
    Log using `Neptune <https://neptune.ai>`_.

    Install it with pip:

    .. code-block:: bash

        pip install neptune-client

    or conda:

    .. code-block:: bash

        conda install -c conda-forge neptune-client

    **Quickstart**

    Pass NeptuneLogger instance to the Trainer to log metadata with Neptune:

    .. code-block:: python


        from pytorch_lightning import Trainer
        from pytorch_lightning.loggers import NeptuneLogger

        neptune_logger = NeptuneLogger(
            api_key="ANONYMOUS",  # replace with your own
            project="common/pytorch-lightning-integration",  # format "<WORKSPACE/PROJECT>"
            tags=["training", "resnet"],  # optional
        )
        trainer = Trainer(max_epochs=10, logger=neptune_logger)

    **How to use NeptuneLogger?**

    Use the logger anywhere in your :class:`~pytorch_lightning.core.module.LightningModule` as follows:

    .. code-block:: python

        from neptune.new.types import File
        from pytorch_lightning import LightningModule


        class LitModel(LightningModule):
            def training_step(self, batch, batch_idx):
                # log metrics
                acc = ...
                self.log("train/loss", loss)

            def any_lightning_module_function_or_hook(self):
                # log images
                img = ...
                self.logger.experiment["train/misclassified_images"].log(File.as_image(img))

                # generic recipe
                metadata = ...
                self.logger.experiment["your/metadata/structure"].log(metadata)

    Note that syntax: ``self.logger.experiment["your/metadata/structure"].log(metadata)`` is specific to Neptune
    and it extends logger capabilities. Specifically, it allows you to log various types of metadata
    like scores, files, images, interactive visuals, CSVs, etc.
    Refer to the `Neptune docs <https://docs.neptune.ai/you-should-know/logging-metadata#essential-logging-methods>`_
    for more detailed explanations.
    You can also use regular logger methods ``log_metrics()``, and ``log_hyperparams()`` with NeptuneLogger
    as these are also supported.

    **Log after fitting or testing is finished**

    You can log objects after the fitting or testing methods are finished:

    .. code-block:: python

        neptune_logger = NeptuneLogger(project="common/pytorch-lightning-integration")

        trainer = pl.Trainer(logger=neptune_logger)
        model = ...
        datamodule = ...
        trainer.fit(model, datamodule=datamodule)
        trainer.test(model, datamodule=datamodule)

        # Log objects after `fit` or `test` methods
        # model summary
        neptune_logger.log_model_summary(model=model, max_depth=-1)

        # generic recipe
        metadata = ...
        neptune_logger.experiment["your/metadata/structure"].log(metadata)

    **Log model checkpoints**

    If you have :class:`~pytorch_lightning.callbacks.ModelCheckpoint` configured,
    Neptune logger automatically logs model checkpoints.
    Model weights will be uploaded to the: "model/checkpoints" namespace in the Neptune Run.
    You can disable this option:

    .. code-block:: python

        neptune_logger = NeptuneLogger(project="common/pytorch-lightning-integration", log_model_checkpoints=False)

    **Pass additional parameters to the Neptune run**

    You can also pass ``neptune_run_kwargs`` to specify the run in the greater detail, like ``tags`` or ``description``:

    .. testcode::
        :skipif: not _NEPTUNE_AVAILABLE

        from pytorch_lightning import Trainer
        from pytorch_lightning.loggers import NeptuneLogger

        neptune_logger = NeptuneLogger(
            project="common/pytorch-lightning-integration",
            name="lightning-run",
            description="mlp quick run with pytorch-lightning",
            tags=["mlp", "quick-run"],
        )
        trainer = Trainer(max_epochs=3, logger=neptune_logger)

    Check `run documentation <https://docs.neptune.ai/essentials/api-reference/run>`_
    for more info about additional run parameters.

    **Details about Neptune run structure**

    Runs can be viewed as nested dictionary-like structures that you can define in your code.
    Thanks to this you can easily organize your metadata in a way that is most convenient for you.

    The hierarchical structure that you apply to your metadata will be reflected later in the UI.

    You can organize this way any type of metadata - images, parameters, metrics, model checkpoint, CSV files, etc.

    See Also:
        - Read about
          `what object you can log to Neptune <https://docs.neptune.ai/you-should-know/what-can-you-log-and-display>`_.
        - Check `example run <https://app.neptune.ai/o/common/org/pytorch-lightning-integration/e/PTL-1/all>`_
          with multiple types of metadata logged.
        - For more detailed info check
          `user guide <https://docs.neptune.ai/integrations-and-supported-tools/model-training/pytorch-lightning>`_.

    Args:
        api_key: Optional.
            Neptune API token, found on https://neptune.ai upon registration.
            Read: `how to find and set Neptune API token <https://docs.neptune.ai/administration/security-and-privacy/
            how-to-find-and-set-neptune-api-token>`_.
            It is recommended to keep it in the `NEPTUNE_API_TOKEN`
            environment variable and then you can drop ``api_key=None``.
        project: Optional.
            Name of a project in a form of "my_workspace/my_project" for example "tom/mask-rcnn".
            If ``None``, the value of `NEPTUNE_PROJECT` environment variable will be taken.
            You need to create the project in https://neptune.ai first.
        name: Optional. Editable name of the run.
            Run name appears in the "all metadata/sys" section in Neptune UI.
        run: Optional. Default is ``None``. The Neptune ``Run`` object.
            If specified, this `Run`` will be used for logging, instead of a new Run.
            When run object is passed you can't specify other neptune properties.
        log_model_checkpoints: Optional. Default is ``True``. Log model checkpoint to Neptune.
            Works only if ``ModelCheckpoint`` is passed to the ``Trainer``.
        prefix: Optional. Default is ``"training"``. Root namespace for all metadata logging.
        \**neptune_run_kwargs: Additional arguments like ``tags``, ``description``, ``capture_stdout``, etc.
            used when run is created.

    Raises:
        ModuleNotFoundError:
            If required Neptune package is not installed.
        ValueError:
            If argument passed to the logger's constructor is incorrect.
    /ZhyperparamsZ	artifactsNTZtraining)api_keyprojectnamerunlog_model_checkpointsprefixr   )r   r   r   r   r   r   neptune_run_kwargsc                   s~   t sttt | ||||| t   || _|| _|| _|| _	|| _
|| _|| _d | _| jd urz|   tj| jt< d S N)_NEPTUNE_AVAILABLEModuleNotFoundErrorstr_verify_input_argumentssuper__init___log_model_checkpoints_prefix	_run_name_project_name_api_key_run_instance_neptune_run_kwargs_run_short_id_retrieve_run_datapl__version___INTEGRATION_VERSION_KEY)selfr   r   r   r   r   r   r   	__class__ j/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/pytorch_lightning/loggers/neptune.pyr%      s    

zNeptuneLogger.__init__)returnc                 C   sV   | j d usJ | j   | j drF| j d  | _| j d  | _nd| _d| _d S )Nzsys/idzsys/nameZOFFLINEzoffline-name)r+   waitexistsfetchr-   r(   r2   r5   r5   r6   r.      s    
z NeptuneLogger._retrieve_run_datac                 C   s   i }z
| j }W n ty    Y n0 | jd ur6| j|d< | jd urJ| j|d< | jd ur^| j|d< z| jd urt| j|d< W n ty   Y n0 |S )Nr   Z	api_tokenr   r   )r,   AttributeErrorr)   r*   r-   r(   )r2   argsr5   r5   r6   _neptune_init_args  s"    







z NeptuneLogger._neptune_init_args)keysr7   c                 G   s&   | j r| j| j g|S | j|S )zXReturn sequence of keys joined by `LOGGER_JOIN_CHAR`, started with `_prefix` if defined.)r'   LOGGER_JOIN_CHARjoin)r2   r?   r5   r5   r6   _construct_path_with_prefix!  s    z)NeptuneLogger._construct_path_with_prefix)r   r   r   r   r   r7   c                 C   sN   |d urt |tstdtdd | ||fD p4|}|d urJ|rJtdd S )Nz7Run parameter expected to be of type `neptune.new.Run`.c                 s   s   | ]}|d uV  qd S r   r5   ).0argr5   r5   r6   	<genexpr>3      z8NeptuneLogger._verify_input_arguments.<locals>.<genexpr>zjWhen an already initialized run object is provided you can't provide other neptune.init_run() parameters.
)
isinstancer   
ValueErrorany)r   r   r   r   r   Zany_neptune_init_arg_passedr5   r5   r6   r#   '  s    	z%NeptuneLogger._verify_input_argumentsc                 C   s   | j  }d |d< |S )Nr+   )__dict__copyr2   stater5   r5   r6   __getstate__:  s    
zNeptuneLogger.__getstate__)rM   r7   c                 C   s   || _ tjf i | j| _d S r   )rJ   neptuneinit_runr>   r+   rL   r5   r5   r6   __setstate__@  s    zNeptuneLogger.__setstate__c                 C   s   | j S )a  
        Actual Neptune run object. Allows you to use neptune logging features in your
        :class:`~pytorch_lightning.core.module.LightningModule`.

        Example::

            class LitModel(LightningModule):
                def training_step(self, batch, batch_idx):
                    # log metrics
                    acc = ...
                    self.logger.experiment["train/acc"].log(acc)

                    # log images
                    img = ...
                    self.logger.experiment["train/misclassified_images"].log(File.as_image(img))

        Note that syntax: ``self.logger.experiment["your/metadata/structure"].log(metadata)``
        is specific to Neptune and it extends logger capabilities.
        Specifically, it allows you to log various types of metadata like scores, files,
        images, interactive visuals, CSVs, etc. Refer to the
        `Neptune docs <https://docs.neptune.ai/you-should-know/logging-metadata#essential-logging-methods>`_
        for more detailed explanations.
        You can also use regular logger methods ``log_metrics()``, and ``log_hyperparams()``
        with NeptuneLogger as these are also supported.
        )r   r;   r5   r5   r6   
experimentD  s    zNeptuneLogger.experimentc                 C   s4   | j s.tjf i | j| _ |   tj| j t< | j S r   )r+   rO   rP   r>   r.   r/   r0   r1   r;   r5   r5   r6   r   b  s
    zNeptuneLogger.run)paramsr7   c                 C   s.   t |}t|}| j}| |}|| j|< dS )a  
        Log hyper-parameters to the run.

        Hyperparams will be logged under the "<prefix>/hyperparams" namespace.

        Note:

            You can also log parameters by directly using the logger instance:
            ``neptune_logger.experiment["model/hyper-parameters"] = params_dict``.

            In this way you can keep hierarchical structure of the parameters.

        Args:
            params: `dict`.
                Python dictionary structure with parameters.

        Example::

            from pytorch_lightning.loggers import NeptuneLogger

            PARAMS = {
                "batch_size": 64,
                "lr": 0.07,
                "decay_factor": 0.97
            }

            neptune_logger = NeptuneLogger(
                api_key="ANONYMOUS",
                project="common/pytorch-lightning-integration"
            )

            neptune_logger.log_hyperparams(PARAMS)
        N)r   r   PARAMETERS_KEYrB   r   )r2   rS   Zparameters_keyr5   r5   r6   log_hyperparamsm  s
    #
zNeptuneLogger.log_hyperparams)metricsstepr7   c                 C   sH   t jdkrtdt|| j| j}| D ]\}}| j| | q*dS )zLog metrics (numeric values) in Neptune runs.

        Args:
            metrics: Dictionary with metric names as keys and measured quantities as values.
            step: Step number at which the metrics should be recorded, currently ignored.
        r   z&run tried to log from global_rank != 0N)	r   ZrankrH   r   r'   r@   itemsr   log)r2   rV   rW   keyvalr5   r5   r6   log_metrics  s
    
zNeptuneLogger.log_metrics)statusr7   c                    s.   | j s
d S |r|| j| d< t | d S )Nr]   )r+   r   rB   r$   finalize)r2   r]   r3   r5   r6   r^     s
    zNeptuneLogger.finalizec                 C   s   t jt  dS )zGets the save directory of the experiment which in this case is ``None`` because Neptune does not save
        locally.

        Returns:
            the root directory where experiment logs get saved
        z.neptune)ospathrA   getcwdr;   r5   r5   r6   save_dir  s    zNeptuneLogger.save_dirzpl.LightningModule)model	max_depthr7   c                 C   s2   t t||d}tjjj|dd| j| d< d S )N)rd   re   txt)content	extensionzmodel/summary)r"   r   rO   typesFileZfrom_contentr   rB   )r2   rd   re   Z	model_strr5   r5   r6   log_model_summary  s    zNeptuneLogger.log_model_summary)checkpoint_callbackr7   c           
      C   s~  | j s
dS t }| d}t|dr^|jr^| |j|}|| | j| d|  |j t|dr|j	
 D ]4}| ||}|| | j| d|  | qrt|dr|jr|j| j| d< | |j|}|| | j| d|  |j | j|rH| j }| ||}t|| D ]}	| j| d|	 = q.t|drz|jrz|j   | j| d	< dS )
zAutomatically log checkpointed model. Called after model checkpoint callback saves a new checkpoint.

        Args:
            checkpoint_callback: the model checkpoint callback instance
        Nzmodel/checkpointslast_model_pathr   best_k_modelsbest_model_pathzmodel/best_model_pathbest_model_scorezmodel/best_model_score)r&   setrB   hasattrrm   _get_full_model_nameaddr   uploadrn   r?   ro   r9   Zget_structure(_get_full_model_names_from_exp_structurelistrp   cpudetachnumpy)
r2   rl   Z
file_namesZcheckpoints_namespaceZmodel_last_namerZ   Z
model_nameexp_structureZuploaded_model_namesZfile_to_dropr5   r5   r6   after_save_checkpoint  s4    





z#NeptuneLogger.after_save_checkpoint)
model_pathrl   r7   c                 C   s`   t |drX|j tjj }| |s:t|  d| dtj| t|d \}}n| }|S )zZReturns model name which is string `model_path` appended to `checkpoint_callback.dirpath`.dirpathz was expected to start with .N)	rr   r~   r_   r`   sep
startswithrH   splitextlen)r}   rl   Zexpected_model_pathfilepath_r5   r5   r6   rs     s    

z"NeptuneLogger._get_full_model_name)r{   	namespacer7   c                 C   s0   | | j}|D ]}|| }q|}t| |S )zHReturns all paths to properties which were already logged in `namespace`)splitr@   rq   _dict_paths)clsr{   r   Zstructure_keysrZ   Zuploaded_models_dictr5   r5   r6   rv     s
    
z6NeptuneLogger._get_full_model_names_from_exp_structure)dpath_in_buildr7   c                 c   sT   |  D ]F\}}|d ur&| d| n|}t|ts<|V  q| ||E d H  qd S )Nr   )rX   rG   dictr   )r   r   r   kvr`   r5   r5   r6   r     s
    
zNeptuneLogger._dict_pathsc                 C   s   | j S )zMReturn the experiment name or 'offline-name' when exp is run in offline mode.)r(   r;   r5   r5   r6   r     s    zNeptuneLogger.namec                 C   s   | j S )zLReturn the experiment version.

        It's Neptune Run's short_id
        )r-   r;   r5   r5   r6   version  s    zNeptuneLogger.version)N)rc   )N)0__name__
__module____qualname____doc__r@   rT   ZARTIFACTS_KEYr   r"   boolr   r%   r.   propertyr   r>   rB   staticmethodr   r#   rN   rQ   r   r   rR   r   r   r
   r   rU   r   floatintr\   r^   rb   rk   r   r|   rs   classmethodr	   rv   r   r   r   r   __classcell__r5   r5   r3   r6   r   2   s|    &	"**
	/ )+r   __all__loggingr_   argparser   typingr   r   r   r   r   r	   r
   Z lightning_utilities.core.importsr   Ztorchr   Zpytorch_lightningr/   Z!lightning_fabric.utilities.loggerr   r   r   Zpytorch_lightning.callbacksr   Z pytorch_lightning.loggers.loggerr   r   Z)pytorch_lightning.utilities.model_summaryr   Z%pytorch_lightning.utilities.rank_zeror   r    rO   r   Zneptune.new.runr   	getLoggerr   rY   r1   r   r5   r5   r5   r6   <module>   s,   $
