Silentlin/DiffSinger

1from multiprocessing.pool import Pool2 3import matplotlib4 5from utils.pl_utils import data_loader6from utils.training_utils import RSQRTSchedule7from vocoders.base_vocoder import get_vocoder_cls, BaseVocoder8from modules.fastspeech.pe import PitchExtractor9 10matplotlib.use('Agg')11import os12import numpy as np13from tqdm import tqdm14import torch.distributed as dist15 16from tasks.base_task import BaseTask17from utils.hparams import hparams18from utils.text_encoder import TokenTextEncoder19import json20 21import torch22import torch.optim23import torch.utils.data24import utils25 26 27 28class TtsTask(BaseTask):29    def __init__(self, *args, **kwargs):30        self.vocoder = None31        self.phone_encoder = self.build_phone_encoder(hparams['binary_data_dir'])32        self.padding_idx = self.phone_encoder.pad()33        self.eos_idx = self.phone_encoder.eos()34        self.seg_idx = self.phone_encoder.seg()35        self.saving_result_pool = None36        self.saving_results_futures = None37        self.stats = {}38        super().__init__(*args, **kwargs)39 40    def build_scheduler(self, optimizer):41        return RSQRTSchedule(optimizer)42 43    def build_optimizer(self, model):44        self.optimizer = optimizer = torch.optim.AdamW(45            model.parameters(),46            lr=hparams['lr'])47        return optimizer48 49    def build_dataloader(self, dataset, shuffle, max_tokens=None, max_sentences=None,50                         required_batch_size_multiple=-1, endless=False, batch_by_size=True):51        devices_cnt = torch.cuda.device_count()52        if devices_cnt == 0:53            devices_cnt = 154        if required_batch_size_multiple == -1:55            required_batch_size_multiple = devices_cnt56 57        def shuffle_batches(batches):58            np.random.shuffle(batches)59            return batches60 61        if max_tokens is not None:62            max_tokens *= devices_cnt63        if max_sentences is not None:64            max_sentences *= devices_cnt65        indices = dataset.ordered_indices()66        if batch_by_size:67            batch_sampler = utils.batch_by_size(68                indices, dataset.num_tokens, max_tokens=max_tokens, max_sentences=max_sentences,69                required_batch_size_multiple=required_batch_size_multiple,70            )71        else:72            batch_sampler = []73            for i in range(0, len(indices), max_sentences):74                batch_sampler.append(indices[i:i + max_sentences])75 76        if shuffle:77            batches = shuffle_batches(list(batch_sampler))78            if endless:79                batches = [b for _ in range(1000) for b in shuffle_batches(list(batch_sampler))]80        else:81            batches = batch_sampler82            if endless:83                batches = [b for _ in range(1000) for b in batches]84        num_workers = dataset.num_workers85        if self.trainer.use_ddp:86            num_replicas = dist.get_world_size()87            rank = dist.get_rank()88            batches = [x[rank::num_replicas] for x in batches if len(x) % num_replicas == 0]89        return torch.utils.data.DataLoader(dataset,90                                           collate_fn=dataset.collater,91                                           batch_sampler=batches,92                                           num_workers=num_workers,93                                           pin_memory=False)94 95    def build_phone_encoder(self, data_dir):96        phone_list_file = os.path.join(data_dir, 'phone_set.json')97 98        phone_list = json.load(open(phone_list_file))99        return TokenTextEncoder(None, vocab_list=phone_list, replace_oov=',')100 101    def build_optimizer(self, model):102        self.optimizer = optimizer = torch.optim.AdamW(103            model.parameters(),104            lr=hparams['lr'])105        return optimizer106 107    def test_start(self):108        self.saving_result_pool = Pool(8)109        self.saving_results_futures = []110        self.vocoder: BaseVocoder = get_vocoder_cls(hparams)()111        if hparams.get('pe_enable') is not None and hparams['pe_enable']:112            self.pe = PitchExtractor().cuda()113            utils.load_ckpt(self.pe, hparams['pe_ckpt'], 'model', strict=True)114            self.pe.eval()115    def test_end(self, outputs):116        self.saving_result_pool.close()117        [f.get() for f in tqdm(self.saving_results_futures)]118        self.saving_result_pool.join()119        return {}120 121    ##########122    # utils123    ##########124    def weights_nonzero_speech(self, target):125        # target : B x T x mel126        # Assign weight 1.0 to all labels except for padding (id=0).127        dim = target.size(-1)128        return target.abs().sum(-1, keepdim=True).ne(0).float().repeat(1, 1, dim)129 130if __name__ == '__main__':131    TtsTask.start()132