jSrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSKrSSK r SSK J r SSK J r SSKJrJr SSKJr SSKrSSKrSSKJr SSKJr SS KJrJr SS KJr SS K J!r!J"r"J#r# SS K$J$r$ SSK%r&SS K'J(r(J)r) SSK*J+r+J,r, SSK-J.r.J/r/ SSK0J1r1 SSK2J3r3 SSK4J5r5 SSK6J7r7 SSK8J9r9 \&RtRvRy5 \3"\=5r>\R~R\R~R\B55rC\R~R\CS5rCSrESrFSrGS0\CS4SjrHSrI"SS\5rJ\&RtRR"SS55rMg) zCore SpeechBrain code for running experiments. Authors * Peter Plantinga 2020, 2023 * Abdel Heba 2020 * Mirco Ravanelli 2020 * Aku Rouhe 2021 * Andreas Nautsch 2022 * Sylvain de Langen 2023 * Adel Moumen 2023, 2024 N)contextmanager)date)Enumauto)SimpleNamespace)resolve_references)version) DataParallel SyncBatchNorm)DistributedDataParallel) DataLoaderDistributedSamplerIterableDataset)tqdm) LoopedLoaderSaveableDataLoader)DistributedSamplerWrapperReproducibleRandomSampler) AMPConfig TorchAutocast)is_distributed_initialized) get_logger)rm_vector_weight_decay)prepare_profiler) RunOptionszlog-config.yamlbrain_intra_epoch_ckptTc[RRR5(Ga[R R U5(d[R"U5 Ub[R RUS5n[USS9n[Xb5nSSS5 [USSS9n[S[R"5-US9 [S[R RU5-US9 [S US9 [R "WU5 SSS5 ["R$"["R&"5R(5n U b?[R R+U R,5n [R."X5 [R RUS 5n S S S U 000n [RR0R3X<5 [4[6l[RR:R=5 [0R?S5 [0R?SU35 U(ak[RR0RA5n [[R RUS5SSS9nURCU 5 SSS5 [RRRE5 g!,(df  GNS=f!,(df  GN=f!,(df  N[=f![RRRE5 f=f)aCreate the output folder and relevant experimental files. Arguments --------- experiment_directory : str The place where the experiment directory should be created. hyperparams_to_save : str A filename of a yaml file representing the parameters for this experiment. If passed, references are resolved, and the result is written to a file in the experiment directory called "hyperparams.yaml". overrides : dict A mapping of replacements made in the yaml file, to save in yaml. log_config : str A yaml filename containing configuration options for the logger. save_env_desc : bool If True, an environment state description is saved to the experiment directory, in a file called env.log in the experiment directory. Nzhyperparams.yamlutf-8encodingwz# Generated %s from:)filez# %sz# yamllint disablezlog.txthandlers file_handlerfilenamezBeginning experiment!zExperiment folder: zenv.log)#sbutils distributedif_main_processospathisdirmakedirsjoinopenrprintrtodayabspathshutil copyfileobjinspect getmodule currentframef_backrealpath__file__copylogger setup_logging_logging_excepthooksys excepthookquirkslog_applied_quirksinfoget_environment_descriptionwrite ddp_barrier)experiment_directoryhyperparams_to_save overrides log_config save_env_deschyperparams_filenamef resolved_yamlr#module callingfilelog_filelogger_overridesdescription_strfos ڞC:\Users\johnn\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.13_qbz5n2kfra8p0\LocalCache\local-packages\Python313\site-packages\speechbrain/core.pycreate_experiment_directoryrX>sX25+ 88   / / 1 177==!566 01#.')ww||(*<($-@A$6q$DMA.gF!04::<?aH&277??3F#GGaP.Q7&&}a8 G&&w';';'='D'DEF! gg..v?  K>ww||$8)DH^j(-CD   HHOO ) )* G0CN HHOO . . 0 KK/ 0 KK-.B-CD E"$((//"M"M"OGGLL!5yA$HH_-  ((*UA@FF@ ((*sVBL$ K/#L$6A,L"F L$,L>L$/ K>9L$ L L$ L!L$$*Mc0[RSXU4S9 g)z-Interrupt exception raising to log the error.z Exception:)exc_infoN)r>error)exc_type exc_value exc_tracebacks rWr@r@s LLm(LLMcB\rSrSrSr\"5r\"5r\"5rSr g)Stagez*Simple enum to track stage of experiments.N) __name__ __module__ __qualname____firstlineno____doc__rTRAINVALIDTEST__static_attributes__rcr_rWraras4 FE FE 6Dr_rac\rSrSrSrS&SjrSrSrSrS'Sjr S'S jr S(S jr S r S r S rS)SjrS*SjrSrSrSrSrSrSrSr\R2"5S5rSrSrSrSS004Sjr\S5r Sr!Sr"Sr#SSS04S jr$S!r%\&S+S"j5r'\(RRRTRVS#5r,\(RRRTRZS$5r.S%r/g),Brainai Brain class abstracts away the details of data loops. The primary purpose of the `Brain` class is the implementation of the ``fit()`` method, which iterates epochs and datasets for the purpose of "fitting" a set of modules to a set of data. In order to use the ``fit()`` method, one should sub-class the ``Brain`` class and override any methods for which the default behavior does not match the use case. For a simple use case (e.g., training a single model with a single dataset) the only methods that need to be overridden are: * ``compute_forward()`` * ``compute_objectives()`` The example below illustrates how overriding these two methods is done. For more complicated use cases, such as multiple modules that need to be updated, the following methods can be overridden: * ``fit_batch()`` * ``evaluate_batch()`` Arguments --------- modules : dict[str, torch.nn.Module] These modules are passed to the optimizer by default if they have trainable parameters, and will have ``train()``/``eval()`` called on them. opt_class : Optional[Type[torch.optim]] A torch optimizer constructor that takes only the list of parameters (e.g. a lambda or partial function definition). By default, this will be passed all modules in ``modules`` at the beginning of the ``fit()`` method. This behavior can be changed by overriding the ``configure_optimizers()`` method. hparams : Optional[dict] Each key:value pair should consist of a string key and a hyperparameter that is used within the overridden methods. These will be accessible via an ``hparams`` attribute, using "dot" notation: e.g., self.hparams.model(x). run_opts : Optional[Union[RunOptions, dict]] A set of options to change the runtime environment, see ``RunOptions`` for a list. Typically in a script this comes from ``speechbrain.parse_args``, an alias for ``RunOptions.from_command_line_args``. If an option is not defined here (keep in mind that `parse_args` will inject some options by default), then the option is also searched for in hparams (by key). checkpointer : Optional[speechbrain.utils.checkpoints.Checkpointer] By default, this will be used to load checkpoints, and will have the optimizer added to continue training if interrupted. Example ------- >>> from torch.optim import SGD >>> class SimpleBrain(Brain): ... def compute_forward(self, batch, stage): ... return self.modules.model(batch[0] * self.hparams.scalar) ... ... def compute_objectives(self, predictions, batch, stage): ... return torch.nn.functional.l1_loss(predictions, batch[0]) >>> model = torch.nn.Linear(in_features=10, out_features=10) >>> brain = SimpleBrain( ... modules={"model": model}, ... opt_class=lambda x: SGD(x, lr=0.1), ... hparams={"scalar": 5}, ... run_opts={"device": "cpu"}, ... ) >>> brain.fit(range(1), ([torch.rand(10, 10), torch.rand(10, 10)],)) Nc>SUlX lXPl[U[5(a[ R "U5n[ 5nUR5R5HupxUbGXtR;a8Ub%Xs;a [RUS3SXG3-5 [XXG5 MOUb.Xs;a)[RSUS35 [XX75 M[XU5 M [RR[ :Xa"[RR"[$:d[R'S[)[RR5-S-[)[RR"5-S-[)[ 5-S-[)[$5-5 [*R,R/S5SL=(a! [*R,R/S 5SLUlUR2(aUR0(a [5S 5eUR6S :a&UR8S :a[R:"S 5 UR<c-[>R@RBRE5UlUR<S :XaS Ul#OSUR<;aNSUl#UR<RIS5up[JRLRO[QU 55 O#UR<RIS5S Ul#UR2(a-[JRLRS5S :Xa [5S5e[JRTRWU5RYUR<5Ul-URZHXn []URZU S5(dM#URZU RYUR<5URZU 'MZ Ub[_S-0UD6Ul0URb(aURd(dURb[]URS5(av[fRh"5n [RSU Rj35 [lRn"U Rj5URl8XRl9SUl:URv(a[R'S5 SUl<URz(a[R'S5 SUl<URFS :Xa+URxS:XdUR|S:Xa [5S5eURxS:H=(a URFS:Hn UR~(aU (a[R'S5 [RSU S35 [RSURxS35 [RSUR|S35 [R"[JR5[R"S5:a-[JRLRRU S9UlEO$[JR"UR<U S9UlE[R"URx5Rn[URFUS 9UlJ[R"UR|5Rn[URFUS 9UlKU (a2URb%URRS!URS"S#9 UR5 UR0(au[Q[*R,S5UlN[5(dEURS :a [5S$5e[R'S%5 [RS&5 S'UlPS UlQS UlRURbURRS(U5 UR(d%[RURS)5UlUSUlVUR(ap[RS*5 URUR-S+- UlZ[URURUR`R5UlV[]URZS,5(aURZRUl^gURZUl^g![4a# [JRLROS 5 GNVf=f).Nz. which is specified in hparams was overridden zby command line input to: z Run option z from hparams is usedzDetected Python .z-. We suggest using SpeechBrain with Python >=RANK LOCAL_RANKzTo use data_parallel backend, start your script with: python experiment.py hyperparams.yaml --data_parallel_backend=True To use DDP backend, start your script with: torchrun [args] experiment.py hyperparams.yamlrzThe options `ckpt_interval_minutes` and `ckpt_interval_steps` are mutually exclusive. Please keep only one active per experiment run.cpucuda:z0You must have at least 1 GPU to use DataParalleltocheckpoints_dirzRSince debug mode is active, switching checkpointer output to temporary directory: ztThe option `--auto_mix_prec` is deprecated and will be removed in the future. Please use `--precision=fp16` instead.fp16zxThe option `--bfloat16_mix_prec` is deprecated and will be removed in the future. Please use `--precision=bf16` instead.bf16zThe option `--precision` or `--eval_precision` is set to fp16. This option is not yet supported on CPU. Please use `--precision=bf16` or `--eval_precision=bf16` instead to enable mixed precision on CPU.zThe option `skip_nonfinite_grads` will be ignored because GradScaler is enabled and will automatically skip nonfinite gradients.zGradscaler enabled: ``z'Using training precision: `--precision=z.Using evaluation precision: `--eval_precision=z2.4.0)enabled) device_typedtypescalerT) optional_loadz ================ WARNING ===============Please add sb.ddp_init_group() into your exp.pyTo use DDP backend, start your script with: torchrun [args] experiment.py hyperparams.yamlzMTo use DDP, please add sb.utils.distributed.ddp_init_group() into your exp.pyzAOnly the main process is alive, all other subprocess were killed.brainz$Pytorch profiler has been activated.rQrc)_optimizers_dict opt_class checkpointer isinstancedictrfrom_dictionaryas_dictitemsoverridden_argsr>rEsetattrrA version_infomajorPYTHON_VERSION_MAJORminorPYTHON_VERSION_MINORwarningstrr,environgetdistributed_launchdata_parallel_backend ValueErrorckpt_interval_minutesckpt_interval_stepsexitdevicer(r)r* infer_devicer}splittorchru set_deviceint device_countnn ModuleDictrwmoduleshasattrrhparamsdebugdebug_persistentlytempfileTemporaryDirectorynamepathlibPathrxtempdir train_sampler auto_mix_prec precisionbfloat16_mix_preceval_precisionskip_nonfinite_gradsr parse __version__amp GradScalerrr from_namer~r training_ctxevaluation_ctxadd_recoverableprint_trainable_parametersrankravg_train_lossstepoptimizer_steptqdm_colored_barfromkeys tqdm_barcolorprofilerprofile_training profile_stepsprofile_warmuptot_prof_stepsr output_folderrQ raw_modules)selfrrrrun_optsrrun_opt_defaultsargdefault_ device_indexrQrgradscaler_enabled train_dtype eval_dtypes rW__init__Brain.__init__s $"( h % %!11(;H&<,446<<>LC#/G/G(G&3>KK%MN6x}oFG8=1$ k#.CDE7<07+?&    " "&: :  &&*>> NN"c&&,,-.c&&,,-.B B *+ ,  *+,  JJNN6 "$ . 9 |,D8   % %$*A*AA   % % )d.F.F.J HHB  ;; ((..;;=DK ;;% $D  t{{ "%D  )"&++"3"3C"8 %%c,&78 ${{005a8D   % %%***A*A*Cq*HOP Pxx**7366t{{C llFt||F+T22'+||F';'>'>t{{'K V$#  *5W5DL JJ++!!-))+<==113G KK229,,A 18 W\\0JD   -)0   % "    NN9 $DN  ! ! NN9 $DN   u $ NNf $(;(;v(E4  NNf $ C)9)9V)C   $ $); NN,   +,>+?qAB =dnn=MQOP  ?#'#5#58K8K#Kq"PD ,##"" **DMt||X.. LL    o ) %%a( )sAe//)ffc *SnSnURR5H:nX#R5- nUR(dM(XR5- nM< URR nUS:Xa8[ RS5 [ RUSUSUSSSS35 g US:Xaa[ RS5 [RR RU5n[ RUSUSUSSS S35 g S U-U- n[RR RU5n[RR RU5n[ RUSUSUSUS S35 g ) z7Prints the number of trainable parameters in the model.rzThe model has no parameters!z; Model Statistics: * Total Number of Trainable Parameters: z * Total Number of Parameters: z" * Trainable Parameters represent z.2fz% of the total size.z&The model has no trainable parameters!z.4fdN) r parametersnumel requires_grad __class__rdr>rrEr(r)format_order_of_magnitude)rtotal_trainable_paramstotal_parameters parameter class_nameformatted_total_paramspercentage_trainableformatted_trainable_paramss rWr Brain.print_trainable_parameterss!"002I  1 1 &&&&//*;;&3^^,, q NN9 : KK,;;Q:RS11A0BC445c7:NP  $q ( NNC D%'XX__%N%N & " KK,;;Q:RS11G0HI445c7:NP ,,/?? !99* ' &(XX__%N%N & " KK,;;U:VW11G0HI44H3MMac r_c[e)aForward pass, to be overridden by sub-classes. Arguments --------- batch : torch.Tensor or tensors An element from the dataloader, including inputs for processing. stage : Stage The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST Returns ------- torch.Tensor or torch.Tensors The outputs after all processing is complete. Directly passed to ``compute_objectives()``. NotImplementedError)rbatchstages rWcompute_forwardBrain.compute_forwards "!r_c[e)aCompute loss, to be overridden by sub-classes. Arguments --------- predictions : torch.Tensor or torch.Tensors The output tensor or tensors to evaluate. Comes directly from ``compute_forward()``. batch : torch.Tensor or tensors An element from the dataloader, including targets for comparison. stage : Stage The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST Returns ------- loss : torch.Tensor A tensor with the computed loss. r)r predictionsrrs rWcompute_objectivesBrain.compute_objectivess $"!r_cg)zGets called when a stage starts. Useful for defining class variables used during the stage. Arguments --------- stage : Stage The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST epoch : int The current epoch count. Nrc)rrepochs rWon_stage_startBrain.on_stage_start+s r_cg)a8Gets called at the end of a stage. Useful for computing stage statistics, saving checkpoints, etc. Arguments --------- stage : Stage The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST stage_loss : float The average loss over the completed stage. epoch : int The current epoch count. Nrc)rr stage_lossrs rW on_stage_endBrain.on_stage_end9 r_c U[RR:XaURX5n[RR R "U40UD6nURbVUbS[U[5(d[U[5(a)X2R-nURRXe5 U$)aCreates DataLoaders for Datasets. This is used by ``fit()`` and ``evaluate()`` if they just receive Datasets. Alternatively, this can be called from outside the Brain subclass. In that case, the DataLoader should be passed to ``fit()`` in place of the dataset. The Stage.TRAIN DataLoader is handled specially. It has extra args for shuffle and drop_last. In DDP a DistributedSampler is created (unless the dataset is an IterableDataset). NOTE ---- Some important DataLoader arguments are passed via **loader_kwargs, e.g., batch_size, num_workers, pin_memory. NOTE ---- By default, ``evaluate()`` specifies ckpt_prefix=None to stop the test DataLoader being added to the checkpointer. If you need to add a recoverable after saving checkpoints (e.g., at test time, after checkpointing the training), and still be able to recover reasonably, you should probably specify ``allow_partial_load=True``. Arguments --------- dataset : Dataset A set of data to use to create data loader. If the Dataset is a DynamicItemDataset, PaddedBatch is used as the default collate_fn, unless specified in loader_kwargs. stage : Stage The stage of the experiment: Stage.TRAIN, Stage.VALID, Stage.TEST ckpt_prefix : str, None Prefix to use for SaveableDataLoader Checkpoint name. The Stage name is added to this to create the full key. Set to None to not save the DataLoader. **loader_kwargs : dict Additional keyword arguments to the DataLoader. E.g., batch_size, num_workers, pin_memory. Returns ------- DataLoader for the input dataset ) r(rari_train_loader_specificsdataio dataloadermake_dataloaderrrrrrr)rdatasetr ckpt_prefix loader_kwargsrckpt_keys rWrBrain.make_dataloaderIsd BHHNN " 88PMYY))99  $    )':'9::j,77#ZZ/H    - -h Cr_cURSS5nURSS5nU(a`UR(dOUb [S5e[RRSS5n[ XS9nX0lUR US'US UR(Ga[U[5(d[URS5(aURRS :HnOS nURS S5nUb1[UURUUS 9UlSUS'UR US'U$URS 5c1[UURUUS9UlSUS'UR US'U$[URS S5URUS9UlUR US 'U$UR(a*[U[5(a[R!S5 U$)NsamplershuffleFz>Cannot specify both shuffle=Trueand a sampler in loader_kwargsSB_GLOBAL_SEEDi&l!)seedsortingrandomT drop_last)rrr  batch_sampler)rr r)rr zDCannot automatically solve distributed sampling for IterableDataset.)rrrr,rrrrrrrr rrrr>r)rrrr r r  shuffle_ddprs rWrBrain._train_loader_specificss##It4 ##Iu5 422" 5::>>"2I>D/CG!( '+'9'9M) $i(  " " ":g+O+Ot||Y//LL((H4# %))+u=I"%>'# &",1 i(+/+=+= i(21""?3;%7'' &",1 i(+/+=+= i(&?!%%ot<'&" 261C1C o.   $ $G_)M)M NN' r_cUR5 UR5 UR5 URbURR 5 gg)zGets called at the beginning of ``fit()``, on multiple processes if ``distributed_count > 0`` and backend is ddp. Default implementation compiles the jit modules, initializes optimizers, and loads the latest checkpoint to resume training. N)_compile_wrap_distributedinit_optimizersrrecover_if_possiblers rW on_fit_startBrain.on_fit_startsN        (    1 1 3 )r_cURR5nURbUR(a[ UR5nURU5UlSUR 0UlURb'URRSUR 5 gg[RS5 g)aCalled during ``on_fit_start()``, initialize optimizers after parameters are fully configured (e.g. DDP, jit). The default implementation of this method depends on an optimizer class being passed at initialization that takes only a list of parameters (e.g., a lambda or a partial function definition). This creates a single optimizer that optimizes all trainable params. Override this class if there are multiple optimizers. Nr optimizerzSNo `opt_class` was provided to this Brain class, skipping optimizer initialization.) rrrremove_vector_weight_decayrrrrrr>rE)r all_paramss rWrBrain.init_optimizerss\\,,. >> %..3DLLA !^^J7DN$/#@D   ,!!11+t~~N- KK5 r_cURb@URUR5R5HnURUS9 M gURbUR RUS9 gg)zSets the gradients of all optimized ``torch.Tensor``s to zero if ``set_to_none=False`` (default) or to None otherwise. Setting gradients to None should save the memory, e.g. during ``evaluate()`` and thus larger batch might be used. N set_to_none)rfreeze_optimizersvalues zero_gradrr)rr"opts rWr%Brain.zero_grad sg    +--d.B.BCJJL + 6M ^^ ' NN $ $ $ =(r_cRURbURRXS9 gg)aGets called at the beginning of ``evaluate()`` Default implementation loads the best-performing checkpoint for evaluation, based on stored metrics. Arguments --------- max_key : str Key to use for finding best checkpoint (higher is better). By default, passed to ``self.checkpointer.recover_if_possible()``. min_key : str Key to use for finding best checkpoint (lower is better). By default, passed to ``self.checkpointer.recover_if_possible()``. Nmax_keymin_key)rr)rr*r+s rWon_evaluate_startBrain.on_evaluate_starts1"    (    1 1 2  )r_cURUR-S:HnURX5 URU(+5 UR UR U[ RR5nURX1[ RR5nSSS5 URRWUR- 5nURU5 UR5 SSS5 U(aUR5 URUWWU5 UR!5R#5$!,(df  N=f!,(df  Nh=f)aFit one batch, override to do multiple updates. The default implementation depends on a few methods being defined with a particular behavior: * ``compute_forward()`` * ``compute_objectives()`` * ``optimizers_step()`` Also depends on having optimizers passed at initialization. Arguments --------- batch : list of torch.Tensors Batch of data to use for training. Default implementation assumes this batch has two elements: inputs and targets. Returns ------- detached loss rN)rgrad_accumulation_factoron_fit_batch_startno_syncrrr(rarirrscalecheck_loss_isfinitebackwardoptimizers_stepon_fit_batch_enddetachrt)rr should_stepoutputsloss scaled_losss rW fit_batchBrain.fit_batch0s ,yy4#@#@@QF  3 \\k/ *""..ubhhnnE..wrxx~~N#++++t444K  $ $[ 1  "+   " eWdK@{{}  ""#"+ *s& EAE&AE E E E'c[R"U5(dPU=RS- slURUR:a [ S5e[ R S5 gg)a"Check if the loss is finite. If the loss is not finite, log a helpful message and increment the `nonfinite_count`. If the `nonfinite_count` exceeds the `--nonfinite_patience` threshold, stop the training and raise an error. This check is particularly useful when the loss becomes NaN or inf, while the parameters and gradients remain finite. It helps prevent getting stuck in an infinite loop during training. Arguments --------- loss : tensor The loss tensor after ``backward()`` has been called but before the optimizers ``step()``. rzLoss is not finite and patience is exhausted. To debug, wrap `fit()` with autograd's `detect_anomaly()`, e.g. with torch.autograd.detect_anomaly(): brain.fit(...)zPatience not yet exhausted.N)risfinitenonfinite_countnonfinite_patiencerr>r)rr:s rWr3Brain.check_loss_isfiniteYsa"~~d##  A % ##d&=&== I<=$r_cHURR5HnUR(dMURcM%[R "UR5R 5(aMZSUl[RSURS35 M g)zXChecks if the gradients are finite. If not, it will emit a warning and set them to zero.Nz Gradients z% contain NaN or Inf. Setting to None.) rrrgradrr?allr>rr)rparams rWcheck_gradientsBrain.check_gradientsxsn\\,,.E"""uzz'=~~ejj15577!%EJNN$UZZL0UV /r_cU$)zBy default, this method returns the passed optimizers. Override this method if you want to freeze some optimizers during training. To do so, return a of active optimizers. rc)r optimizerss rWr#Brain.freeze_optimizerss r_c8URbURUR5nOURbSUR0nOgUR 5HnUR R U5 M UR 5HGn[RRRURSSUR5 MI UR R5(d!UR(aUR5 UR 5HnUR R!U5 M UR R#5 UR 5HnUR%SS9 M U=R&S- slg)zvPerforms a step of gradient descent on the optimizers. This method is called every ``grad_accumulation_factor`` steps.NrrparamsTr!r)rr#rrr$runscale_rrr)clip_grad_norm_ param_groups max_grad_norm is_enabledrrGrupdater%r)rvalid_optimizersr&s rWr5Brain.optimizers_stepsJ    +#55d6J6JK  ^^ '!,T^^<  $**,C KK  %-$**,C HHNN * *  #H-t/A/A -{{%%''D,E,E  "$**,C KK  S !- #**,C MMdM +- q r_cg)aCalled at the beginning of ``fit_batch()``. This method is not called under the AMP context manager. Do not assume automatic casting of the input batch to a lower precision (e.g. fp16). Arguments --------- batch : list of torch.Tensors Batch of data to use for training. Default implementation assumes this batch has two elements: inputs and targets. should_step : boolean Whether optimizer.step() was called or not. Nrc)rrr8s rWr0Brain.on_fit_batch_startrr_cg)aCalled after ``fit_batch()``. Arguments --------- batch : list of torch.Tensors Batch of data to use for training. Default implementation assumes this batch has two elements: inputs and targets. outputs : list or dictionary of torch.Tensors Returned value of compute_forward(). loss : torch.Tensor Returned value of compute_objectives(). should_step : boolean Whether optimizer.step() was called or not. Nrc)rrr9r:r8s rWr6Brain.on_fit_batch_ends r_cUR URXS9nURX1US9nSSS5 WR5R 5$!,(df  N,=f)aEvaluate one batch, override for different procedure than train. The default implementation depends on two methods being defined with a particular behavior: * ``compute_forward()`` * ``compute_objectives()`` Arguments --------- batch : list of torch.Tensors Batch of data to use for evaluation. Default implementation assumes this batch has two elements: inputs and targets. stage : Stage The stage of the experiment: Stage.VALID, Stage.TEST Returns ------- detached loss rN)rrrr7rt)rrroutr:s rWevaluate_batchBrain.evaluate_batchsZ, &&u&:C**3U*CD!{{}  ""! s A A!cUR[RU5 URR 5 UR 5 SUlURb6[URS5(aURRU5 [R"5nSn[UURSU(+URSS9nURbURR5 UGHnUR (a["R%S5 GOZU=RS- sl US- nUR'U5nUR)XR*5UlUR-UR*S9 URbwURR5 URR.UR0:a9["R%S 5 URR35 [55 UR6(aURUR8:Xa OEUR;XE5(dGM[UR=5 [R"5nSnGM SSS5 UR SS 9 UR?[RUR*U5 S UlSUl g!,(df  NW=f) Nr set_epochTtrain)initial dynamic_ncolsdisablecolourzTrain iteration limit exceededr) train_lossz+The profiler finished, training is stopped.r!r) rrarirrar%r@rrr`timerrrrstart_optimizer_step_limit_exceededr>rEr<update_averager set_postfixstep_numrstopquitr debug_batches_should_save_intra_epoch_ckpt_save_intra_epoch_ckptr) r train_setrenablelast_ckpt_timesteps_since_ckpttrr:s rW _fit_trainBrain._fit_trains? EKK/   !    )g    / /     ( ( /  IIJ%%g.  }}( ##%66KK @A Q  A% ~~e,&*&9&9--'# )<)< ===,MM&&(}}--0C0CC I **,::$))t/A/A"A55"//1%)YY[N'($? V 4( %++t':':EB! ]  sFK,K K%cURcgURS::aURS::ag[R"5U- S- nSURs=:=(a U:Os nU=(d SURs=:=(a U:*Os n[ 5(dU$U/n[ R RUSS9 US$)zDetermines if an intra-epoch checkpoint should be saved. Returns True if there's a checkpointer and time or steps has exceeded limit. FrgN@)src)rrrrgrrr*broadcast_object_list)rrtruelapsed_minutesdecisionbroadcast_lists rWrp#Brain._should_save_intra_epoch_ckpt5s    $  % % *t/G/G1/L 99;74?t11CCOCOq4#;#;OO?OO*++O 'ZN    3 3N 3 J!!$ $r_cbUGbUR[RU5 URR 5 Sn[ R "5 [USU(+URSS9HvnU=RS- sl URU[RS9nURXd5nUR(dMZURUR:XdMv O SUl UR[RXB5 SSS5 gg!,(df  g=f)NrTvalidrcrdrerr[r)rrarjrevalrno_gradrrrr]rjrror)r valid_setrrsavg_valid_lossrr:s rW _fit_validBrain._fit_validTs     U 3 LL    N!"& &J--g6 E IINI..uEKK.HD%)%8%8%NNzzzdii43E3E&E !!%++~E!! !!sA5D D ++D  D.cUR(a[RS5 g[U[5(dB[U[ 5(d-UR "U4S[RR0UD6nUbX[U[5(dC[U[ 5(d.UR "U4[RRSS.UD6nUR5 UcUR(+nU=(a( [RRR5nUHVnUR!X(US9 UR#X8US9 UR$(aXR&:XdUR((dMV g g)aIterate epochs and datasets to improve objective. Relies on the existence of multiple functions that can (or should) be overridden. The following methods are used and expected to have a certain behavior: * ``fit_batch()`` * ``evaluate_batch()`` * ``update_average()`` If the initialization was done with distributed_count > 0 and the distributed_backend is ddp, this will generally handle multiprocess logic, like splitting the training data into subsets for each device and only saving a checkpoint on the main process. Arguments --------- epoch_counter : iterable Each call should return an integer indicating the epoch count. train_set : Dataset, DataLoader A set of data to use for training. If a Dataset is given, a DataLoader is automatically created. If a DataLoader is given, it is used directly. valid_set : Dataset, DataLoader A set of data to use for validation. If a Dataset is given, a DataLoader is automatically created. If a DataLoader is given, it is used directly. progressbar : bool Whether to display the progress of each epoch in a progressbar. train_loader_kwargs : dict Kwargs passed to `make_dataloader()` for making the train_loader (if train_set is a Dataset, not DataLoader). E.G. batch_size, num_workers. DataLoader kwargs are all valid. valid_loader_kwargs : dict Kwargs passed to `make_dataloader()` for making the valid_loader (if valid_set is a Dataset, not DataLoader). E.g., batch_size, num_workers. DataLoader kwargs are all valid. Returns ------- None z8Test only mode, skipping training and validation stages.Nr)rr)rrrrs)rrrs) test_onlyr>rErr rrr(rarirjr noprogressbarr)r*r+rwrr debug_epochsri) r epoch_counterrrr progressbartrain_loader_kwargsvalid_loader_kwargsrsrs rWfit Brain.fitlsJj >> KKJ   y* - -)\22,,!#3FI  y* - -)\22,,hhnn & I   "000KG!5!5!E!E!G#E OOiVO L OOiVO L ...666#r_c^URSL=(a URUR:$N)optimizer_step_limitrrs rWri$Brain._optimizer_step_limit_exceededs1  % %T 1 A##t'@'@@ r_chURRSSS[S0[RS9 g)z,Saves a CKPT with specific intra-epoch flag.Frc([UR;$r)INTRA_EPOCH_CKPT_FLAGmeta)cs rW.Brain._save_intra_epoch_ckpt..s%:aff%Dr_T) end_of_epoch num_to_keepckpt_predicater verbosityN)rsave_and_keep_onlyrloggingDEBUGrs rWrqBrain._save_intra_epoch_ckpts3 ,,D'.mm - r_c[[S5nU(dURb [S5e[ 5nUR (aMURc[ UR 5nO*[ UR5n[RS5 [ 5nUR(aMURc[ UR 5nO*[ UR5n[RS5 X#-H!nX@R ;dM[SUS35e UHn[R "UR UURURURS9nURUR 5UR U'UR#U5 M UHWn[RR%UR U5nURUR 5UR U'MY g![a'n[RS US U35 SnAGM SnAff=f) z;Compile requested modules with either JIT or TorchInductor.compileNz_'compile_module_keys' specified, but this install of PyTorch seems to be too old to support it.zy--compile and --compile_module_keys are both specified. Only modules specified in --compile_module_keys will be compiled.zm--jit and --jit_module_keys are both specified. Only modules specified in --jit_module_keys will be compiled.zmodule z% is not defined in your hparams file.)mode fullgraphdynamic'zh' in 'compile_module_keys' failed to compile and will be skipped (may fallback onto JIT, if specified): )rrcompile_module_keysrsetrrr>rjitjit_module_keys compile_modecompile_using_fullgraph#compile_using_dynamic_shape_tracing Exceptionrwrdiscardscript)rcompile_availablerrrrQes rWrBrain._compiles#E95 T%=%=%I5  "e <<''/&)$,,&7#&)$*B*B&C#X % 88##+"%dll"3"%d&:&:";T (9D<<' dV#HI:(D LL&**":: DD "(4;;!7DLL   # #D )#(&$DYY%%dll4&89F!'4;;!7DLL $ v##$#'   s/AH  H=H88H=cUR(dUR(dgUR(aURR5Hup[ SUR 555(dM,[ R"U5nURS:Xa[USURS9nO [UUR/URS9nX RU'M gURR5HEup[ SUR 555(dM,[U5nX RU'MG g)z5Wrap modules with distributed wrapper when requested.Nc38# UHoRv M g7frr.0ps rW *Brain._wrap_distributed..%D0C10Cgloo) device_idsfind_unused_parametersc38# UHoRv M g7frrrs rWrr7rr) rrrranyrr convert_sync_batchnormdistributed_backendDDPrrDP)rrrQs rWrBrain._wrap_distributeds&&t/I/I   $ $ $ 2 2 4 D0A0A0CDDD*AA&IF//69!$"'+373N3N" "%"(, }373N3N" *0LL&!5$!% 2 2 4 D0A0A0CDDDZF)/LL&!5r_cUcUR(+nU=(a( [RRR 5n[ U[ 5(d<[ U[5(d'SUS'UR"U[R40UD6nURX#S9 UR[RSS9 URR5 Sn[R "5 [#USU(+UR$SS9HvnU=R&S - slUR)U[RS 9n UR+X5nUR,(dMZUR&UR.:XdMv O UR1[RUS5 SSS5 S UlU$!,(df  N=f) a]Iterate test_set and evaluate brain performance. By default, loads the best-performing checkpoint (as recorded using the checkpointer). Arguments --------- test_set : Dataset, DataLoader If a DataLoader is given, it is iterated directly. Otherwise passed to ``self.make_dataloader()``. max_key : str Key to use for finding best checkpoint, passed to ``on_evaluate_start()``. min_key : str Key to use for finding best checkpoint, passed to ``on_evaluate_start()``. progressbar : bool Whether to display the progress in a progressbar. test_loader_kwargs : dict Kwargs passed to ``make_dataloader()`` if ``test_set`` is not a DataLoader. NOTE: ``loader_kwargs["ckpt_prefix"]`` gets automatically overwritten to ``None`` (so that the test DataLoader is not added to the checkpointer). Returns ------- average test loss Nrr))rrTtestrrr[r)rr(r)r*r+rr rrrarkr,rrrrrrrrr]rjrror) rtest_setr*r+rtest_loader_kwargsrs avg_test_lossrr:s rWevaluateBrain.evaluate;s{D  "000KG!5!5!E!E!G x , ,(L1104 } -++%**(:H w@ EJJd3  ]]_"" ))&1   Q **5 *C $ 3 3D H :::$))t/A/A"A   ejj- >  #_s5A5G.G %G Gc[R"U5(a,X"UR- -nU[U5UR- - nU$)zUpdate running average of the loss. Arguments --------- loss : torch.tensor detached loss, a single float value. avg_loss : float current running average. Returns ------- avg_loss : float The average loss. )rr?rfloat)rr:avg_losss rWrjBrain.update_averages@ >>$   499, ,H d dii/ /Hr_c#d# U(a/nURR5H8n[US5(dMURUR5 SUlM: Sv SnURR5H$n[US5(dMX$UlUS- nM& gSv g7f)aCopies pytorch's implementation for doing no_sync across all modules. Explanation: nn.module.no_sync() is a context manager for when one does not want to sync gradients, which happens when using both DDP and gradient accumulation. Speechbrain brain's class can contain multiple modules and calling no_sync on these individually would be very awkward, therefore this contextmanager exists. Arguments --------- use : bool If set to `False` will still sync gradients, useful to make behavior toggleable. Yields ------ None require_backward_grad_syncFNrr)rr$rappendr)ruseold_values_listrQis rWr1 Brain.no_syncs$  O,,--/v'CDD&&v'H'HI491 0 A,,--/v'CDD4C4F1Q 0 sB.B0cURURURS.n[USSS9nUR [ R "U55 SSS5 g!,(df  g=f)N)rrrr#r r!)rrrr1rGyamldump)rr- save_dictr#s rW_save Brain._savesUII"11"11 $g .! GGDIIi( )/ . .s &A A-cTA[USS9n[R"U5nSSS5 WSUlUSUlSU;aIUR R nSUS3nUS- n[R"U5 URUl gUSUl g!,(df  N|=f) Nr r!rrrz'optimizer_step' not found in z checkpoint.z7 Using the saved 'step' value (BACKWARDS COMPATIBILITY)) r1r safe_loadrrrrdwarningswarnr)rr-rrOrclsnameMSGs rW_recoverBrain._recovers  $ )Qq)I*f% '(89 9 ,nn--G27)<HC L LC MM# "&))D "+,<"=D * )s B B')rrrr}rrrrr@rrrrrrrrrrrrrr)NNNNNr)z dataloader-)F)NN)T)0rdrerfrgrhrrrrrrrrrrr%r,r<r3rGr#r5r0r6rrr]rwrprrpropertyrirqrrrrjrr1r(r) checkpoints mark_as_saverrmark_as_loaderrrlrcr_rWrnrnsbAJ q f-^&*   "+8HTEN4*< >,'#R>>-!^  " ]]_##4?B%>F8`D   >8@0> EN(!!FXX''*(*XX((>)>r_rn)Nrhr7rr,rr5rArrgr contextlibrdatetimerenumrrtypesrrr hyperpyyamlr packagingr torch.nnr rr torch.nn.parallelr rtorch.utils.datar rrr speechbrainr(speechbrain.dataio.dataloaderrrspeechbrain.dataio.samplerrrspeechbrain.utils.autocastrrspeechbrain.utils.distributedrspeechbrain.utils.loggerrspeechbrain.utils.optimizersrspeechbrain.utils.profilingrspeechbrain.utils.run_optsrr)rC apply_quirksrdr>r-dirnamer4r<DEFAULT_LOG_CONFIGr0rrrrXr@rarregister_checkpoint_hooksrnrcr_rWrs,    %! *=LLJ@D/?81 H WW__RWW__X%>?WW\\"46GH0 ! N+bN D//v>v>0v>r_