Create ddp_train.py

ddp_train.py

@@ -0,0 +1,238 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.tensorboard import SummaryWriter
+import matplotlib.pyplot as plt
+import numpy as np
+import time
+from tqdm import tqdm
+import os.path as osp
+import re
+import sys
+import yaml
+import shutil
+from utils import *
+from optimizers import build_optimizer
+from model import *
+from meldataset import build_dataloader
+from utils import *
+from torch.utils.tensorboard import SummaryWriter
+import click
+
+from accelerate import Accelerator
+from accelerate.utils import LoggerType
+from accelerate import DistributedDataParallelKwargs
+
+import logging
+from logging import StreamHandler
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+handler = StreamHandler()
+handler.setLevel(logging.DEBUG)
+logger.addHandler(handler)
+
+
+import logging
+from accelerate.logging import get_logger
+logger = get_logger(__name__, log_level="DEBUG")
+
+# torch.autograd.detect_anomaly(True)
+torch.backends.cudnn.benchmark = True
+
+@click.command()
+@click.option('-p', '--config_path', default='./Configs/config.yml', type=str)
+def main(config_path):
+
+  config = yaml.safe_load(open(config_path))
+  log_dir = config['log_dir']
+  if not osp.exists(log_dir): os.mkdir(log_dir)
+  shutil.copy(config_path, osp.join(log_dir, osp.basename(config_path)))
+
+  writer = SummaryWriter(log_dir + "/tensorboard")
+  
+  ddp_kwargs = DistributedDataParallelKwargs()
+  accelerator = Accelerator(project_dir=log_dir, split_batches=True, kwargs_handlers=[ddp_kwargs], mixed_precision='bf16')    
+  if accelerator.is_main_process:
+      writer = SummaryWriter(log_dir + "/tensorboard")
+
+
+  # write logs
+  file_handler = logging.FileHandler(osp.join(log_dir, 'train.log'))
+  file_handler.setLevel(logging.DEBUG)
+  file_handler.setFormatter(logging.Formatter('%(levelname)s:%(asctime)s: %(message)s'))
+  logger.logger.addHandler(file_handler)
+
+  epoch = config.get('epoch', 100)
+  save_iter = 1
+  batch_size = config.get('batch_size', 4)
+  log_interval = 1
+  device = accelerator.device
+  train_path = config.get('train_data', None)
+  val_path = config.get('val_data', None)
+  epochs = config.get('epochs', 1000)
+
+  train_list, val_list = get_data_path_list(train_path, val_path)
+  train_list = val_list
+  train_dataloader = build_dataloader(train_list,
+                                      batch_size=batch_size,
+                                      num_workers=8,
+                                      dataset_config=config.get('dataset_params', {}),
+                                      device=device)
+
+  val_dataloader = build_dataloader(val_list,
+                                    batch_size=batch_size,
+                                    validation=True,
+                                    num_workers=2,
+                                    device=device,
+                                    dataset_config=config.get('dataset_params', {}))
+  
+
+
+  aligner = AlignerModel()
+  forward_sum_loss = ForwardSumLoss()
+  best_val_loss = float('inf')
+
+
+  scheduler_params = {
+          "max_lr": float(config['optimizer_params'].get('lr', 5e-4)),
+          "pct_start": float(config['optimizer_params'].get('pct_start', 0.0)),
+          "epochs": epochs,
+          "steps_per_epoch": len(train_dataloader),
+      }
+
+
+  optimizer, scheduler = build_optimizer(
+      {"params": aligner.parameters(), "optimizer_params":{}, "scheduler_params": scheduler_params})
+
+  
+  aligner, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
+      aligner, optimizer, train_dataloader, val_dataloader, scheduler
+  )
+
+  with accelerator.main_process_first():
+      if config.get('pretrained_model', '') != '':
+          model, optimizer, start_epoch, iters = load_checkpoint(model,  optimizer, config['pretrained_model'],
+                                      load_only_params=config.get('load_only_params', True))
+      else:
+          start_epoch = 0
+          iters = 0
+  
+  
+  # Training loop
+  for epoch in range(1, epochs + 1):
+      aligner.train()
+      train_losses = []
+      train_fwd_losses = []
+      start_time = time.time()
+      
+      
+      # Training phase
+      pbar = tqdm(train_dataloader, desc=f"Epoch {epoch}/{epochs} [Train]")
+      for i, batch in enumerate(pbar):
+          batch = [b.to(device) for b in batch]
+
+          text_input, text_input_length, mel_input, mel_input_length, attn_prior = batch
+          
+          # Forward pass
+          attn_soft, attn_logprob = aligner(spec=mel_input, 
+                                            spec_len=mel_input_length, 
+                                            text=text_input, 
+                                            text_len=text_input_length,
+                                            attn_prior=attn_prior)
+          
+          # Calculate loss
+          loss = forward_sum_loss(attn_logprob=attn_logprob, 
+                                  in_lens=text_input_length, 
+                                  out_lens=mel_input_length)
+  
+          # Backward pass and optimization
+          optimizer.zero_grad()
+          accelerator.backward(loss)
+          
+          # Optional gradient clipping
+          grad_norm = accelerator.clip_grad_norm_(aligner.parameters(), 5.0)
+          
+          optimizer.step()
+          iters = iters + 1 
+
+          if scheduler is not None:
+              scheduler.step()
+          
+
+        if (i+1)%log_interval == 0 and accelerator.is_main_process:
+            log_print('Epoch [%d/%d], Step [%d/%d], Loss: %.5f, Forward Sum Loss: %.5f'
+                    %(epoch+1, epochs, i+1, len(train_list)//batch_size, loss), logger)
+            
+            writer.add_scalar('train/Forward Sum Loss', loss, iters)
+            # writer.add_scalar('train/d_loss', d_loss, iters)
+
+            train_losses.append(loss.item())
+            train_fwd_losses.append(loss.item())
+
+            running_loss = 0
+            
+            accelerator.print('Time elasped:', time.time()-start_time)
+
+      # Calculate average training loss for this epoch
+      avg_train_loss = sum(train_losses) / len(train_losses)
+      
+      # Validation phase
+      aligner.eval()
+      val_losses = []
+
+      with torch.no_grad():
+          for batch in tqdm(val_dataloader, desc=f"Epoch {epoch}/{epochs} [Val]"):
+              batch = [b.to(device) for b in batch]
+              
+              text_input, text_input_length, mel_input, mel_input_length, attn_prior = batch
+              
+              # Forward pass
+              attn_soft, attn_logprob = aligner(spec=mel_input, 
+                                              spec_len=mel_input_length, 
+                                              text=text_input, 
+                                              text_len=text_input_length,
+                                              attn_prior=attn_prior)
+              
+              # Calculate loss
+              val_loss = forward_sum_loss(attn_logprob=attn_logprob, 
+                                        in_lens=text_input_length, 
+                                        out_lens=mel_input_length)
+              
+              val_losses.append(val_loss.item())
+      
+        # Calculate average validation loss
+        avg_val_loss = sum(val_losses) / len(val_losses)
+        
+        # Log to TensorBoard
+        writer.add_scalar('epoch/train_loss', avg_train_loss, epoch)
+        writer.add_scalar('epoch/val_loss', avg_val_loss, epoch)
+        
+        # Save checkpoint every N epochs
+            
+        if (i+1)%save_iter == 0 and accelerator.is_main_process:
+
+            print(f'Saving on step {epoch*len(train_dataloader)+i}...')
+            state = {
+                'net':  {key: aligner[key].state_dict() for key in model}, 
+                'optimizer': optimizer.state_dict(),
+                'iters': iters,
+                'epoch': epoch,
+            }
+            save_path = os.path.join(log_dir, 'checkpoints', f'TextAligner_checkpoint_epoch_{epoch}.pt')
+            torch.save(state, save_path)    
+        # Print summary for this epoch
+        epoch_time = time.time() - start_time
+        accelerator.print(f"Epoch {epoch}/{epochs} completed in {epoch_time:.2f}s | "
+              f"Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
+        
+        # # Plot and save attention matrices for visualization
+        # if epoch % config.get('plot_every', 10) == 0:
+        #     plot_attention_matrices(aligner, val_dataloader, device, 
+        #                           os.path.join(log_dir, 'attention_plots', f'epoch_{epoch}'),
+        #                           num_samples=4)
+    
+    writer.close()
+
+if __name__=="__main__":
+    main()