o ϯiX@sddlZddlZddlZddlZddlZddlmZddlZddl Z ddl m m Z zddlZWn ey<dZYnwddlmZmZddlmZGdddeZdd Z dd d Zdd d Z    dddZ dddZddZddZdS)N)suppress)ClipLossgather_features) is_masterc@s*eZdZdZddZddZd ddZd S) AverageMeterz1Computes and stores the average and current valuecCs |dSN)resetselfr M/home/ubuntu/.local/lib/python3.10/site-packages/laion_clap/training/train.py__init__s zAverageMeter.__init__cCsd|_d|_d|_d|_dS)Nr)valavgsumcountr r r r r s zAverageMeter.resetrcCs8||_|j||7_|j|7_|j|j|_dSr)rrrr)r rnr r r update!szAverageMeter.updateN)r)__name__ __module__ __qualname____doc__rr rr r r r rs rcCst|dr|jS|S)Nmodule)hasattrr)modelr r r unwrap_model(s rc+Cst|j}|jdkrtjjjnt} |t|j |j d|j |j |j |j|jd} |dj|dj} } |jr@| dur@| || j} tt| jdd}|jdkrY| jt}t}t}t}t| D]\}}| ||}t|t r|!D]}||q~n|||}|d}|"t|t|t r|!D]}|#qn|#| )||||\}}}}}}|jr| ||||||d } n| |||d } Wdn1swYt|t r>|dur-|$| %|!D]/}|j r"|&|'||( |)|Wdn 1swYq|)|q|"nW| %|!D]}|)q5nF|dur||$| %|j rr|&|'||( |)|Wdn 1slwYn|)||"n| %|)t*$t+|j,-d td |jrt+|j.-d td Wdn 1swY|"t|t}|d}!t/|r}|d d ks|!| kr}t|t rt0|d }"nt0|}"|!|"|j }#| j}$d|!| }%|"| 1|"|1}&|1}'t|t r|jrbt23d|d|#d|d|$d|%dd|j4dd|j5dd|j5dd|j5dddd|!Dd|&dd|'d|j4|j4|j4|&|'d d|!Dd!}(nt23d|d|#d|d|$d|%dd|j4dd|j5dd|j5dd|j5ddd"d|!Dd|&d|j4|j4|j4|&d#d|!Dd$}(n|jrt23d|d|#d|d|$d|%dd|j4dd|j5dd|j5dd|j5dd|j6d d%d&d|&dd|'d|j4|j4|j4|&|'|j6d d%d!}(nGt23d|d|#d|d|$d|%dd|j4dd|j5dd|j5dd|j5dd|j6d d%d&d|&d|j4|j4|j4|&|j6d d%d$}(|(7D]+\})}*d'|)})|dur]|8|)|*||j9rst9dusjJd(t9|)|*d)|iqI|:|:qjdS)*NampT) local_lossgather_with_grad cache_labelsrank world_size use_horovodmlp_lossweight_loss_kappatrainr toytext)audio_features text_features logit_scale_a logit_scale_taudio_features_mlptext_features_mlp)r*r+r,rdwaveformgY@z Train Epoch:  [>/z (z.0fz %)] Loss: z#.5gz#.4gz ) Data (t): z.3fz Batch (t): z LR: cSg|] }|jddqSrlr param_groups.0o_r r r z#train_one_epoch..z Logit Scale Audio: zLogit Scale Text: cSr5r6r8r:r r r r=r>)loss data_time batch_time scale_audio scale_textr7cSr5r6r8r:r r r r=r>cSr5r6r8r:r r r r=r>)r?r@rArBr7r75fztrain/Please install wandb.step);torchdevice precisioncudarautocastrr&rrrr!r"horovod clap_mlplosskappa dataloadersampler distributed set_epoch num_batchesmathceillog num_samples dataset_typedatasetgenerate_queuertime enumerate isinstancedictvaluesr zero_gradscalebackward synchronizeunscale_skip_synchronizerFno_gradrr,clamp_r-rlenitemlogginginforrr9items add_scalarwandbr )+rdataepoch optimizerscaler schedulerargs tb_writerrHrKr?rOrPnum_batches_per_epoch sample_digitsloss_m batch_time_m data_time_mendibatchrFsaudiostextsr<r*r+r.r/r,r- total_loss batch_count batch_sizerWsamples_per_epochpercent_completelogit_scale_scalar_alogit_scale_scalar_tlog_datanamerr r r train_one_epoch/s                               "   $  $  $   $   rcs2i}|js t|s |St|j}|t|rtd|jdkr'tjjj nt }|j ddgkrb|jr7t dt |||||||}|D]} || qEd|vrZ|d|it|||}n+d|vr|jr||jdksx||jkr|dj} d} | j} i} |jrd dggggd | d <n d dggd | d <tt| D]Z\}}|}|d }ttdd|dD}|D] }|| vr|jrd dggggd | |<qd dggd | |<q|||||\}}}}}}|jr"|jrt||||dd|j|j|j|jd \}}}}nt||dd|j|j|j|jd\}}t|r| |jd7} g|d D]d krm| d |!| d |!|jrk| d |!| d |!q4t"#t"$dd|dDkd}| d |!%dt&|'| d |!%dt&|'|jr| d |!%dt&|'| d |!%dt&|'q4Wdn 1swYt|r|ddkrt()d|d| d| dqt|r}i}| D]o|jr@t*t+| dt+| d|!t+| dt+| d|!|jd}nt*t+| dt+| d|!|jd}fd d!|,D|<||d|vr{|d|iq Wdn 1swYt|r|s|St()d|d"d#-d$d|D|j.r|,D]\}}|dur|/d%|||qt0t1j2-|j3d&d'}|4t56||4d#Wdn 1swY|j7rt7dusJd(|,D]\}}t78d%||d|iq|S|S))Nz Evaluating...rClotho audiocapszEParallel evaluation not supported for eval only Clotho and audiocaps.rprr)cumulative_lossrWall_audio_featuresall_text_featuresall_audio_features_mlpall_text_features_mlpallrrWrrr)cS$g|]}d|dddqS-r4joinsplitr;br r r r=F$zevaluate..__url__F) r*r+r.r/rrr!r"r#r$)r*r+rrr!r"r#r$rrrrcSrrrrr r r r=rr0z Eval Epoch: r2z / ])r*r+r,r.r/r-r$)r*r+r,r$ci|] \}}d||qSr4r r;kvrr r zevaluate..  cSs$g|]}ddd|DqS) cSs&g|]\}}|dt|ddqS)z: z.4f)roundrr r r r=&z'evaluate...)rrl)r;mr r r r=szval/z results.jsonlza+rE)9 parallel_evalrrGrHevalprintrIrJrrKrval_dataset_namesNotImplementedErrorevaluate_clotho_audiocapsr_rkeys"select_top_metric_clotho_audiocaps val_frequencyepochsrOrWrMrfr\listsetrr!r"rLshapeappendcpunpwherearray index_selecttensorlongrjrk get_metricscatrlr save_logsrmopenospathcheckpoint_pathwritejsondumpsrnrV)rrorprtrumetricsrHrKval_metrics_per_datasetrrOrWsamples_per_val eval_infor|r}rr all_namesrr*r+r.r/r,r-idxmetrics_single_datasetrfr rr evaluate s                        ^           rFc CsHi}|rv|||}|} |||} | } t|jd} t|| t| | t| | t| | d} | |d<|jd|d<|| d| | dd}tt | dd}nI|||}|}t|jd} t|| t|| d} | |d<|jd|d<||d}tt | dd}| D]^\}}tj |d d }t||kd}|}|d||d <tt|d||d <d D]}t||k||d|<qtt|dkd|dd||d<q|S)NrrrrW) audio_to_text text_to_audiorrT descending _mean_rank _median_rankrr'z_R@r'rz_mAP@10)tdetachrrGarangerrF cross_entropyrirhviewrlargsortrnumpymeanrfloormedian)r*r+r,r.r/r-r$ra_logits_per_audioa_logits_per_textt_logits_per_audiot_logits_per_textlabelsrlogits ground_truthlogits_per_audiologits_per_textrlogitrankingpredsrr r r rs\            ,rc' sx|dj}ti}t|D]\} } | } jdkr3ddlmfdd| dDtndd lm fd d| dDfd d d Dt t d d| dD} | D]} | | vrrddggd|| <qa||| d|}|d|}t j|dd}t j|dd}t t dd| dD} | D]Nttdd| dDkd}|d|dt||d|dd|jdgdt|d|jdgqWdn1swYqi}| D]|dd|\}}|}tj|ddd}tj|ddd}|||tddjdji}|jd|d<t|jdfddtdD}fddtdD}t|t|d}| |d <g}tdD]9}ddd|ddf}tt!|"dd}tj#|d!d"}t||kd}||$qtj%|dd#}|d|d$<t&t'|d|d%<d&D]}t||k|d'|<qtt|d(kd|dd|d)<g} g}!tD]g}|ddf}"tj#|"d!d"}t|d|ddd}tt(|gd|"ddkd}#t)|#}$|!|$$|#|#d(k$}%t*tdt!|%d|%dd}&| |&qt| |d*<d&D]}tt|!|k|d+|<qfd,d |+D|<qWd|S1swY|S)-a Adapted from https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py. 1. for text-to-audio retrieval, do 5 times and average the results 2. for R@1, R@5, R@10 in audio-to-text retrieval, take the best rank among 5 text 3. for map@10 in audio-to-text retrieval: 3.1: sort the rank of 5 text 3.2: exclude the rank >=10 (0-index) 3.3: compute the map regarding the remaining ranks: np.mean(np.arange(1, len(ranks)+1) / ranks). (3.3) That is, take the top ranks of 5 text that is < 10, and assign the descending number as ground truth. (3.3) E.g.: the ground truth of first rank of the 5 text should be 1, the second rank should be 2, etc. r transformerrtokenizecsg|]}|qSr r r;rrr r r=Yz-evaluate_clotho_audiocaps.. full_textr) tokenizercsg|] }|jdqS)tmodelrr)rtrr r r=]r>cs&i|]tfddDqS)csg|]}|qSr r rrr r r=^rz8evaluate_clotho_audiocaps...)rGr)r;)rrr r^rz-evaluate_clotho_audiocaps..cSrrrrr r r r=drrrrNr)dimcSrrrrr r r r=trcSrrrrr r r r=xrrrrzdataset z, logits_per_audio shape: z, logits_per_text shape: rWc s4g|]}tddddd|fqSrNrrreshaper;d)rrrWr r r=  c s4g|]}tddd|ddfqSrrr)rrrWr r r=r rrTr)axistext_to_audio_mean_ranktext_to_audio_median_rankrztext_to_audio_R@r'ztext_to_audio_mAP@10zaudio_to_text_mAP@10zaudio_to_text_R@crrr rrr r rr),rOrGrfr\r clap_modulerrrorrrrr normalizerrrrrrrrrrrrrjrkrrangerrirhrrr concatenaterrstackminrrl)'rrorprtrKrHrurOrr|r}rrrr*r+rval_metrics_allr,r-raudio_to_text_losstext_to_audio_lossr pred_textrrrrrpred_text_concatrmap_allpred_audio_all logit_singleall_predmin_predall_pred_filter map_singler ) rtrrrrrWrrrr rAs                $ & &"  rcCsNg}|D]}|||d|||dd}||qt|S)zI Calculate performance for Clotho+AudioCaps for model selection. z/audio_to_text_mAP@10z/text_to_audio_mAP@10r)rrrr)rselection_performance_allrselection_performancer r r 0calculate_selection_performance_clotho_audiocapss   r!c Cs<t|dsJt|}i}|D]#}||D]}|||||dddd|dd<qq||d<|d|d<||||_||_|St|}|j}||kri}|D]#}||D]}|||||dddd|dd<qcq[||d<|d|d<||||_||_|S||j|S)Ntop_selection_performancer4rz-toprrpztop-selection-epoch)rr!rrr top_metricr") rrrtr  metric_updaterrselection_performance_newselection_performance_oldr r r rs:  2   2   rr)NNNF)rrjrTrr[ contextlibrrrrGtorch.nn.functionalnn functionalrrn ImportErrorr rrrQrobjectrrrrrrr!rr r r r sB      \s J "