class_docstr: JETS module (generator + discriminator).
  
  This is a module of JETS described in `JETS: Jointly Training FastSpeech2
  and HiFi-GAN for End to End Text to Speech'_.
  
  .. _`JETS: Jointly Training FastSpeech2 and HiFi-GAN for End to End Text to Speech`
      : https://arxiv.org/abs/2203.16852
class_docstr_tok: ["JETS","module","(","generator","+","discriminator",")",".","This","is","a","module","of","JETS","described","in","`","JETS",":","Jointly","Training","FastSpeech2","and","HiFi-GAN","for","End","to","End","Text","to","Speech","'","_",".","..","_","`","JETS",":","Jointly","Training","FastSpeech2","and","HiFi-GAN","for","End","to","End","Text","to","Speech","`",":","https",":","\/\/arxiv.org\/abs\/2203.16852"]
docstr: Perform generator forward.
  
  Args:
      text (Tensor): Text index tensor (B, T_text).
      text_lengths (Tensor): Text length tensor (B,).
      feats (Tensor): Feature tensor (B, T_feats, aux_channels).
      feats_lengths (Tensor): Feature length tensor (B,).
      speech (Tensor): Speech waveform tensor (B, T_wav).
      speech_lengths (Tensor): Speech length tensor (B,).
      sids (Optional[Tensor]): Speaker index tensor (B,) or (B, 1).
      spembs (Optional[Tensor]): Speaker embedding tensor (B, spk_embed_dim).
      lids (Optional[Tensor]): Language index tensor (B,) or (B, 1).
  
  Returns:
      Dict[str, Any]:
          * loss (Tensor): Loss scalar tensor.
          * stats (Dict[str, float]): Statistics to be monitored.
          * weight (Tensor): Weight tensor to summarize losses.
          * optim_idx (int): Optimizer index (0 for G and 1 for D).
docstr_tok: ["Perform","generator","forward",".","Args",":","text","(","Tensor",")",":","Text","index","tensor","(","B",",","T_text",")",".","text_lengths","(","Tensor",")",":","Text","length","tensor","(","B",",",")",".","feats","(","Tensor",")",":","Feature","tensor","(","B",",","T_feats",",","aux_channels",")",".","feats_lengths","(","Tensor",")",":","Feature","length","tensor","(","B",",",")",".","speech","(","Tensor",")",":","Speech","waveform","tensor","(","B",",","T_wav",")",".","speech_lengths","(","Tensor",")",":","Speech","length","tensor","(","B",",",")",".","sids","(","Optional","[","Tensor","]",")",":","Speaker","index","tensor","(","B",",",")","or","(","B",",","1",")",".","spembs","(","Optional","[","Tensor","]",")",":","Speaker","embedding","tensor","(","B",",","spk_embed_dim",")",".","lids","(","Optional","[","Tensor","]",")",":","Language","index","tensor","(","B",",",")","or","(","B",",","1",")",".","Returns",":","Dict","[","str",",","Any","]",":","*","loss","(","Tensor",")",":","Loss","scalar","tensor",".","*","stats","(","Dict","[","str",",","float","]",")",":","Statistics","to","be","monitored",".","*","weight","(","Tensor",")",":","Weight","tensor","to","summarize","losses",".","*","optim_idx","(","int",")",":","Optimizer","index","(","0","for","G","and","1","for","D",")","."]
code: def _forward_generator(
      self,
      text: torch.Tensor,
      text_lengths: torch.Tensor,
      feats: torch.Tensor,
      feats_lengths: torch.Tensor,
      speech: torch.Tensor,
      speech_lengths: torch.Tensor,
      sids: Optional[torch.Tensor] = None,
      spembs: Optional[torch.Tensor] = None,
      lids: Optional[torch.Tensor] = None,
      **kwargs,
  ) -> Dict[str, Any]:
      """Perform generator forward.
  
      Args:
          text (Tensor): Text index tensor (B, T_text).
          text_lengths (Tensor): Text length tensor (B,).
          feats (Tensor): Feature tensor (B, T_feats, aux_channels).
          feats_lengths (Tensor): Feature length tensor (B,).
          speech (Tensor): Speech waveform tensor (B, T_wav).
          speech_lengths (Tensor): Speech length tensor (B,).
          sids (Optional[Tensor]): Speaker index tensor (B,) or (B, 1).
          spembs (Optional[Tensor]): Speaker embedding tensor (B, spk_embed_dim).
          lids (Optional[Tensor]): Language index tensor (B,) or (B, 1).
  
      Returns:
          Dict[str, Any]:
              * loss (Tensor): Loss scalar tensor.
              * stats (Dict[str, float]): Statistics to be monitored.
              * weight (Tensor): Weight tensor to summarize losses.
              * optim_idx (int): Optimizer index (0 for G and 1 for D).
  
      """
      # setup
      batch_size = text.size(0)
      speech = speech.unsqueeze(1)
  
      # calculate generator outputs
      reuse_cache = True
      if not self.cache_generator_outputs or self._cache is None:
          reuse_cache = False
          outs = self.generator(
              text=text,
              text_lengths=text_lengths,
              feats=feats,
              feats_lengths=feats_lengths,
              sids=sids,
              spembs=spembs,
              lids=lids,
              **kwargs,
          )
      else:
          outs = self._cache
  
      # store cache
      if (
          self.training
          and self.cache_generator_outputs
          and not reuse_cache
      ):
          self._cache = outs
  
      # parse outputs
      (
          speech_hat_,
          bin_loss,
          log_p_attn,
          start_idxs,
          d_outs,
          ds,
          p_outs,
          ps,
          e_outs,
          es,
      ) = outs
      speech_ = get_segments(
          x=speech,
          start_idxs=start_idxs * self.generator.upsample_factor,
          segment_size=self.generator.segment_size
          * self.generator.upsample_factor,
      )
  
      # calculate discriminator outputs
      p_hat = self.discriminator(speech_hat_)
      with torch.no_grad():
          # do not store discriminator gradient in generator turn
          p = self.discriminator(speech_)
  
      # calculate losses
      mel_loss = self.mel_loss(speech_hat_, speech_)
      adv_loss = self.generator_adv_loss(p_hat)
      feat_match_loss = self.feat_match_loss(p_hat, p)
      dur_loss, pitch_loss, energy_loss = self.var_loss(
          d_outs, ds, p_outs, ps, e_outs, es, text_lengths
      )
      forwardsum_loss = self.forwardsum_loss(
          log_p_attn, text_lengths, feats_lengths
      )
  
      mel_loss = mel_loss * self.lambda_mel
      adv_loss = adv_loss * self.lambda_adv
      feat_match_loss = feat_match_loss * self.lambda_feat_match
      g_loss = mel_loss + adv_loss + feat_match_loss
      var_loss = (dur_loss + pitch_loss + energy_loss) * self.lambda_var
      align_loss = (forwardsum_loss + bin_loss) * self.lambda_align
  
      loss = g_loss + var_loss + align_loss
  
      stats = dict(
          generator_loss=loss.item(),
          generator_g_loss=g_loss.item(),
          generator_var_loss=var_loss.item(),
          generator_align_loss=align_loss.item(),
          generator_g_mel_loss=mel_loss.item(),
          generator_g_adv_loss=adv_loss.item(),
          generator_g_feat_match_loss=feat_match_loss.item(),
          generator_var_dur_loss=dur_loss.item(),
          generator_var_pitch_loss=pitch_loss.item(),
          generator_var_energy_loss=energy_loss.item(),
          generator_align_forwardsum_loss=forwardsum_loss.item(),
          generator_align_bin_loss=bin_loss.item(),
      )
  
      if self.plot_pred_mos:
          # Caltulate predicted MOS from generated speech waveform.
          with torch.no_grad():
              # speech_hat_: (B, 1, T)
              pmos = self.predictor(
                  speech_hat_.squeeze(1), self.fs
              ).mean()
          stats["generator_predicted_mos"] = pmos.item()
  
      loss, stats, weight = force_gatherable(
          (loss, stats, batch_size), loss.device
      )
  
      # reset cache
      if reuse_cache or not self.training:
          self._cache = None
  
      return {
          "loss": loss,
          "stats": stats,
          "weight": weight,
          "optim_idx": 0,  # needed for trainer
      }
  
file_path: espnet/espnet2/gan_tts/jets/jets.py
imports: from typing import Any, Dict, Optional
  import torch
  from typeguard import check_argument_types
  from espnet2.gan_tts.abs_gan_tts import AbsGANTTS
  from espnet2.gan_tts.hifigan import HiFiGANMultiPeriodDiscriminator, HiFiGANMultiScaleDiscriminator, HiFiGANMultiScaleMultiPeriodDiscriminator, HiFiGANPeriodDiscriminator, HiFiGANScaleDiscriminator
  from espnet2.gan_tts.hifigan.loss import DiscriminatorAdversarialLoss, FeatureMatchLoss, GeneratorAdversarialLoss, MelSpectrogramLoss
  from espnet2.gan_tts.jets.generator import JETSGenerator
  from espnet2.gan_tts.jets.loss import ForwardSumLoss, VarianceLoss
  from espnet2.gan_tts.utils import get_segments
  from espnet2.torch_utils.device_funcs import force_gatherable
label_desc: Use image node_id 5 for calling the JETS obj's underlying member method code with example usage: obj._forward_generator(text, text_lengths, feats, feats_lengths, speech, speech_lengths, sids, spembs, lids) and returns: dict