Update infer/modules/train/train.py

infer/modules/train/train.py

@@ -8,6 +8,7 @@ now_dir = os.getcwd()
 sys.path.append(os.path.join(now_dir))
 
 import datetime
+from tqdm import tqdm  # Added import
 
 from infer.lib.train import utils
 
@@ -105,6 +106,7 @@ def main():
     os.environ["MASTER_PORT"] = str(randint(20000, 55555))
     children = []
     logger = utils.get_logger(hps.model_dir)
+    logger.info(f"Starting training with {n_gpus} GPU(s)")
     for i in range(n_gpus):
         subproc = mp.Process(
             target=run,
@@ -120,9 +122,8 @@ def main():
 def run(rank, n_gpus, hps, logger: logging.Logger):
     global global_step
     if rank == 0:
-        # logger = utils.get_logger(hps.model_dir)
+        logger.info(f"Process {rank}/{n_gpus-1} started")
         logger.info(hps)
-        # utils.check_git_hash(hps.model_dir)
         writer = SummaryWriter(log_dir=hps.model_dir)
         writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
 
@@ -140,18 +141,17 @@ def run(rank, n_gpus, hps, logger: logging.Logger):
     train_sampler = DistributedBucketSampler(
         train_dataset,
         hps.train.batch_size * n_gpus,
-        # [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1200,1400],  # 16s
-        [100, 200, 300, 400, 500, 600, 700, 800, 900],  # 16s
+        [100, 200, 300, 400, 500, 600, 700, 800, 900],
         num_replicas=n_gpus,
         rank=rank,
         shuffle=True,
     )
-    # It is possible that dataloader's workers are out of shared memory. Please try to raise your shared memory limit.
-    # num_workers=8 -> num_workers=4
+    
     if hps.if_f0 == 1:
         collate_fn = TextAudioCollateMultiNSFsid()
     else:
         collate_fn = TextAudioCollate()
+    
     train_loader = DataLoader(
         train_dataset,
         num_workers=4,
@@ -162,6 +162,11 @@ def run(rank, n_gpus, hps, logger: logging.Logger):
         persistent_workers=True,
         prefetch_factor=8,
     )
+    
+    if rank == 0:
+        logger.info(f"Training dataset size: {len(train_dataset)}")
+        logger.info(f"Number of batches per epoch: {len(train_loader)}")
+
     if hps.if_f0 == 1:
         net_g = RVC_Model_f0(
             hps.data.filter_length // 2 + 1,
@@ -177,11 +182,14 @@ def run(rank, n_gpus, hps, logger: logging.Logger):
             **hps.model,
             is_half=hps.train.fp16_run,
         )
+    
     if torch.cuda.is_available():
         net_g = net_g.cuda(rank)
+    
     net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm)
     if torch.cuda.is_available():
         net_d = net_d.cuda(rank)
+    
     optim_g = torch.optim.AdamW(
         net_g.parameters(),
         hps.train.learning_rate,
@@ -194,8 +202,7 @@ def run(rank, n_gpus, hps, logger: logging.Logger):
         betas=hps.train.betas,
         eps=hps.train.eps,
     )
-    # net_g = DDP(net_g, device_ids=[rank], find_unused_parameters=True)
-    # net_d = DDP(net_d, device_ids=[rank], find_unused_parameters=True)
+    
     if hasattr(torch, "xpu") and torch.xpu.is_available():
         pass
     elif torch.cuda.is_available():
@@ -205,52 +212,43 @@ def run(rank, n_gpus, hps, logger: logging.Logger):
         net_g = DDP(net_g)
         net_d = DDP(net_d)
 
-    try:  # 如果能加载自动resume
+    try:
         _, _, _, epoch_str = utils.load_checkpoint(
             utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_d
-        )  # D多半加载没事
+        )
         if rank == 0:
-            logger.info("loaded D")
-        # _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g,load_opt=0)
+            logger.info("Loaded discriminator checkpoint")
+        
         _, _, _, epoch_str = utils.load_checkpoint(
             utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g
         )
         global_step = (epoch_str - 1) * len(train_loader)
-        # epoch_str = 1
-        # global_step = 0
-    except:  # 如果首次不能加载，加载pretrain
-        # traceback.print_exc()
+        if rank == 0:
+            logger.info(f"Resuming from epoch {epoch_str}, global step {global_step}")
+    except:
         epoch_str = 1
         global_step = 0
         if hps.pretrainG != "":
             if rank == 0:
-                logger.info("loaded pretrained %s" % (hps.pretrainG))
+                logger.info(f"Loading pretrained generator from {hps.pretrainG}")
             if hasattr(net_g, "module"):
-                logger.info(
-                    net_g.module.load_state_dict(
-                        torch.load(hps.pretrainG, map_location="cpu")["model"]
-                    )
-                )  ##测试不加载优化器
+                net_g.module.load_state_dict(
+                    torch.load(hps.pretrainG, map_location="cpu")["model"]
+                )
             else:
-                logger.info(
-                    net_g.load_state_dict(
-                        torch.load(hps.pretrainG, map_location="cpu")["model"]
-                    )
-                )  ##测试不加载优化器
+                net_g.load_state_dict(
+                    torch.load(hps.pretrainG, map_location="cpu")["model"]
+                )
         if hps.pretrainD != "":
             if rank == 0:
-                logger.info("loaded pretrained %s" % (hps.pretrainD))
+                logger.info(f"Loading pretrained discriminator from {hps.pretrainD}")
             if hasattr(net_d, "module"):
-                logger.info(
-                    net_d.module.load_state_dict(
-                        torch.load(hps.pretrainD, map_location="cpu")["model"]
-                    )
+                net_d.module.load_state_dict(
+                    torch.load(hps.pretrainD, map_location="cpu")["model"]
                 )
             else:
-                logger.info(
-                    net_d.load_state_dict(
-                        torch.load(hps.pretrainD, map_location="cpu")["model"]
-                    )
+                net_d.load_state_dict(
+                    torch.load(hps.pretrainD, map_location="cpu")["model"]
                 )
 
     scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
@@ -263,6 +261,11 @@ def run(rank, n_gpus, hps, logger: logging.Logger):
     scaler = GradScaler(enabled=hps.train.fp16_run)
 
     cache = []
+    
+    if rank == 0:
+        logger.info(f"Starting training from epoch {epoch_str} to {hps.train.epochs}")
+        logger.info(f"Total epochs to train: {hps.train.epochs - epoch_str + 1}")
+    
     for epoch in range(epoch_str, hps.train.epochs + 1):
         if rank == 0:
             train_and_evaluate(
@@ -313,12 +316,16 @@ def train_and_evaluate(
 
     # Prepare data iterator
     if hps.if_cache_data_in_gpu == True:
-        # Use Cache
-        data_iterator = cache
         if cache == []:
-            # Make new cache
+            if rank == 0:
+                logger.info("Caching data in GPU...")
+                cache_progress = tqdm(total=len(train_loader), 
+                                    desc="Caching", 
+                                    position=0, 
+                                    leave=True,
+                                    disable=(rank != 0))
+            
             for batch_idx, info in enumerate(train_loader):
-                # Unpack
                 if hps.if_f0 == 1:
                     (
                         phone,
@@ -341,7 +348,7 @@ def train_and_evaluate(
                         wave_lengths,
                         sid,
                     ) = info
-                # Load on CUDA
+                
                 if torch.cuda.is_available():
                     phone = phone.cuda(rank, non_blocking=True)
                     phone_lengths = phone_lengths.cuda(rank, non_blocking=True)
@@ -352,8 +359,7 @@ def train_and_evaluate(
                     spec = spec.cuda(rank, non_blocking=True)
                     spec_lengths = spec_lengths.cuda(rank, non_blocking=True)
                     wave = wave.cuda(rank, non_blocking=True)
-                    wave_lengths = wave_lengths.cuda(rank, non_blocking=True)
-                # Cache on list
+                
                 if hps.if_f0 == 1:
                     cache.append(
                         (
@@ -386,18 +392,31 @@ def train_and_evaluate(
                             ),
                         )
                     )
-        else:
-            # Load shuffled cache
-            shuffle(cache)
+                
+                if rank == 0:
+                    cache_progress.update(1)
+            
+            if rank == 0:
+                cache_progress.close()
+                logger.info(f"Cached {len(cache)} batches in GPU")
+        
+        shuffle(cache)
+        data_iterator = cache
     else:
-        # Loader
         data_iterator = enumerate(train_loader)
 
-    # Run steps
+    # Initialize tqdm progress bar for training
+    if rank == 0:
+        epoch_progress = tqdm(total=len(train_loader), 
+                            desc=f"Epoch {epoch}/{hps.train.epochs}", 
+                            position=0, 
+                            leave=True,
+                            bar_format='{l_bar}{bar:20}{r_bar}{bar:-20b}')
+    
     epoch_recorder = EpochRecorder()
+    
     for batch_idx, info in data_iterator:
-        # Data
-        ## Unpack
+        # Unpack data
         if hps.if_f0 == 1:
             (
                 phone,
@@ -412,7 +431,7 @@ def train_and_evaluate(
             ) = info
         else:
             phone, phone_lengths, spec, spec_lengths, wave, wave_lengths, sid = info
-        ## Load on CUDA
+        
         if (hps.if_cache_data_in_gpu == False) and torch.cuda.is_available():
             phone = phone.cuda(rank, non_blocking=True)
             phone_lengths = phone_lengths.cuda(rank, non_blocking=True)
@@ -423,9 +442,8 @@ def train_and_evaluate(
             spec = spec.cuda(rank, non_blocking=True)
             spec_lengths = spec_lengths.cuda(rank, non_blocking=True)
             wave = wave.cuda(rank, non_blocking=True)
-            # wave_lengths = wave_lengths.cuda(rank, non_blocking=True)
 
-        # Calculate
+        # Forward pass
         with autocast(enabled=hps.train.fp16_run):
             if hps.if_f0 == 1:
                 (
@@ -443,6 +461,7 @@ def train_and_evaluate(
                     z_mask,
                     (z, z_p, m_p, logs_p, m_q, logs_q),
                 ) = net_g(phone, phone_lengths, spec, spec_lengths, sid)
+            
             mel = spec_to_mel_torch(
                 spec,
                 hps.data.filter_length,
@@ -454,6 +473,7 @@ def train_and_evaluate(
             y_mel = commons.slice_segments(
                 mel, ids_slice, hps.train.segment_size // hps.data.hop_length
             )
+            
             with autocast(enabled=False):
                 y_hat_mel = mel_spectrogram_torch(
                     y_hat.float().squeeze(1),
@@ -465,26 +485,30 @@ def train_and_evaluate(
                     hps.data.mel_fmin,
                     hps.data.mel_fmax,
                 )
+            
             if hps.train.fp16_run == True:
                 y_hat_mel = y_hat_mel.half()
+            
             wave = commons.slice_segments(
                 wave, ids_slice * hps.data.hop_length, hps.train.segment_size
-            )  # slice
+            )
 
-            # Discriminator
+            # Discriminator forward
             y_d_hat_r, y_d_hat_g, _, _ = net_d(wave, y_hat.detach())
             with autocast(enabled=False):
                 loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(
                     y_d_hat_r, y_d_hat_g
                 )
+        
+        # Discriminator backward
         optim_d.zero_grad()
         scaler.scale(loss_disc).backward()
         scaler.unscale_(optim_d)
         grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
         scaler.step(optim_d)
 
+        # Generator forward
         with autocast(enabled=hps.train.fp16_run):
-            # Generator
             y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(wave, y_hat)
             with autocast(enabled=False):
                 loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
@@ -492,6 +516,8 @@ def train_and_evaluate(
                 loss_fm = feature_loss(fmap_r, fmap_g)
                 loss_gen, losses_gen = generator_loss(y_d_hat_g)
                 loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl
+        
+        # Generator backward
         optim_g.zero_grad()
         scaler.scale(loss_gen_all).backward()
         scaler.unscale_(optim_g)
@@ -499,39 +525,43 @@ def train_and_evaluate(
         scaler.step(optim_g)
         scaler.update()
 
+        # Update progress bar and logging
         if rank == 0:
+            if epoch_progress is not None:
+                epoch_progress.update(1)
+                
+                # Update progress bar description with current losses
+                if batch_idx % hps.train.log_interval == 0:
+                    postfix_dict = {
+                        'G': f'{loss_gen_all:.3f}',
+                        'D': f'{loss_disc:.3f}',
+                        'Mel': f'{loss_mel:.3f}',
+                        'KL': f'{loss_kl:.3f}',
+                        'Step': global_step
+                    }
+                    epoch_progress.set_postfix(postfix_dict)
+            
             if global_step % hps.train.log_interval == 0:
                 lr = optim_g.param_groups[0]["lr"]
-                logger.info(
-                    "Train Epoch: {} [{:.0f}%]".format(
-                        epoch, 100.0 * batch_idx / len(train_loader)
-                    )
-                )
-                # Amor For Tensorboard display
-                if loss_mel > 75:
-                    loss_mel = 75
-                if loss_kl > 9:
-                    loss_kl = 9
-
-                logger.info([global_step, lr])
-                logger.info(
-                    f"loss_disc={loss_disc:.3f}, loss_gen={loss_gen:.3f}, loss_fm={loss_fm:.3f},loss_mel={loss_mel:.3f}, loss_kl={loss_kl:.3f}"
-                )
+                
+                logger.info(f"\nEpoch: {epoch} [{batch_idx}/{len(train_loader)}]")
+                logger.info(f"Global Step: {global_step}")
+                logger.info(f"Learning Rate: {lr:.6f}")
+                logger.info(f"Generator Loss: {loss_gen_all:.3f} (FM: {loss_fm:.3f}, Mel: {loss_mel:.3f}, KL: {loss_kl:.3f})")
+                logger.info(f"Discriminator Loss: {loss_disc:.3f}")
+                logger.info(f"Grad Norm - G: {grad_norm_g:.3f}, D: {grad_norm_d:.3f}")
+                
+                # Tensorboard logging
                 scalar_dict = {
                     "loss/g/total": loss_gen_all,
                     "loss/d/total": loss_disc,
                     "learning_rate": lr,
                     "grad_norm_d": grad_norm_d,
                     "grad_norm_g": grad_norm_g,
+                    "loss/g/fm": loss_fm,
+                    "loss/g/mel": loss_mel,
+                    "loss/g/kl": loss_kl,
                 }
-                scalar_dict.update(
-                    {
-                        "loss/g/fm": loss_fm,
-                        "loss/g/mel": loss_mel,
-                        "loss/g/kl": loss_kl,
-                    }
-                )
-
                 scalar_dict.update(
                     {"loss/g/{}".format(i): v for i, v in enumerate(losses_gen)}
                 )
@@ -541,6 +571,7 @@ def train_and_evaluate(
                 scalar_dict.update(
                     {"loss/d_g/{}".format(i): v for i, v in enumerate(losses_disc_g)}
                 )
+                
                 image_dict = {
                     "slice/mel_org": utils.plot_spectrogram_to_numpy(
                         y_mel[0].data.cpu().numpy()
@@ -552,89 +583,100 @@ def train_and_evaluate(
                         mel[0].data.cpu().numpy()
                     ),
                 }
+                
                 utils.summarize(
                     writer=writer,
                     global_step=global_step,
                     images=image_dict,
                     scalars=scalar_dict,
                 )
+        
         global_step += 1
-    # /Run steps
-
+    
+    # Close progress bar
+    if rank == 0 and epoch_progress is not None:
+        epoch_progress.close()
+    
+    # Save checkpoints
     if epoch % hps.save_every_epoch == 0 and rank == 0:
         if hps.if_latest == 0:
+            save_path_g = os.path.join(hps.model_dir, f"G_{global_step}.pth")
+            save_path_d = os.path.join(hps.model_dir, f"D_{global_step}.pth")
             utils.save_checkpoint(
                 net_g,
                 optim_g,
                 hps.train.learning_rate,
                 epoch,
-                os.path.join(hps.model_dir, "G_{}.pth".format(global_step)),
+                save_path_g,
             )
             utils.save_checkpoint(
                 net_d,
                 optim_d,
                 hps.train.learning_rate,
                 epoch,
-                os.path.join(hps.model_dir, "D_{}.pth".format(global_step)),
+                save_path_d,
             )
+            logger.info(f"Saved checkpoints: {save_path_g}, {save_path_d}")
         else:
+            save_path_g = os.path.join(hps.model_dir, "G_2333333.pth")
+            save_path_d = os.path.join(hps.model_dir, "D_2333333.pth")
             utils.save_checkpoint(
                 net_g,
                 optim_g,
                 hps.train.learning_rate,
                 epoch,
-                os.path.join(hps.model_dir, "G_{}.pth".format(2333333)),
+                save_path_g,
             )
             utils.save_checkpoint(
                 net_d,
                 optim_d,
                 hps.train.learning_rate,
                 epoch,
-                os.path.join(hps.model_dir, "D_{}.pth".format(2333333)),
+                save_path_d,
             )
+            logger.info(f"Saved latest checkpoints: {save_path_g}, {save_path_d}")
+        
         if rank == 0 and hps.save_every_weights == "1":
             if hasattr(net_g, "module"):
                 ckpt = net_g.module.state_dict()
             else:
                 ckpt = net_g.state_dict()
-            logger.info(
-                "saving ckpt %s_e%s:%s"
-                % (
-                    hps.name,
-                    epoch,
-                    savee(
-                        ckpt,
-                        hps.sample_rate,
-                        hps.if_f0,
-                        hps.name + "_e%s_s%s" % (epoch, global_step),
-                        epoch,
-                        hps.version,
-                        hps,
-                    ),
-                )
+            
+            model_name = hps.name + f"_e{epoch}_s{global_step}"
+            save_result = savee(
+                ckpt,
+                hps.sample_rate,
+                hps.if_f0,
+                model_name,
+                epoch,
+                hps.version,
+                hps,
             )
+            logger.info(f"Saved weights checkpoint: {model_name}: {save_result}")
 
+    # Log epoch completion
     if rank == 0:
-        logger.info("====> Epoch: {} {}".format(epoch, epoch_recorder.record()))
+        logger.info(f"Completed Epoch {epoch} {epoch_recorder.record()}")
+        logger.info(f"Global Step: {global_step}")
+    
+    # End training if completed
     if epoch >= hps.total_epoch and rank == 0:
-        logger.info("Training is done. The program is closed.")
-
+        logger.info("Training completed!")
+        
         if hasattr(net_g, "module"):
             ckpt = net_g.module.state_dict()
         else:
             ckpt = net_g.state_dict()
-        logger.info(
-            "saving final ckpt:%s"
-            % (
-                savee(
-                    ckpt, hps.sample_rate, hps.if_f0, hps.name, epoch, hps.version, hps
-                )
-            )
+        
+        final_save = savee(
+            ckpt, hps.sample_rate, hps.if_f0, hps.name, epoch, hps.version, hps
         )
-        sleep(1)
-        os._exit(2333333)
+        logger.info(f"Saved final model: {final_save}")
+        
+        sleep(2)  # Give time for final logging
+        os._exit(0)
 
 
 if __name__ == "__main__":
     torch.multiprocessing.set_start_method("spawn")
-    main()
+    main()