import logging import os import pathlib import shutil import sys from typing import Dict import matplotlib import utils matplotlib.use('Agg') import torch.utils.data from torchmetrics import Metric, MeanMetric import lightning.pytorch as pl from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_info, rank_zero_only from basics.base_module import CategorizedModule from utils.hparams import hparams from utils.training_utils import ( DsModelCheckpoint, DsTQDMProgressBar, DsBatchSampler, DsTensorBoardLogger, get_latest_checkpoint_path, get_strategy ) from utils.phoneme_utils import load_phoneme_dictionary torch.multiprocessing.set_sharing_strategy(os.getenv('TORCH_SHARE_STRATEGY', 'file_system')) log_format = '%(asctime)s %(message)s' logging.basicConfig(stream=sys.stdout, level=logging.INFO, format=log_format, datefmt='%m/%d %I:%M:%S %p') class BaseTask(pl.LightningModule): """ Base class for training tasks. 1. *load_ckpt*: load checkpoint; 2. *training_step*: record and log the loss; 3. *optimizer_step*: run backwards step; 4. *start*: load training configs, backup code, log to tensorboard, start training; 5. *configure_ddp* and *init_ddp_connection*: start parallel training. Subclasses should define: 1. *build_model*, *build_optimizer*, *build_scheduler*: how to build the model, the optimizer and the training scheduler; 2. *_training_step*: one training step of the model; 3. *on_validation_end* and *_on_validation_end*: postprocess the validation output. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.max_batch_frames = hparams['max_batch_frames'] self.max_batch_size = hparams['max_batch_size'] self.max_val_batch_frames = hparams['max_val_batch_frames'] if self.max_val_batch_frames == -1: hparams['max_val_batch_frames'] = self.max_val_batch_frames = self.max_batch_frames self.max_val_batch_size = hparams['max_val_batch_size'] if self.max_val_batch_size == -1: hparams['max_val_batch_size'] = self.max_val_batch_size = self.max_batch_size self.training_sampler = None self.skip_immediate_validation = False self.skip_immediate_ckpt_save = False self.phoneme_dictionary = load_phoneme_dictionary() self.build_model() self.valid_losses: Dict[str, Metric] = {} self.valid_metrics: Dict[str, Metric] = {} def _finish_init(self): self.register_validation_loss('total_loss') self.build_losses_and_metrics() assert len(self.valid_losses) > 0, "No validation loss registered. Please check your configuration file." ########### # Training, validation and testing ########### def setup(self, stage): self.train_dataset = self.dataset_cls('train') self.valid_dataset = self.dataset_cls('valid') self.num_replicas = (self.trainer.distributed_sampler_kwargs or {}).get('num_replicas', 1) def get_need_freeze_state_dict_key(self, model_state_dict) -> list: key_list = [] for i in hparams['frozen_params']: for j in model_state_dict: if j.startswith(i): key_list.append(j) return list(set(key_list)) def freeze_params(self) -> None: model_state_dict = self.state_dict().keys() freeze_key = self.get_need_freeze_state_dict_key(model_state_dict=model_state_dict) for i in freeze_key: params=self.get_parameter(i) params.requires_grad = False def unfreeze_all_params(self) -> None: for i in self.model.parameters(): i.requires_grad = True def load_finetune_ckpt( self, state_dict ) -> None: adapt_shapes = hparams['finetune_strict_shapes'] if not adapt_shapes: cur_model_state_dict = self.state_dict() unmatched_keys = [] for key, param in state_dict.items(): if key in cur_model_state_dict: new_param = cur_model_state_dict[key] if new_param.shape != param.shape: unmatched_keys.append(key) print('| Unmatched keys: ', key, new_param.shape, param.shape) for key in unmatched_keys: del state_dict[key] self.load_state_dict(state_dict, strict=False) def load_pre_train_model(self): pre_train_ckpt_path = hparams['finetune_ckpt_path'] blacklist = hparams['finetune_ignored_params'] # whitelist=hparams['pre_train_whitelist'] if blacklist is None: blacklist = [] # if whitelist is None: # raise RuntimeError("") if pre_train_ckpt_path is not None: ckpt = torch.load(pre_train_ckpt_path) # if ckpt.get('category') is None: # raise RuntimeError("") if isinstance(self.model, CategorizedModule): self.model.check_category(ckpt.get('category')) state_dict = {} for i in ckpt['state_dict']: # if 'diffusion' in i: # if i in rrrr: # continue skip = False for b in blacklist: if i.startswith(b): skip = True break if skip: continue state_dict[i] = ckpt['state_dict'][i] print(i) return state_dict else: raise RuntimeError("") def _build_model(self): raise NotImplementedError() def build_model(self): self.model = self._build_model() # utils.load_warp(self) self.unfreeze_all_params() if hparams['freezing_enabled']: self.freeze_params() if hparams['finetune_enabled'] and get_latest_checkpoint_path(pathlib.Path(hparams['work_dir'])) is None: self.load_finetune_ckpt(self.load_pre_train_model()) self.print_arch() @rank_zero_only def print_arch(self): utils.print_arch(self.model) def build_losses_and_metrics(self): raise NotImplementedError() def register_validation_metric(self, name: str, metric: Metric): assert isinstance(metric, Metric) self.valid_metrics[name] = metric def register_validation_loss(self, name: str, Aggregator: Metric = MeanMetric): assert issubclass(Aggregator, Metric) self.valid_losses[name] = Aggregator() def run_model(self, sample, infer=False): """ steps: 1. run the full model 2. calculate losses if not infer """ raise NotImplementedError() def on_train_epoch_start(self): if self.training_sampler is not None: self.training_sampler.set_epoch(self.current_epoch) def _training_step(self, sample): """ :return: total loss: torch.Tensor, loss_log: dict, other_log: dict """ losses = self.run_model(sample) total_loss = sum(losses.values()) return total_loss, {**losses, 'batch_size': float(sample['size'])} def training_step(self, sample, batch_idx): total_loss, log_outputs = self._training_step(sample) # logs to progress bar self.log_dict(log_outputs, prog_bar=True, logger=False, on_step=True, on_epoch=False) self.log('lr', self.lr_schedulers().get_last_lr()[0], prog_bar=True, logger=False, on_step=True, on_epoch=False) # logs to tensorboard if self.global_step % hparams['log_interval'] == 0: tb_log = {f'training/{k}': v for k, v in log_outputs.items()} tb_log['training/lr'] = self.lr_schedulers().get_last_lr()[0] self.logger.log_metrics(tb_log, step=self.global_step) return total_loss # def on_before_optimizer_step(self, *args, **kwargs): # self.log_dict(grad_norm(self, norm_type=2)) def _on_validation_start(self): pass def on_validation_start(self): if self.skip_immediate_validation: rank_zero_debug("Skip validation") return self._on_validation_start() for metric in self.valid_losses.values(): metric.to(self.device) metric.reset() for metric in self.valid_metrics.values(): metric.to(self.device) metric.reset() def _validation_step(self, sample, batch_idx): """ :param sample: :param batch_idx: :return: loss_log: dict, weight: int """ raise NotImplementedError() def validation_step(self, sample, batch_idx): """ :param sample: :param batch_idx: """ if self.skip_immediate_validation: rank_zero_debug("Skip validation") return if sample['size'] > 0: with torch.autocast(self.device.type, enabled=False): losses, weight = self._validation_step(sample, batch_idx) losses = { 'total_loss': sum(losses.values()), **losses } for k, v in losses.items(): self.valid_losses[k].update(v, weight=weight) def _on_validation_epoch_end(self): pass def on_validation_epoch_end(self): if self.skip_immediate_validation: self.skip_immediate_validation = False self.skip_immediate_ckpt_save = True return self._on_validation_epoch_end() loss_vals = {k: v.compute() for k, v in self.valid_losses.items()} metric_vals = {k: v.compute() for k, v in self.valid_metrics.items()} self.log('val_loss', loss_vals['total_loss'], on_epoch=True, prog_bar=True, logger=False, sync_dist=True) self.logger.log_metrics({f'validation/{k}': v for k, v in loss_vals.items()}, step=self.global_step) self.logger.log_metrics({f'metrics/{k}': v for k, v in metric_vals.items()}, step=self.global_step) # noinspection PyMethodMayBeStatic def build_scheduler(self, optimizer): from utils import build_lr_scheduler_from_config scheduler_args = hparams['lr_scheduler_args'] assert scheduler_args['scheduler_cls'] != '' scheduler = build_lr_scheduler_from_config(optimizer, scheduler_args) return scheduler # noinspection PyMethodMayBeStatic def build_optimizer(self, model): from utils import build_object_from_class_name optimizer_args = hparams['optimizer_args'] assert optimizer_args['optimizer_cls'] != '' if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args: optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2']) optimizer = build_object_from_class_name( optimizer_args['optimizer_cls'], torch.optim.Optimizer, model if optimizer_args['optimizer_cls'] == 'modules.optimizer.muon.Muon_AdamW' else model.parameters(), **optimizer_args ) return optimizer def configure_optimizers(self): optm = self.build_optimizer(self.model) scheduler = self.build_scheduler(optm) if scheduler is None: return optm return { "optimizer": optm, "lr_scheduler": { "scheduler": scheduler, "interval": "step", "frequency": 1 } } def train_dataloader(self): self.training_sampler = DsBatchSampler( self.train_dataset, max_batch_frames=self.max_batch_frames, max_batch_size=self.max_batch_size, num_replicas=self.num_replicas, rank=self.global_rank, sort_by_similar_size=hparams['sort_by_len'], size_reversed=True, required_batch_count_multiple=hparams['accumulate_grad_batches'], shuffle_sample=True, shuffle_batch=True ) return torch.utils.data.DataLoader( self.train_dataset, collate_fn=self.train_dataset.collater, batch_sampler=self.training_sampler, num_workers=hparams['ds_workers'], prefetch_factor=hparams['dataloader_prefetch_factor'], pin_memory=True, persistent_workers=True ) def val_dataloader(self): sampler = DsBatchSampler( self.valid_dataset, max_batch_frames=self.max_val_batch_frames, max_batch_size=self.max_val_batch_size, num_replicas=self.num_replicas, rank=self.global_rank, shuffle_sample=False, shuffle_batch=False, disallow_empty_batch=False, pad_batch_assignment=False ) return torch.utils.data.DataLoader( self.valid_dataset, collate_fn=self.valid_dataset.collater, batch_sampler=sampler, num_workers=hparams['ds_workers'], prefetch_factor=hparams['dataloader_prefetch_factor'], persistent_workers=True ) def test_dataloader(self): return self.val_dataloader() def on_test_start(self): self.on_validation_start() def test_step(self, sample, batch_idx): return self.validation_step(sample, batch_idx) def on_test_end(self): return self.on_validation_end() ########### # Running configuration ########### @classmethod def start(cls): task = cls() # if pre_train is not None: # task.load_state_dict(pre_train,strict=False) # print("load success-------------------------------------------------------------------") work_dir = pathlib.Path(hparams['work_dir']) trainer = pl.Trainer( accelerator=hparams['pl_trainer_accelerator'], devices=hparams['pl_trainer_devices'], num_nodes=hparams['pl_trainer_num_nodes'], strategy=get_strategy( hparams['pl_trainer_devices'], hparams['pl_trainer_num_nodes'], hparams['pl_trainer_accelerator'], hparams['pl_trainer_strategy'], hparams['pl_trainer_precision'], ), precision=hparams['pl_trainer_precision'], callbacks=[ DsModelCheckpoint( dirpath=work_dir, filename='model_ckpt_steps_{step}', auto_insert_metric_name=False, monitor='step', mode='max', save_last=False, # every_n_train_steps=hparams['val_check_interval'], save_top_k=hparams['num_ckpt_keep'], permanent_ckpt_start=hparams['permanent_ckpt_start'], permanent_ckpt_interval=hparams['permanent_ckpt_interval'], verbose=True ), # LearningRateMonitor(logging_interval='step'), DsTQDMProgressBar(), ], logger=DsTensorBoardLogger( save_dir=str(work_dir), name='lightning_logs', version='latest' ), gradient_clip_val=hparams['clip_grad_norm'], val_check_interval=hparams['val_check_interval'] * hparams['accumulate_grad_batches'], # so this is global_steps check_val_every_n_epoch=None, log_every_n_steps=1, max_steps=hparams['max_updates'], use_distributed_sampler=False, num_sanity_val_steps=hparams['num_sanity_val_steps'], accumulate_grad_batches=hparams['accumulate_grad_batches'] ) if not hparams['infer']: # train @rank_zero_only def train_payload_copy(): # Copy files to work_dir binary_dir = pathlib.Path(hparams['binary_data_dir']) spk_map_dst = work_dir / 'spk_map.json' spk_map_src = binary_dir / 'spk_map.json' shutil.copy(spk_map_src, spk_map_dst) print(f'| Copied spk map to {spk_map_dst}.') lang_map_dst = work_dir / 'lang_map.json' lang_map_src = binary_dir / 'lang_map.json' shutil.copy(lang_map_src, lang_map_dst) print(f'| Copied lang map to {lang_map_dst}.') for lang in hparams['dictionaries'].keys(): dict_dst = work_dir / f'dictionary-{lang}.txt' dict_src = binary_dir / f'dictionary-{lang}.txt' shutil.copy(dict_src, dict_dst) print(f'| Copied dictionary for language \'{lang}\' to {dict_dst}.') train_payload_copy() trainer.fit(task, ckpt_path=get_latest_checkpoint_path(work_dir)) else: trainer.test(task) def on_save_checkpoint(self, checkpoint): if isinstance(self.model, CategorizedModule): checkpoint['category'] = self.model.category checkpoint['trainer_stage'] = self.trainer.state.stage.value def on_load_checkpoint(self, checkpoint): from lightning.pytorch.trainer.states import RunningStage from utils import simulate_lr_scheduler if checkpoint.get('trainer_stage', '') == RunningStage.VALIDATING.value: self.skip_immediate_validation = True optimizer_args = hparams['optimizer_args'] scheduler_args = hparams['lr_scheduler_args'] if 'beta1' in optimizer_args and 'beta2' in optimizer_args and 'betas' not in optimizer_args: optimizer_args['betas'] = (optimizer_args['beta1'], optimizer_args['beta2']) if checkpoint.get('optimizer_states', None): opt_states = checkpoint['optimizer_states'] assert len(opt_states) == 1 # only support one optimizer opt_state = opt_states[0] for param_group in opt_state['param_groups']: for k, v in optimizer_args.items(): if k in param_group and param_group[k] != v: if 'lr_schedulers' in checkpoint and checkpoint['lr_schedulers'] and k == 'lr': continue rank_zero_info(f'| Overriding optimizer parameter {k} from checkpoint: {param_group[k]} -> {v}') param_group[k] = v if 'initial_lr' in param_group and param_group['initial_lr'] != optimizer_args['lr']: rank_zero_info( f'| Overriding optimizer parameter initial_lr from checkpoint: {param_group["initial_lr"]} -> {optimizer_args["lr"]}' ) param_group['initial_lr'] = optimizer_args['lr'] if checkpoint.get('lr_schedulers', None): assert checkpoint.get('optimizer_states', False) assert len(checkpoint['lr_schedulers']) == 1 # only support one scheduler checkpoint['lr_schedulers'][0] = simulate_lr_scheduler( optimizer_args, scheduler_args, step_count=checkpoint['global_step'], num_param_groups=len(checkpoint['optimizer_states'][0]['param_groups']) ) for param_group, new_lr in zip( checkpoint['optimizer_states'][0]['param_groups'], checkpoint['lr_schedulers'][0]['_last_lr'], ): if param_group['lr'] != new_lr: rank_zero_info(f'| Overriding optimizer parameter lr from checkpoint: {param_group["lr"]} -> {new_lr}') param_group['lr'] = new_lr