From d53736f0dd372b5a7cfe9990f794002e0eafcbab Mon Sep 17 00:00:00 2001
From: dohe0342 <kimdohe1070@gmail.com>
Date: Sat, 10 Dec 2022 13:43:51 +0900
Subject: [PATCH] from local

---
 .../.train.py.swp                             | Bin 86016 -> 90112 bytes
 .../train.py                                  |  51 ++++++++++++++----
 2 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.train.py.swp b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/.train.py.swp
index e905bdd6ca19c8a29d5b1c0feab7945f7c24ef5a..7a8b76b99fb4c87a691a63b57d8476ab0133bb6c 100644
GIT binary patch
delta 1603
zcmZ|PS!feM7zgmFCXH9L_OMjM4vHpfSnB~|rGi4miU$SjgH_^ocakn{cGumF^=O+4
zDl~$EBZ`6-KKSBQgD41UrTQX>2R;g-f=^QQK~PcrPqv#_H4DG&Zf542Z)SF;??SNe
zL9nKyVN1&fx?E}?xGD&F8a|<%9jN&@rK~zPwjw;HteH=(tgWmO#kiWi8_@%~>6MSe
zH)&u;cvX;)UF9f1bJXs;l&POcNC56kAmk!sVH+%muj2`M0hb^Hy^w@C@NOI-uiz9U
zVHtc05por-z*=~Zo(7;7dY~KjgUF!8tCL%DU}Z2vaWUOtjUr|7A}&YiWNKPxHf^Hq
zg_k%plvuGWt}<mi$=An^X}MZT+vhbYym(EDucMlXz%%pRh*0RNS!G6aRAv6Wzp8b<
zCEIe!kz^&EPQ{b9tU56|@}F+UnPn=S&b&e<qw@VXDh45<^3SiUtNs&(Z#qA1p`Xug
z7DZ<9d|xrD+f7{^l42SOEh$+{-YKfZb$gy@(;$W_dP5sLm2M*n(T4ChJ=J{Eo?z|&
zW0oxcyZz#b<I1*6yJSP{piUQKmcJOXtC@x*Ns;2@<d<6H8!I`+k(vE_fSHm5{Ih|d
zo&Q;A3EI^`HHU6A4e^cn|NHmC^h&a2F++`7#au-lrXxZm!DKg+wppC*acx;~O*=vj
z)1{h0jYi_f^6!P|FA)<F4LQYVQxlD#uBkGF!|Zwel<H>bsD8Wv)Q}uY*W73XbKJhD
zxEm6tO$~gc{AqT^!eWMn(TaUltX)Vr|2LP<)wZ=cnvtXqvzRTr%wJGuJJj5bJvNed
zj3T1QpJJ6~c{7L6jVwqit-D%0zn{Dma<h-T9y(c37RYSIVRR5jLH;ljJ0I8EDTEAy
z1qoOKb0GvDCgYfbYj71*P@n;RgbBG1_uv*>gniHgt04@}ClT@tuEYLGBttYjtbp%0
z|GvQ+aSp<5I1GoN6_&v-oQIEK2u{O#SPegM7CwPN*b9rG4nE;L9EJh70V>qNFwVWZ
zFc69oatsegVF%0w3ia?Av+@x7;b3{LEU-10dwJk}i0As(@ipz0Jd|C)r7Y#^vh}?C
d%y@2Ot9TFIMcL(S3m==UK7J%DZsp!8{|18FupIyZ

delta 291
zcmXBKze_?<90l<2*XPSrZoL$#Na%u-5g|DR)hHSixU{@_w;yXkP;2;FYm$U-O$Gf0
zd50VuY_L_AhM;zZpv9KzPzOG6;GFQ$4R77dUSYGcu9p@Hgm^?xs<D53pBapE+B%+!
z@Zse|YSNZdEmeQhO&#4)4xU~9zhe#WYPlp)9Qp~O8@Ps3*o6|z!B3p%3x;5!2{l-Q
zF;EbNk5QryTtOYSK*KD2$B3Sx1?NzKa*S;8Y}J(`pA!%BJ}*0R5O9g}@@{xqdI8VK
pXTVuW8?;5kQ?g+=8_AQL7_Le?U{78QtG_m)SG2!PHh1yb{sr;2KhOXG

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
index 180cf7f62..aab4d3d37 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless_d2v_v2/train.py
@@ -1127,19 +1127,48 @@ def run(rank, world_size, args, wb=None):
     if world_size > 1:
         logging.info("Using DDP")
         model = DDP(model, device_ids=[rank], find_unused_parameters=True)
+    
+    if params.multi_optim:
+        logging.info("Using seperate optimizers over encoder, decoder ...")
+        enc_param = []
+        dec_param = []
+        for n, p in model.named_parameters():
+            name = n.split('.')[1]
+            if name == 'encoder' and 'feature_extractor' not in n:
+                enc_param.append(p)
+            elif 'feature_extractor' not in n:
+                dec_param.append(p)
 
-    parameters_names = []
-    parameters_names.append(
-        [name_param_pair[0] for name_param_pair in model.named_parameters()]
-    )
-    optimizer = ScaledAdam(
-        model.parameters(),
-        lr=params.base_lr,
-        clipping_scale=2.0,
-        parameters_names=parameters_names,
-    )
+        if wb is None:
+            optimizer_enc = Eve(enc_param, lr=params.peak_enc_lr)
+            optimizer_dec = Eve(dec_param, lr=params.peak_dec_lr)
+        else:
+            logging.info('start wandb sweep optimization...')
+            logging.info(wb.config.peak_enc_lr)
+            logging.info(wb.config.peak_dec_lr)
+            optimizer_enc = Eve(enc_param, lr=wb.config.peak_enc_lr)
+            optimizer_dec = Eve(dec_param, lr=wb.config.peak_dec_lr)
 
-    scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs)
+        scheduler_enc = Eden(optimizer_enc, params.lr_batches*params.accum_grads, params.lr_epochs)
+        scheduler_dec = Eden(optimizer_dec, params.lr_batches*params.acuum_grads, params.lr_epochs)
+        optimizer = [optimizer_enc, optimizer_dec]
+        scheduler = [scheduler_enc, scheduler_dec]
+
+
+
+    else:
+        parameters_names = []
+        parameters_names.append(
+            [name_param_pair[0] for name_param_pair in model.named_parameters()]
+        )
+        optimizer = ScaledAdam(
+            model.parameters(),
+            lr=params.base_lr,
+            clipping_scale=2.0,
+            parameters_names=parameters_names,
+        )
+
+        scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs)
 
     if checkpoints and "optimizer" in checkpoints:
         logging.info("Loading optimizer state dict")