From cdca62d933c4e0766a05044c6cba7cfa0596e615 Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期三, 21 二月 2024 19:22:59 +0800
Subject: [PATCH] Dev gzf (#1377)

---
 funasr/train_utils/trainer.py |   20 ++++++++++----------
 1 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py
index c232642..6a59f91 100644
--- a/funasr/train_utils/trainer.py
+++ b/funasr/train_utils/trainer.py
@@ -181,7 +181,7 @@
 
             time2 = time.perf_counter()
             time_escaped = (time2 - time1)/3600.0
-            print(f"\ntime_escaped_epoch: {time_escaped:.3f} hours, estimated to finish {self.max_epoch} epoch: {(self.max_epoch-epoch)*time_escaped:.3f}\n")
+            print(f"\nrank: {self.local_rank}, time_escaped_epoch: {time_escaped:.3f} hours, estimated to finish {self.max_epoch} epoch: {(self.max_epoch-epoch)*time_escaped:.3f}\n")
 
         if self.rank == 0:
             average_checkpoints(self.output_dir, self.avg_nbest_model)
@@ -302,17 +302,14 @@
                 )
                 pbar.set_description(description)
                 if self.writer:
-                    self.writer.add_scalar(f'rank{self.local_rank}_Loss/train', loss.item(),
-                                           epoch*len(self.dataloader_train) + batch_idx)
+                    self.writer.add_scalar(f'rank{self.local_rank}_Loss/train', loss.item(), self.batch_total)
+                    self.writer.add_scalar(f'rank{self.local_rank}_lr/train', lr, self.batch_total)
                     for key, var in stats.items():
-                        self.writer.add_scalar(f'rank{self.local_rank}_{key}/train', var.item(),
-                                               epoch * len(self.dataloader_train) + batch_idx)
+                        self.writer.add_scalar(f'rank{self.local_rank}_{key}/train', var.item(), self.batch_total)
                     for key, var in speed_stats.items():
-                        self.writer.add_scalar(f'rank{self.local_rank}_{key}/train', eval(var),
-                                               epoch * len(self.dataloader_train) + batch_idx)
-                    
-            # if batch_idx == 2:
-            #     break
+                        self.writer.add_scalar(f'rank{self.local_rank}_{key}/train', eval(var), self.batch_total)
+
+
         pbar.close()
 
     def _validate_epoch(self, epoch):
@@ -356,7 +353,10 @@
                 
                 if (batch_idx+1) % self.log_interval == 0 or (batch_idx+1) == len(self.dataloader_val):
                     pbar.update(self.log_interval)
+                    time_now = datetime.now()
+                    time_now = time_now.strftime("%Y-%m-%d %H:%M:%S")
                     description = (
+                        f"{time_now}, "
                         f"rank: {self.local_rank}, "
                         f"validation epoch: {epoch}/{self.max_epoch}, "
                         f"step: {batch_idx+1}/{len(self.dataloader_val)}, "

--
Gitblit v1.9.1