wehub-resource-sync
515 lines
20 KiB
Python
515 lines
20 KiB
Python
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
|