语帆
2024-02-28 eb92e79fb94e7b3df8f27c8ce3e607a70dff2a2e
funasr/train_utils/trainer.py
@@ -108,7 +108,7 @@
        filename = os.path.join(self.output_dir, f'model.pt.ep{epoch}')
        torch.save(state, filename)
        
        print(f'Checkpoint saved to {filename}')
        print(f'\nCheckpoint saved to {filename}\n')
        latest = Path(os.path.join(self.output_dir, f'model.pt'))
        torch.save(state, latest)
@@ -157,7 +157,7 @@
            self._resume_checkpoint(self.output_dir)
        
        for epoch in range(self.start_epoch, self.max_epoch + 1):
            time1 = time.perf_counter()
            self._train_epoch(epoch)
@@ -179,6 +179,9 @@
            
            self.scheduler.step()
            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")
        if self.rank == 0:
            average_checkpoints(self.output_dir, self.avg_nbest_model)
@@ -285,7 +288,7 @@
                                             )
                lr = self.scheduler.get_last_lr()[0]
                time_now = datetime.now()
                time_now = now.strftime("%Y-%m-%d %H:%M:%S")
                time_now = time_now.strftime("%Y-%m-%d %H:%M:%S")
                description = (
                    f"{time_now}, "
                    f"rank: {self.local_rank}, "