349 lines
13 KiB
Python
349 lines
13 KiB
Python
# Copyright (c) 2023 Amphion.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import collections
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch.nn.parallel import DistributedDataParallel
|
|
from torch.utils.data import ConcatDataset, DataLoader
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
from models.base.base_sampler import BatchSampler
|
|
from utils.util import (
|
|
Logger,
|
|
remove_older_ckpt,
|
|
save_config,
|
|
set_all_random_seed,
|
|
ValueWindow,
|
|
)
|
|
|
|
|
|
class BaseTrainer(object):
|
|
def __init__(self, args, cfg):
|
|
self.args = args
|
|
self.log_dir = args.log_dir
|
|
self.cfg = cfg
|
|
|
|
self.checkpoint_dir = os.path.join(args.log_dir, "checkpoints")
|
|
os.makedirs(self.checkpoint_dir, exist_ok=True)
|
|
if not cfg.train.ddp or args.local_rank == 0:
|
|
self.sw = SummaryWriter(os.path.join(args.log_dir, "events"))
|
|
self.logger = self.build_logger()
|
|
self.time_window = ValueWindow(50)
|
|
|
|
self.step = 0
|
|
self.epoch = -1
|
|
self.max_epochs = self.cfg.train.epochs
|
|
self.max_steps = self.cfg.train.max_steps
|
|
|
|
# set random seed & init distributed training
|
|
set_all_random_seed(self.cfg.train.random_seed)
|
|
if cfg.train.ddp:
|
|
dist.init_process_group(backend="nccl")
|
|
|
|
if cfg.model_type not in ["AutoencoderKL", "AudioLDM"]:
|
|
self.singers = self.build_singers_lut()
|
|
|
|
# setup data_loader
|
|
self.data_loader = self.build_data_loader()
|
|
|
|
# setup model & enable distributed training
|
|
self.model = self.build_model()
|
|
print(self.model)
|
|
|
|
if isinstance(self.model, dict):
|
|
for key, value in self.model.items():
|
|
value.cuda(self.args.local_rank)
|
|
if key == "PQMF":
|
|
continue
|
|
if cfg.train.ddp:
|
|
self.model[key] = DistributedDataParallel(
|
|
value, device_ids=[self.args.local_rank]
|
|
)
|
|
else:
|
|
self.model.cuda(self.args.local_rank)
|
|
if cfg.train.ddp:
|
|
self.model = DistributedDataParallel(
|
|
self.model, device_ids=[self.args.local_rank]
|
|
)
|
|
|
|
# create criterion
|
|
self.criterion = self.build_criterion()
|
|
if isinstance(self.criterion, dict):
|
|
for key, value in self.criterion.items():
|
|
self.criterion[key].cuda(args.local_rank)
|
|
else:
|
|
self.criterion.cuda(self.args.local_rank)
|
|
|
|
# optimizer
|
|
self.optimizer = self.build_optimizer()
|
|
self.scheduler = self.build_scheduler()
|
|
|
|
# save config file
|
|
self.config_save_path = os.path.join(self.checkpoint_dir, "args.json")
|
|
|
|
def build_logger(self):
|
|
log_file = os.path.join(self.checkpoint_dir, "train.log")
|
|
logger = Logger(log_file, level=self.args.log_level).logger
|
|
|
|
return logger
|
|
|
|
def build_dataset(self):
|
|
raise NotImplementedError
|
|
|
|
def build_data_loader(self):
|
|
Dataset, Collator = self.build_dataset()
|
|
# build dataset instance for each dataset and combine them by ConcatDataset
|
|
datasets_list = []
|
|
for dataset in self.cfg.dataset:
|
|
subdataset = Dataset(self.cfg, dataset, is_valid=False)
|
|
datasets_list.append(subdataset)
|
|
train_dataset = ConcatDataset(datasets_list)
|
|
|
|
train_collate = Collator(self.cfg)
|
|
# TODO: multi-GPU training
|
|
if self.cfg.train.ddp:
|
|
raise NotImplementedError("DDP is not supported yet.")
|
|
|
|
# sampler will provide indices to batch_sampler, which will perform batching and yield batch indices
|
|
batch_sampler = BatchSampler(
|
|
cfg=self.cfg, concat_dataset=train_dataset, dataset_list=datasets_list
|
|
)
|
|
|
|
# use batch_sampler argument instead of (sampler, shuffle, drop_last, batch_size)
|
|
train_loader = DataLoader(
|
|
train_dataset,
|
|
collate_fn=train_collate,
|
|
num_workers=self.args.num_workers,
|
|
batch_sampler=batch_sampler,
|
|
pin_memory=False,
|
|
)
|
|
if not self.cfg.train.ddp or self.args.local_rank == 0:
|
|
datasets_list = []
|
|
for dataset in self.cfg.dataset:
|
|
subdataset = Dataset(self.cfg, dataset, is_valid=True)
|
|
datasets_list.append(subdataset)
|
|
valid_dataset = ConcatDataset(datasets_list)
|
|
valid_collate = Collator(self.cfg)
|
|
batch_sampler = BatchSampler(
|
|
cfg=self.cfg, concat_dataset=valid_dataset, dataset_list=datasets_list
|
|
)
|
|
valid_loader = DataLoader(
|
|
valid_dataset,
|
|
collate_fn=valid_collate,
|
|
num_workers=1,
|
|
batch_sampler=batch_sampler,
|
|
)
|
|
else:
|
|
raise NotImplementedError("DDP is not supported yet.")
|
|
# valid_loader = None
|
|
data_loader = {"train": train_loader, "valid": valid_loader}
|
|
return data_loader
|
|
|
|
def build_singers_lut(self):
|
|
# combine singers
|
|
if not os.path.exists(os.path.join(self.log_dir, self.cfg.preprocess.spk2id)):
|
|
singers = collections.OrderedDict()
|
|
else:
|
|
with open(
|
|
os.path.join(self.log_dir, self.cfg.preprocess.spk2id), "r"
|
|
) as singer_file:
|
|
singers = json.load(singer_file)
|
|
singer_count = len(singers)
|
|
for dataset in self.cfg.dataset:
|
|
singer_lut_path = os.path.join(
|
|
self.cfg.preprocess.processed_dir, dataset, self.cfg.preprocess.spk2id
|
|
)
|
|
with open(singer_lut_path, "r") as singer_lut_path:
|
|
singer_lut = json.load(singer_lut_path)
|
|
for singer in singer_lut.keys():
|
|
if singer not in singers:
|
|
singers[singer] = singer_count
|
|
singer_count += 1
|
|
with open(
|
|
os.path.join(self.log_dir, self.cfg.preprocess.spk2id), "w"
|
|
) as singer_file:
|
|
json.dump(singers, singer_file, indent=4, ensure_ascii=False)
|
|
print(
|
|
"singers have been dumped to {}".format(
|
|
os.path.join(self.log_dir, self.cfg.preprocess.spk2id)
|
|
)
|
|
)
|
|
return singers
|
|
|
|
def build_model(self):
|
|
raise NotImplementedError()
|
|
|
|
def build_optimizer(self):
|
|
raise NotImplementedError
|
|
|
|
def build_scheduler(self):
|
|
raise NotImplementedError()
|
|
|
|
def build_criterion(self):
|
|
raise NotImplementedError
|
|
|
|
def get_state_dict(self):
|
|
raise NotImplementedError
|
|
|
|
def save_config_file(self):
|
|
save_config(self.config_save_path, self.cfg)
|
|
|
|
# TODO, save without module.
|
|
def save_checkpoint(self, state_dict, saved_model_path):
|
|
torch.save(state_dict, saved_model_path)
|
|
|
|
def load_checkpoint(self):
|
|
checkpoint_path = os.path.join(self.checkpoint_dir, "checkpoint")
|
|
assert os.path.exists(checkpoint_path)
|
|
checkpoint_filename = open(checkpoint_path).readlines()[-1].strip()
|
|
model_path = os.path.join(self.checkpoint_dir, checkpoint_filename)
|
|
assert os.path.exists(model_path)
|
|
if not self.cfg.train.ddp or self.args.local_rank == 0:
|
|
self.logger.info(f"Re(store) from {model_path}")
|
|
checkpoint = torch.load(model_path, map_location="cpu")
|
|
return checkpoint
|
|
|
|
def load_model(self, checkpoint):
|
|
raise NotImplementedError
|
|
|
|
def restore(self):
|
|
checkpoint = self.load_checkpoint()
|
|
self.load_model(checkpoint)
|
|
|
|
def train_step(self, data):
|
|
raise NotImplementedError(
|
|
f"Need to implement function {sys._getframe().f_code.co_name} in "
|
|
f"your sub-class of {self.__class__.__name__}. "
|
|
)
|
|
|
|
@torch.no_grad()
|
|
def eval_step(self):
|
|
raise NotImplementedError(
|
|
f"Need to implement function {sys._getframe().f_code.co_name} in "
|
|
f"your sub-class of {self.__class__.__name__}. "
|
|
)
|
|
|
|
def write_summary(self, losses, stats):
|
|
raise NotImplementedError(
|
|
f"Need to implement function {sys._getframe().f_code.co_name} in "
|
|
f"your sub-class of {self.__class__.__name__}. "
|
|
)
|
|
|
|
def write_valid_summary(self, losses, stats):
|
|
raise NotImplementedError(
|
|
f"Need to implement function {sys._getframe().f_code.co_name} in "
|
|
f"your sub-class of {self.__class__.__name__}. "
|
|
)
|
|
|
|
def echo_log(self, losses, mode="Training"):
|
|
message = [
|
|
"{} - Epoch {} Step {}: [{:.3f} s/step]".format(
|
|
mode, self.epoch + 1, self.step, self.time_window.average
|
|
)
|
|
]
|
|
|
|
for key in sorted(losses.keys()):
|
|
if isinstance(losses[key], dict):
|
|
for k, v in losses[key].items():
|
|
message.append(
|
|
str(k).split("/")[-1] + "=" + str(round(float(v), 5))
|
|
)
|
|
else:
|
|
message.append(
|
|
str(key).split("/")[-1] + "=" + str(round(float(losses[key]), 5))
|
|
)
|
|
self.logger.info(", ".join(message))
|
|
|
|
def eval_epoch(self):
|
|
self.logger.info("Validation...")
|
|
valid_losses = {}
|
|
for i, batch_data in enumerate(self.data_loader["valid"]):
|
|
for k, v in batch_data.items():
|
|
if isinstance(v, torch.Tensor):
|
|
batch_data[k] = v.cuda()
|
|
valid_loss, valid_stats, total_valid_loss = self.eval_step(batch_data, i)
|
|
for key in valid_loss:
|
|
if key not in valid_losses:
|
|
valid_losses[key] = 0
|
|
valid_losses[key] += valid_loss[key]
|
|
|
|
# Add mel and audio to the Tensorboard
|
|
# Average loss
|
|
for key in valid_losses:
|
|
valid_losses[key] /= i + 1
|
|
self.echo_log(valid_losses, "Valid")
|
|
return valid_losses, valid_stats
|
|
|
|
def train_epoch(self):
|
|
for i, batch_data in enumerate(self.data_loader["train"]):
|
|
start_time = time.time()
|
|
# Put the data to cuda device
|
|
for k, v in batch_data.items():
|
|
if isinstance(v, torch.Tensor):
|
|
batch_data[k] = v.cuda(self.args.local_rank)
|
|
|
|
# Training step
|
|
train_losses, train_stats, total_loss = self.train_step(batch_data)
|
|
self.time_window.append(time.time() - start_time)
|
|
|
|
if self.args.local_rank == 0 or not self.cfg.train.ddp:
|
|
if self.step % self.args.stdout_interval == 0:
|
|
self.echo_log(train_losses, "Training")
|
|
|
|
if self.step % self.cfg.train.save_summary_steps == 0:
|
|
self.logger.info(f"Save summary as step {self.step}")
|
|
self.write_summary(train_losses, train_stats)
|
|
|
|
if (
|
|
self.step % self.cfg.train.save_checkpoints_steps == 0
|
|
and self.step != 0
|
|
):
|
|
saved_model_name = "step-{:07d}_loss-{:.4f}.pt".format(
|
|
self.step, total_loss
|
|
)
|
|
saved_model_path = os.path.join(
|
|
self.checkpoint_dir, saved_model_name
|
|
)
|
|
saved_state_dict = self.get_state_dict()
|
|
self.save_checkpoint(saved_state_dict, saved_model_path)
|
|
self.save_config_file()
|
|
# keep max n models
|
|
remove_older_ckpt(
|
|
saved_model_name,
|
|
self.checkpoint_dir,
|
|
max_to_keep=self.cfg.train.keep_checkpoint_max,
|
|
)
|
|
|
|
if self.step != 0 and self.step % self.cfg.train.valid_interval == 0:
|
|
if isinstance(self.model, dict):
|
|
for key in self.model.keys():
|
|
self.model[key].eval()
|
|
else:
|
|
self.model.eval()
|
|
# Evaluate one epoch and get average loss
|
|
valid_losses, valid_stats = self.eval_epoch()
|
|
if isinstance(self.model, dict):
|
|
for key in self.model.keys():
|
|
self.model[key].train()
|
|
else:
|
|
self.model.train()
|
|
# Write validation losses to summary.
|
|
self.write_valid_summary(valid_losses, valid_stats)
|
|
self.step += 1
|
|
|
|
def train(self):
|
|
for epoch in range(max(0, self.epoch), self.max_epochs):
|
|
self.train_epoch()
|
|
self.epoch += 1
|
|
if self.step > self.max_steps:
|
|
self.logger.info("Training finished!")
|
|
break
|