clean

speech/cosyvoice/utils/executor.py

@@ -235,8 +235,8 @@ class Executor:
         info_dict["loss_dict"] = total_loss_dict
         log_per_save(writer, info_dict)
         model_name = (
-            "epoch_{}_whole".format(self.epoch)
+            f"epoch_{self.epoch}_whole"
             if on_batch_end
-            else "epoch_{}_step_{}".format(self.epoch, self.step + 1)
+            else f"epoch_{self.epoch}_step_{self.step + 1}"
         )
         save_model(model, model_name, info_dict)

speech/cosyvoice/utils/train_utils.py

@@ -187,7 +187,7 @@ def init_optimizer_and_scheduler(args, configs, model, gan):
 
 
 def init_summarywriter(args):
-    
+    """Init summary writer"""    
     writer = None
     if int(os.environ.get('RANK', 0)) == 0:
         os.makedirs(args.model_dir, exist_ok=True)
@@ -196,6 +196,7 @@ def init_summarywriter(args):
 
 
 def save_model(model, model_name, info_dict):
+    """Save model"""
     rank = int(os.environ.get('RANK', 0))
     model_dir = info_dict["model_dir"]
     save_model_path = os.path.join(model_dir, '{}.pt'.format(model_name))
@@ -280,6 +281,7 @@ def batch_forward(model, batch, scaler, info_dict, ref_model=None, dpo_loss=None
 
 
 def batch_backward(model, scaler, info_dict):
+    """Backward batch"""
     if info_dict["train_engine"] == "deepspeed":
         scaled_loss = model.backward(info_dict['loss_dict']['loss'])    
     else:
@@ -294,6 +296,7 @@ def batch_backward(model, scaler, info_dict):
 
 
 def update_parameter_and_lr(model, optimizer, scheduler, scaler, info_dict):
+    """Update parameters and learning rate"""
     grad_norm = 0.0
     if info_dict['train_engine'] == "deepspeed":
         info_dict["is_gradient_accumulation_boundary"] = model.is_gradient_accumulation_boundary()
@@ -326,6 +329,7 @@ def update_parameter_and_lr(model, optimizer, scheduler, scaler, info_dict):
 
 
 def log_per_step(writer, info_dict):
+    """Log per step"""
     tag = info_dict["tag"]
     epoch = info_dict.get('epoch', 0)
     step = info_dict["step"]
@@ -338,23 +342,23 @@ def log_per_step(writer, info_dict):
         if (info_dict['train_engine'] == 'deepspeed' and info_dict['is_gradient_accumulation_boundary'] is True) or \
            (info_dict['train_engine'] == 'torch_ddp' and (info_dict['batch_idx'] + 1) % info_dict['accum_grad'] == 0):
             for k in ['epoch', 'lr', 'grad_norm']:
-                writer.add_scalar('{}/{}'.format(tag, k), info_dict[k], step + 1)
+                writer.add_scalar(f'{tag}/{k}', info_dict[k], step + 1)
             for k, v in loss_dict.items():
-                writer.add_scalar('{}/{}'.format(tag, k), v, step + 1)
+                writer.add_scalar(f'{tag}/{k}', v, step + 1)
 
     # TRAIN & CV, Shell log (stdout)
     if (info_dict['batch_idx'] + 1) % info_dict['log_interval'] == 0:
-        log_str = '{} Batch {}/{} '.format(tag, epoch, batch_idx + 1)
+        log_str = f'{tag} Batch {epoch}/{batch_idx + 1} '
         for name, value in loss_dict.items():
-            log_str += '{} {:.6f} '.format(name, value)
+            log_str += f'{name} {value:.6f} '
         if tag == "TRAIN":
-            log_str += 'lr {:.8f} grad_norm {:.6f}'.format(
-                info_dict["lr"], info_dict['grad_norm'])
-        log_str += ' rank {}'.format(rank)
+            log_str += f'lr {info_dict["lr"]:.8f} grad_norm {info_dict["grad_norm"]:.6f}'
+        log_str += f' rank {rank}'
         logging.debug(log_str)
 
 
 def log_per_save(writer, info_dict):
+    """Log per save"""
     tag = info_dict["tag"]
     epoch = info_dict["epoch"]
     step = info_dict["step"]
@@ -366,6 +370,6 @@ def log_per_save(writer, info_dict):
 
     if writer is not None:
         for k in ['epoch', 'lr']:
-            writer.add_scalar('{}/{}'.format(tag, k), info_dict[k], step + 1)
+            writer.add_scalar(f'{tag}/{k}', info_dict[k], step + 1)
         for k, v in loss_dict.items():
-            writer.add_scalar('{}/{}'.format(tag, k), v, step + 1)
+            writer.add_scalar(f'{tag}/{k}', v, step + 1)

speech/train.py

@@ -13,82 +13,97 @@
 # limitations under the License.
 
 from __future__ import print_function
+
 import argparse
 import datetime
 import logging
-logging.getLogger('matplotlib').setLevel(logging.WARNING)
-from copy import deepcopy
+
+logging.getLogger("matplotlib").setLevel(logging.WARNING)
 import os
+from copy import deepcopy
+
+import deepspeed
 import torch
 import torch.distributed as dist
-import deepspeed
-from loguru import logger
-
 from hyperpyyaml import load_hyperpyyaml
-
+from loguru import logger
 from torch.distributed.elastic.multiprocessing.errors import record
 
-from cosyvoice.utils.losses import DPOLoss
 from cosyvoice.utils.executor import Executor
-from cosyvoice.utils.train_utils import (
-    init_distributed,
-    init_dataset_and_dataloader,
-    init_optimizer_and_scheduler,
-    init_summarywriter, save_model,
-    check_modify_and_save_config)
+from cosyvoice.utils.losses import DPOLoss
+from cosyvoice.utils.train_utils import (check_modify_and_save_config,
+                                         init_dataset_and_dataloader,
+                                         init_distributed,
+                                         init_optimizer_and_scheduler,
+                                         init_summarywriter, save_model)
 
 
 def get_args():
-    parser = argparse.ArgumentParser(description='training your network')
-    parser.add_argument('--train_engine',
-                        default='torch_ddp',
-                        choices=['torch_ddp', 'deepspeed'],
-                        help='Engine for paralleled training')
-    parser.add_argument('--model', required=True, help='model which will be trained')
-    parser.add_argument('--ref_model', required=False, help='ref model used in dpo')
-    parser.add_argument('--config', required=True, help='config file')
-    parser.add_argument('--train_data', required=True, help='train data file')
-    parser.add_argument('--cv_data', required=True, help='cv data file')
-    parser.add_argument('--qwen_pretrain_path', required=False, help='qwen pretrain path')
-    parser.add_argument('--checkpoint', help='checkpoint model')
-    parser.add_argument('--model_dir', required=True, help='save model dir')
-    parser.add_argument('--tensorboard_dir',
-                        default='tensorboard',
-                        help='tensorboard log dir')
-    parser.add_argument('--ddp.dist_backend',
-                        dest='dist_backend',
-                        default='nccl',
-                        choices=['nccl', 'gloo'],
-                        help='distributed backend')
-    parser.add_argument('--num_workers',
-                        default=0,
-                        type=int,
-                        help='num of subprocess workers for reading')
-    parser.add_argument('--prefetch',
-                        default=100,
-                        type=int,
-                        help='prefetch number')
-    parser.add_argument('--pin_memory',
-                        action='store_true',
-                        default=False,
-                        help='Use pinned memory buffers used for reading')
-    parser.add_argument('--use_amp',
-                        action='store_true',
-                        default=False,
-                        help='Use automatic mixed precision training')
-    parser.add_argument('--dpo',
-                        action='store_true',
-                        default=False,
-                        help='Use Direct Preference Optimization')
-    parser.add_argument('--deepspeed.save_states',
-                        dest='save_states',
-                        default='model_only',
-                        choices=['model_only', 'model+optimizer'],
-                        help='save model/optimizer states')
-    parser.add_argument('--timeout',
-                        default=60,
-                        type=int,
-                        help='timeout (in seconds) of cosyvoice_join.')
+    parser = argparse.ArgumentParser(description="training your network")
+    parser.add_argument(
+        "--train_engine",
+        default="torch_ddp",
+        choices=["torch_ddp", "deepspeed"],
+        help="Engine for paralleled training",
+    )
+    parser.add_argument("--model", required=True, help="model which will be trained")
+    parser.add_argument("--ref_model", required=False, help="ref model used in dpo")
+    parser.add_argument("--config", required=True, help="config file")
+    parser.add_argument("--train_data", required=True, help="train data file")
+    parser.add_argument("--cv_data", required=True, help="cv data file")
+    parser.add_argument(
+        "--qwen_pretrain_path", required=False, help="qwen pretrain path"
+    )
+    parser.add_argument("--checkpoint", help="checkpoint model")
+    parser.add_argument("--model_dir", required=True, help="save model dir")
+    parser.add_argument(
+        "--tensorboard_dir", default="tensorboard", help="tensorboard log dir"
+    )
+    parser.add_argument(
+        "--ddp.dist_backend",
+        dest="dist_backend",
+        default="nccl",
+        choices=["nccl", "gloo"],
+        help="distributed backend",
+    )
+    parser.add_argument(
+        "--num_workers",
+        default=0,
+        type=int,
+        help="num of subprocess workers for reading",
+    )
+    parser.add_argument("--prefetch", default=100, type=int, help="prefetch number")
+    parser.add_argument(
+        "--pin_memory",
+        action="store_true",
+        default=False,
+        help="Use pinned memory buffers used for reading",
+    )
+    parser.add_argument(
+        "--use_amp",
+        action="store_true",
+        default=False,
+        help="Use automatic mixed precision training",
+    )
+    parser.add_argument(
+        "--dpo",
+        action="store_true",
+        default=False,
+        help="Use Direct Preference Optimization",
+    )
+    parser.add_argument(
+        "--deepspeed.save_states",
+        dest="save_states",
+        default="model_only",
+        choices=["model_only", "model+optimizer"],
+        help="save model/optimizer states",
+    )
+    parser.add_argument(
+        "--timeout",
+        default=60,
+        type=int,
+        help="timeout (in seconds) of cosyvoice_join.",
+    )
     parser = deepspeed.add_config_arguments(parser)
     args = parser.parse_args()
     return args
@@ -97,30 +112,41 @@ def get_args():
 @record
 def main():
     args = get_args()
-    logging.basicConfig(level=logging.DEBUG,
-                        format='%(asctime)s %(levelname)s %(message)s')
+    logging.basicConfig(
+        level=logging.DEBUG, format="%(asctime)s %(levelname)s %(message)s"
+    )
     # gan train has some special initialization logic
-    gan = True if args.model == 'hifigan' else False
+    gan = True if args.model == "hifigan" else False
 
-    override_dict = {k: None for k in ['llm', 'flow', 'hift', 'hifigan'] if k != args.model}
+    override_dict = {
+        k: None for k in ["llm", "flow", "hift", "hifigan"] if k != args.model
+    }
     if gan is True:
-        override_dict.pop('hift')
+        override_dict.pop("hift")
     try:
-        with open(args.config, 'r') as f:
-            configs = load_hyperpyyaml(f, overrides={**override_dict, 'qwen_pretrain_path': args.qwen_pretrain_path})
-    except Exception:
-        with open(args.config, 'r') as f:
+        with open(args.config, "r", encoding="utf-8") as f:
+            configs = load_hyperpyyaml(
+                f,
+                overrides={
+                    **override_dict,
+                    "qwen_pretrain_path": args.qwen_pretrain_path,
+                },
+            )
+    except Exception as e:
+        logger.error(f"Error loading config: {e}")
+        with open(args.config, "r", encoding="utf-8") as f:
             configs = load_hyperpyyaml(f, overrides=override_dict)
     if gan is True:
-        configs['train_conf'] = configs['train_conf_gan']
-    configs['train_conf'].update(vars(args))
+        configs["train_conf"] = configs["train_conf_gan"]
+    configs["train_conf"].update(vars(args))
 
     # Init env for ddp
     init_distributed(args)
 
     # Get dataset & dataloader
-    train_dataset, _, train_data_loader, cv_data_loader = \
-        init_dataset_and_dataloader(args, configs, gan, args.dpo)
+    train_dataset, _, train_data_loader, cv_data_loader = init_dataset_and_dataloader(
+        args, configs, gan, args.dpo
+    )
 
     # Do some sanity checks and save config to arsg.model_dir
     configs = check_modify_and_save_config(args, configs)
@@ -136,40 +162,45 @@ def main():
     start_step, start_epoch = 0, -1
     if args.checkpoint is not None:
         if os.path.exists(args.checkpoint):
-            state_dict = torch.load(args.checkpoint, map_location='cpu')
+            state_dict = torch.load(args.checkpoint, map_location="cpu")
             model.load_state_dict(state_dict, strict=False)
-            if 'step' in state_dict:
-                start_step = state_dict['step']
-            if 'epoch' in state_dict:
-                start_epoch = state_dict['epoch']
+            if "step" in state_dict:
+                start_step = state_dict["step"]
+            if "epoch" in state_dict:
+                start_epoch = state_dict["epoch"]
         else:
-            logging.warning('checkpoint {} do not exsist!'.format(args.checkpoint))
+            logger.warning(f"checkpoint {args.checkpoint} do not exsist!")
 
     # Dispatch model from cpu to gpu
     model = model.cuda()
-    model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True)
-   
+    model = torch.nn.parallel.DistributedDataParallel(
+        model, find_unused_parameters=True
+    )
 
     # Get optimizer & scheduler
-    model, optimizer, scheduler, optimizer_d, scheduler_d = init_optimizer_and_scheduler(args, configs, model, gan)
+    model, optimizer, scheduler, optimizer_d, scheduler_d = (
+        init_optimizer_and_scheduler(args, configs, model, gan)
+    )
     scheduler.set_step(start_step)
     if scheduler_d is not None:
         scheduler_d.set_step(start_step)
 
     # Save init checkpoints
-    info_dict = deepcopy(configs['train_conf'])
-    info_dict['step'] = start_step
-    info_dict['epoch'] = start_epoch
-    save_model(model, 'init', info_dict)
+    info_dict = deepcopy(configs["train_conf"])
+    info_dict["step"] = start_step
+    info_dict["epoch"] = start_epoch
+    save_model(model, "init", info_dict)
 
     # DPO related
     if args.dpo is True:
         ref_model = deepcopy(configs[args.model])
-        state_dict = torch.load(args.ref_model, map_location='cpu')
+        state_dict = torch.load(args.ref_model, map_location="cpu")
         ref_model.load_state_dict(state_dict, strict=False)
         dpo_loss = DPOLoss(beta=0.01, label_smoothing=0.0, ipo=False)
         ref_model = ref_model.cuda()
-        ref_model = torch.nn.parallel.DistributedDataParallel(ref_model, find_unused_parameters=True)
+        ref_model = torch.nn.parallel.DistributedDataParallel(
+            ref_model, find_unused_parameters=True
+        )
     else:
         ref_model, dpo_loss = None, None
 
@@ -179,21 +210,44 @@ def main():
 
     # Init scaler, used for pytorch amp mixed precision training
     scaler = torch.amp.GradScaler() if args.use_amp else None
-    logger.info(f'start step {start_step} start epoch {start_epoch}')
+    logger.info(f"start step {start_step} start epoch {start_epoch}")
 
     # Start training loop
-    for epoch in range(start_epoch + 1, info_dict['max_epoch']):
+    for epoch in range(start_epoch + 1, info_dict["max_epoch"]):
         executor.epoch = epoch
         train_dataset.set_epoch(epoch)
         dist.barrier()
-        group_join = dist.new_group(backend="nccl", timeout=datetime.timedelta(seconds=args.timeout))
+        group_join = dist.new_group(
+            backend="nccl", timeout=datetime.timedelta(seconds=args.timeout)
+        )
         if gan is True:
-            executor.train_one_epoc_gan(model, optimizer, scheduler, optimizer_d, scheduler_d, train_data_loader, cv_data_loader,
-                                        writer, info_dict, scaler, group_join)
+            executor.train_one_epoc_gan(
+                model,
+                optimizer,
+                scheduler,
+                optimizer_d,
+                scheduler_d,
+                train_data_loader,
+                cv_data_loader,
+                writer,
+                info_dict,
+                scaler,
+                group_join,
+            )
         else:
-            executor.train_one_epoc(model, optimizer, scheduler, train_data_loader, cv_data_loader, writer, info_dict, scaler, group_join)
+            executor.train_one_epoc(
+                model,
+                optimizer,
+                scheduler,
+                train_data_loader,
+                cv_data_loader,
+                writer,
+                info_dict,
+                scaler,
+                group_join,
+            )
         dist.destroy_process_group(group_join)
 
 
-if __name__ == '__main__':
-    main()
+if __name__ == "__main__":
+    main()