o 8wi\:@sdZddlZddlZddlmZmZddlmZeeZ GdddZ Gddde Z Gd d d e Z Gd d d e Z d dZGdddZdddZddZdS)zXLoggers for experiment monitoring. Authors * Peter Plantinga 2020 * Jarod Duret 2023 N)if_main_processmain_process_only) get_loggerc@s"eZdZdZ    dddZdS) TrainLoggerz:Abstract class defining an interface for training loggers.NFcCst)a'Log the stats for one epoch. Arguments --------- stats_meta : dict of str:scalar pairs Meta information about the stats (e.g., epoch, learning-rate, etc.). train_stats : dict of str:list pairs Each loss type is represented with a str : list pair including all the values for the training pass. valid_stats : dict of str:list pairs Each loss type is represented with a str : list pair including all the values for the validation pass. test_stats : dict of str:list pairs Each loss type is represented with a str : list pair including all the values for the test pass. verbose : bool Whether to also put logging information to the standard logger. )NotImplementedError)self stats_meta train_stats valid_stats test_statsverboser [/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/speechbrain/utils/train_logger.py log_statsszTrainLogger.log_statsNNNF)__name__ __module__ __qualname____doc__rr r r rrsrc@sDeZdZdZd ddZdddZddd Ze    dd d ZdS)FileTrainLoggerzText logger of training information. Arguments --------- save_file : str The file to use for logging train information. precision : int Number of decimal places to display. Default 2, example: 1.35e-5. cCs||_||_dSN) save_file precision)rrrr r r__init__=s zFileTrainLogger.__init__NcCsvt|trd|krdkrnn |d|jd}nt|tr)|d|jd}|dur4|d|}|d|S) z+Convert one item to string, handling floatsg?gY@.feN z: ) isinstancefloatr)rkeyvaluedatasetr r r_item_to_stringAs" zFileTrainLogger._item_to_stringcsdfdd|DS)z,Convert all stats to a single string summaryz, csg|] \}}||qSr )r$).0kvr#rr r Nsz4FileTrainLogger._stats_to_string..)joinitems)rstatsr#r r(r_stats_to_stringKsz FileTrainLogger._stats_to_stringTc Cs||}d|fd|fd|ffD]\}}|dur"|d|||7}qt|jddd} t|| d Wdn1s||_ddlm}d|_tr||j|_iiidd|_dS)Nr) SummaryWriter)r/r0r1meta)save_dirtorch.utils.tensorboardr<writerr global_step)rr>r<r r rrws   zTensorboardLogger.__init__NFcCs|jdd7<|D]\}}|j|||jdq d|fd|fd|ffD]d\}} | dur1q(| D]V\} } | |j|vrGd|j|| <| d|} t| trq| D]}|j|| d} |j| || | |j|| <qUq5| }|j|| d} |j| || | |j|| <q5q(dS) r.r=r/r0r1Nr/)rAr+r@ add_scalarrlist)rrr r r r namer"r#r,stat value_listtagnew_global_stepr r rrs2   zTensorboardLogger.log_statscCs|jj|||jd|ddS)zAdd audio signal in the logs.r=) sample_rateN)r@ add_audiorA)rrFr"rKr r r log_audios zTensorboardLogger.log_audiocCs.t|}|dur|j|||jddSdS)zAdd a figure in the logs.Nr=)plot_spectrogramr@ add_figurerA)rrFr"figr r r log_figureszTensorboardLogger.log_figurer) rrrrrrrrMrQr r r rr;js   % r;c@s.eZdZdZddZe    dddZdS) WandBLoggera Logger for WandB (Weights & Biases). This logger is designed to be used in the same way as TrainLogger and supports handling nested dictionaries as well. Arguments --------- initializer: callable A callable function that initializes the WandB run. For more information on the parameters that can be passed to the initializer, refer to the documentation: https://docs.wandb.ai/ref/python/init *args: tuple Positional arguments to be passed to the initializer function. **kwargs: dict Keyword arguments to be passed to the initializer function. Example ------- To initialize the logger, use the following pattern in hparams.yaml: ``` train_logger: !new:speechbrain.utils.train_logger.WandBLogger initializer: !name:wandb.init entity: speechbrain project: sb_project name: sb_run reinit: True resume: False dir: !ref /wandb id: sb_run resume: allow ``` NOTE ---- If there is an issue with the WandB Logger initialization, it raises an exception. c OsJzd|_tr||i||_WdSWdSty$}z|dd}~ww)Nz7There was an issue with the WandB Logger initialization)runr Exception)r initializerargskwargsrr r rrszWandBLogger.__init__NFc Cs~i}d|fd|fd|ffD] \}}|durq |||<q |dd} | dur3|jji||| ddS|ji||dS)r.r/r0r1Nepoch)step)getrSlog) rrr r r r logsr#r,rYr r rrs    zWandBLogger.log_statsr)rrrrrrrr r r rrRs%rRcCs2z ddl}|jjWStytdYdSw)zeReturns the TorchVision image saver, if available or None if it is not - optional dependency rNz2torchvision is not available - cannot save figures) torchvisionutils save_image ImportErrorr7warning)r]r r r_get_image_savers   rbc@sjeZdZdZdejiddeiddZdZ dd d Z d d Z d dZ ddZ ddZe ddZdS)ProgressSampleLoggeraA logger that outputs samples during training progress, used primarily in speech synthesis but customizable, reusable and applicable to any other generative task Natively, this logger supports images and raw PyTorch output. Other custom formats can be added as needed. Example: In hparams.yaml progress_sample_logger: !new:speechbrain.utils.progress_samples.ProgressSampleLogger output_path: output/samples progress_batch_sample_size: 3 format_defs: foo: extension: bar saver: !speechbrain.dataio.mystuff.save_my_format kwargs: baz: qux formats: foobar: foo In the brain: Run the following to "remember" a sample (e.g. from compute_objectives) self.hparams.progress_sample_logger.remember( target=spectrogram_target, output=spectrogram_output, alignments=alignments_output, my_output= raw_batch={ "inputs": inputs, "spectrogram_target": spectrogram_target, "spectrogram_output": spectrogram_output, "alignments": alignments_output } ) Run the following at the end of the epoch (e.g. from on_stage_end) self.progress_sample_logger.save(epoch) Arguments --------- output_path: str the filesystem path to which samples will be saved. formats: dict A mapping from keys to formats. format_defs: dict a dictionary with format identifiers as keys and dictionaries with handler callables and extensions as values. The signature of the handler should be similar to torch.save Example: { "myformat": { "extension": "myf", "saver": somemodule.save_my_format } } batch_sample_size: int The number of items to retrieve when extracting a batch sample pth) extensionsaverrWpng)rawimageriNrBcCs@i|_|pi|_t|j|_|dur|j|||_||_dSr)progress_samplesformatsdict_DEFAULT_FORMAT_DEFS format_defsupdatebatch_sample_size output_path)rrqrkrnrpr r rr[s    zProgressSampleLogger.__init__cCs i|_dS)z.Initializes the collection of progress samplesN)rjrr r rresetfs zProgressSampleLogger.resetcKs|jdd|DdS)zUpdates the internal dictionary of snapshots with the provided values Arguments --------- **kwargs: dict the parameters to be saved with cSi|] \}}|t|qSr detach)r%r!r"r r r tz1ProgressSampleLogger.remember..N)rjror+)rrWr r rrememberjs zProgressSampleLogger.remembercsNt|trfdd|D}|St|tjtfr#|dj}|S|}|S)aObtains a sample of a batch for saving. This can be useful to monitor raw data (both samples and predictions) over the course of training Arguments --------- value: dict|torch.Tensor|list the raw values from the batch Returns ------- result: object the same type of object as the provided value csi|] \}}||qSr )get_batch_sampler%r! item_valuerrr rrws z9ProgressSampleLogger.get_batch_sample..N)rrlr+torchTensorrErp)rr"resultr rrrrzws  z%ProgressSampleLogger.get_batch_samplecCs&|jD] \}}||||qdS)zSaves all items previously saved with remember() calls Arguments --------- epoch: int The epoch number N)rjr+ save_item)rrXr!datar r rsaveszProgressSampleLogger.savec Cstj|jt|}tj|st||j||j }|j |}|dur+t d|d|d}tj||}|d} | durQ| ||fi|ddSdS)aSaves a single sample item Arguments --------- key: str the key/identifier of the item data: torch.Tensor the data to save epoch: int the epoch number (used in file path calculations) NzUnsupported format {format}rrerfrW) ospathr*rqstrexistsmakedirsrkrZDEFAULT_FORMATrn ValueError) rr!rrX target_pathformat format_def file_nameeffective_file_namerfr r rrs    zProgressSampleLogger.save_item)NNrB)rrrrr}rrbrmrrrsryrzrrr r r rrcs" C    rc FcCszddl}|dddlm}Wnty tdYdSw| }|j |d}|j |ddd| ||sG||S) zcReturns the matplotlib spectrogram if available or None if it is not - optional dependency rNAggz0matplotlib is not available - cannot log figures)figsizeautolower)aspectorigin) matplotlibusematplotlib.pyplotpyplotr`r7rarvcpunumpysqueezefigureimshowcolorbar tight_layoutclose) spectrogramapfig_size output_figrpltrPr r rrNs     rNcCsDt|tjr|}|St|trdd|D}|S|}|S)alDetaches the specified object from the graph, which can be a single tensor or a dictionary of tensors. Dictionaries of tensors are converted recursively Arguments --------- value: torch.Tensor|dict a tensor or a dictionary of tensors Returns ------- result: torch.Tensor|dict a tensor of dictionary of tensors cSrtr rur{r r rrwrxzdetach..)rr}r~rvrrlr+)r"rr r rrvs   rv)NrF)rrr}speechbrain.utils.distributedrrspeechbrain.utils.loggerrrr7rrr;rRrbrcrNrvr r r rs  8NI  *