add generator and discriminator log info

wenet-e2e · Jan 9, 2024 · 89a4a5d · 89a4a5d
1 parent 25c60ba
commit 89a4a5d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/wenet/tts/vits/models.py b/wenet/tts/vits/models.py
@@ -885,7 +885,6 @@ def forward(self, batch: dict, device: torch.device):
                 'loss_dur': loss_dur,
                 'loss_kl': loss_kl,
             }
-        print('optimizer_idx', optimizer_idx)
         return losses
 
     def infer(self, text: torch.Tensor):

diff --git a/wenet/utils/executor.py b/wenet/utils/executor.py
@@ -70,6 +70,7 @@ def train(self, model, optimizer, scheduler, train_data_loader,
                 else:
                     context = nullcontext
                 num_opt = len(optimizer) if isinstance(optimizer, list) else 1
+                loss_dict = {}
                 for opt_idx in range(num_opt):
                     batch_dict['optimizer_idx'] = opt_idx
                     with context():
@@ -84,6 +85,8 @@ def train(self, model, optimizer, scheduler, train_data_loader,
                         scaler,
                         info_dict,
                     )
+                    loss_dict.update(info_dict['loss_dict'])
+                info_dict['loss_dict'] = loss_dict
                 save_interval = info_dict.get('save_interval', 10000)
                 if self.step % save_interval == 0 and self.step != 0 \
                         and (batch_idx + 1) % info_dict["accum_grad"] == 0: