PaddleSpeech/paddlespeech/t2s/models/vits/vits_updater.py

# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
from typing import Dict

import paddle
from paddle import distributed as dist
from paddle.io import DataLoader
from paddle.nn import Layer
from paddle.optimizer import Optimizer
from paddle.optimizer.lr import LRScheduler

from paddlespeech.t2s.modules.nets_utils import get_segments
from paddlespeech.t2s.training.extensions.evaluator import StandardEvaluator
from paddlespeech.t2s.training.reporter import report
from paddlespeech.t2s.training.updaters.standard_updater import StandardUpdater
from paddlespeech.t2s.training.updaters.standard_updater import UpdaterState

logging.basicConfig(
    format='%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s',
    datefmt='[%Y-%m-%d %H:%M:%S]')
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


class VITSUpdater(StandardUpdater):
    def __init__(self,
                 model: Layer,
                 optimizers: Dict[str, Optimizer],
                 criterions: Dict[str, Layer],
                 schedulers: Dict[str, LRScheduler],
                 dataloader: DataLoader,
                 generator_train_start_steps: int=0,
                 discriminator_train_start_steps: int=100000,
                 lambda_adv: float=1.0,
                 lambda_mel: float=45.0,
                 lambda_feat_match: float=2.0,
                 lambda_dur: float=1.0,
                 lambda_kl: float=1.0,
                 generator_first: bool=False,
                 output_dir=None):
        # it is designed to hold multiple models
        # 因为输入的是单模型，但是没有用到父类的 init(), 所以需要重新写这部分
        models = {"main": model}
        self.models: Dict[str, Layer] = models
        # self.model = model

        self.model = model._layers if isinstance(model,
                                                 paddle.DataParallel) else model

        self.optimizers = optimizers
        self.optimizer_g: Optimizer = optimizers['generator']
        self.optimizer_d: Optimizer = optimizers['discriminator']

        self.criterions = criterions
        self.criterion_mel = criterions['mel']
        self.criterion_feat_match = criterions['feat_match']
        self.criterion_gen_adv = criterions["gen_adv"]
        self.criterion_dis_adv = criterions["dis_adv"]
        self.criterion_kl = criterions["kl"]

        self.schedulers = schedulers
        self.scheduler_g = schedulers['generator']
        self.scheduler_d = schedulers['discriminator']

        self.dataloader = dataloader

        self.generator_train_start_steps = generator_train_start_steps
        self.discriminator_train_start_steps = discriminator_train_start_steps

        self.lambda_adv = lambda_adv
        self.lambda_mel = lambda_mel
        self.lambda_feat_match = lambda_feat_match
        self.lambda_dur = lambda_dur
        self.lambda_kl = lambda_kl

        if generator_first:
            self.turns = ["generator", "discriminator"]
        else:
            self.turns = ["discriminator", "generator"]

        self.state = UpdaterState(iteration=0, epoch=0)
        self.train_iterator = iter(self.dataloader)

        log_file = output_dir / 'worker_{}.log'.format(dist.get_rank())
        self.filehandler = logging.FileHandler(str(log_file))
        logger.addHandler(self.filehandler)
        self.logger = logger
        self.msg = ""

    def update_core(self, batch):
        self.msg = "Rank: {}, ".format(dist.get_rank())
        losses_dict = {}

        for turn in self.turns:
            speech = batch["speech"]
            speech = speech.unsqueeze(1)
            outs = self.model(
                text=batch["text"],
                text_lengths=batch["text_lengths"],
                feats=batch["feats"],
                feats_lengths=batch["feats_lengths"],
                sids=batch.get("spk_id", None),
                spembs=batch.get("spk_emb", None),
                forward_generator=turn == "generator")
            # Generator
            if turn == "generator":
                # parse outputs
                speech_hat_, dur_nll, _, start_idxs, _, z_mask, outs_ = outs
                _, z_p, m_p, logs_p, _, logs_q = outs_
                speech_ = get_segments(
                    x=speech,
                    start_idxs=start_idxs *
                    self.model.generator.upsample_factor,
                    segment_size=self.model.generator.segment_size *
                    self.model.generator.upsample_factor, )

                # calculate discriminator outputs
                p_hat = self.model.discriminator(speech_hat_)
                with paddle.no_grad():
                    # do not store discriminator gradient in generator turn
                    p = self.model.discriminator(speech_)

                # calculate losses
                mel_loss = self.criterion_mel(speech_hat_, speech_)
                kl_loss = self.criterion_kl(z_p, logs_q, m_p, logs_p, z_mask)
                dur_loss = paddle.sum(dur_nll)
                adv_loss = self.criterion_gen_adv(p_hat)
                feat_match_loss = self.criterion_feat_match(p_hat, p)

                mel_loss = mel_loss * self.lambda_mel
                kl_loss = kl_loss * self.lambda_kl
                dur_loss = dur_loss * self.lambda_dur
                adv_loss = adv_loss * self.lambda_adv
                feat_match_loss = feat_match_loss * self.lambda_feat_match
                gen_loss = mel_loss + kl_loss + dur_loss + adv_loss + feat_match_loss

                report("train/generator_loss", float(gen_loss))
                report("train/generator_mel_loss", float(mel_loss))
                report("train/generator_kl_loss", float(kl_loss))
                report("train/generator_dur_loss", float(dur_loss))
                report("train/generator_adv_loss", float(adv_loss))
                report("train/generator_feat_match_loss",
                       float(feat_match_loss))

                losses_dict["generator_loss"] = float(gen_loss)
                losses_dict["generator_mel_loss"] = float(mel_loss)
                losses_dict["generator_kl_loss"] = float(kl_loss)
                losses_dict["generator_dur_loss"] = float(dur_loss)
                losses_dict["generator_adv_loss"] = float(adv_loss)
                losses_dict["generator_feat_match_loss"] = float(
                    feat_match_loss)

                self.optimizer_g.clear_grad()
                gen_loss.backward()

                self.optimizer_g.step()
                self.scheduler_g.step()

                # reset cache
                if self.model.reuse_cache_gen or not self.model.training:
                    self.model._cache = None

            # Disctiminator
            elif turn == "discriminator":
                # parse outputs
                speech_hat_, _, _, start_idxs, *_ = outs
                speech_ = get_segments(
                    x=speech,
                    start_idxs=start_idxs *
                    self.model.generator.upsample_factor,
                    segment_size=self.model.generator.segment_size *
                    self.model.generator.upsample_factor, )

                # calculate discriminator outputs
                p_hat = self.model.discriminator(speech_hat_.detach())
                p = self.model.discriminator(speech_)

                # calculate losses
                real_loss, fake_loss = self.criterion_dis_adv(p_hat, p)
                dis_loss = real_loss + fake_loss

                report("train/real_loss", float(real_loss))
                report("train/fake_loss", float(fake_loss))
                report("train/discriminator_loss", float(dis_loss))
                losses_dict["real_loss"] = float(real_loss)
                losses_dict["fake_loss"] = float(fake_loss)
                losses_dict["discriminator_loss"] = float(dis_loss)

                self.optimizer_d.clear_grad()
                dis_loss.backward()

                self.optimizer_d.step()
                self.scheduler_d.step()

                # reset cache
                if self.model.reuse_cache_dis or not self.model.training:
                    self.model._cache = None

        self.msg += ', '.join('{}: {:>.6f}'.format(k, v)
                              for k, v in losses_dict.items())


class VITSEvaluator(StandardEvaluator):
    def __init__(self,
                 model,
                 criterions: Dict[str, Layer],
                 dataloader: DataLoader,
                 lambda_adv: float=1.0,
                 lambda_mel: float=45.0,
                 lambda_feat_match: float=2.0,
                 lambda_dur: float=1.0,
                 lambda_kl: float=1.0,
                 generator_first: bool=False,
                 output_dir=None):
        # 因为输入的是单模型，但是没有用到父类的 init(), 所以需要重新写这部分
        models = {"main": model}
        self.models: Dict[str, Layer] = models
        # self.model = model
        self.model = model._layers if isinstance(model,
                                                 paddle.DataParallel) else model

        self.criterions = criterions
        self.criterion_mel = criterions['mel']
        self.criterion_feat_match = criterions['feat_match']
        self.criterion_gen_adv = criterions["gen_adv"]
        self.criterion_dis_adv = criterions["dis_adv"]
        self.criterion_kl = criterions["kl"]

        self.dataloader = dataloader

        self.lambda_adv = lambda_adv
        self.lambda_mel = lambda_mel
        self.lambda_feat_match = lambda_feat_match
        self.lambda_dur = lambda_dur
        self.lambda_kl = lambda_kl

        if generator_first:
            self.turns = ["generator", "discriminator"]
        else:
            self.turns = ["discriminator", "generator"]

        log_file = output_dir / 'worker_{}.log'.format(dist.get_rank())
        self.filehandler = logging.FileHandler(str(log_file))
        logger.addHandler(self.filehandler)
        self.logger = logger
        self.msg = ""

    def evaluate_core(self, batch):
        # logging.debug("Evaluate: ")
        self.msg = "Evaluate: "
        losses_dict = {}

        for turn in self.turns:
            speech = batch["speech"]
            speech = speech.unsqueeze(1)
            outs = self.model(
                text=batch["text"],
                text_lengths=batch["text_lengths"],
                feats=batch["feats"],
                feats_lengths=batch["feats_lengths"],
                sids=batch.get("spk_id", None),
                spembs=batch.get("spk_emb", None),
                forward_generator=turn == "generator")
            # Generator
            if turn == "generator":
                # parse outputs
                speech_hat_, dur_nll, _, start_idxs, _, z_mask, outs_ = outs
                _, z_p, m_p, logs_p, _, logs_q = outs_
                speech_ = get_segments(
                    x=speech,
                    start_idxs=start_idxs *
                    self.model.generator.upsample_factor,
                    segment_size=self.model.generator.segment_size *
                    self.model.generator.upsample_factor, )

                # calculate discriminator outputs
                p_hat = self.model.discriminator(speech_hat_)
                with paddle.no_grad():
                    # do not store discriminator gradient in generator turn
                    p = self.model.discriminator(speech_)

                # calculate losses
                mel_loss = self.criterion_mel(speech_hat_, speech_)
                kl_loss = self.criterion_kl(z_p, logs_q, m_p, logs_p, z_mask)
                dur_loss = paddle.sum(dur_nll)
                adv_loss = self.criterion_gen_adv(p_hat)
                feat_match_loss = self.criterion_feat_match(p_hat, p)

                mel_loss = mel_loss * self.lambda_mel
                kl_loss = kl_loss * self.lambda_kl
                dur_loss = dur_loss * self.lambda_dur
                adv_loss = adv_loss * self.lambda_adv
                feat_match_loss = feat_match_loss * self.lambda_feat_match
                gen_loss = mel_loss + kl_loss + dur_loss + adv_loss + feat_match_loss

                report("eval/generator_loss", float(gen_loss))
                report("eval/generator_mel_loss", float(mel_loss))
                report("eval/generator_kl_loss", float(kl_loss))
                report("eval/generator_dur_loss", float(dur_loss))
                report("eval/generator_adv_loss", float(adv_loss))
                report("eval/generator_feat_match_loss", float(feat_match_loss))

                losses_dict["generator_loss"] = float(gen_loss)
                losses_dict["generator_mel_loss"] = float(mel_loss)
                losses_dict["generator_kl_loss"] = float(kl_loss)
                losses_dict["generator_dur_loss"] = float(dur_loss)
                losses_dict["generator_adv_loss"] = float(adv_loss)
                losses_dict["generator_feat_match_loss"] = float(
                    feat_match_loss)

                # reset cache
                if self.model.reuse_cache_gen or not self.model.training:
                    self.model._cache = None

            # Disctiminator
            elif turn == "discriminator":
                # parse outputs
                speech_hat_, _, _, start_idxs, *_ = outs
                speech_ = get_segments(
                    x=speech,
                    start_idxs=start_idxs *
                    self.model.generator.upsample_factor,
                    segment_size=self.model.generator.segment_size *
                    self.model.generator.upsample_factor, )

                # calculate discriminator outputs
                p_hat = self.model.discriminator(speech_hat_.detach())
                p = self.model.discriminator(speech_)

                # calculate losses
                real_loss, fake_loss = self.criterion_dis_adv(p_hat, p)
                dis_loss = real_loss + fake_loss

                report("eval/real_loss", float(real_loss))
                report("eval/fake_loss", float(fake_loss))
                report("eval/discriminator_loss", float(dis_loss))
                losses_dict["real_loss"] = float(real_loss)
                losses_dict["fake_loss"] = float(fake_loss)
                losses_dict["discriminator_loss"] = float(dis_loss)

                # reset cache
                if self.model.reuse_cache_dis or not self.model.training:
                    self.model._cache = None

        self.msg += ', '.join('{}: {:>.6f}'.format(k, v)
                              for k, v in losses_dict.items())
        self.logger.info(self.msg)