class_docstr: Speech enhancement or separation Frontend model
docstr: Main entry of speech enhancement/separation model training.
  
  Args:
      encoder: waveform encoder that converts waveforms to feature representations
      separator: separator that enhance or separate the feature representations
      decoder: waveform decoder that converts the feature back to waveforms
      mask_module: mask module that converts the feature to masks
          NOTE: Only used for compatibility with joint speaker diarization.
          See test/espnet2/enh/test_espnet_enh_s2t_model.py for details.
      loss_wrappers: list of loss wrappers
          Each loss wrapper contains a criterion for loss calculation and
          the corresonding loss weight.
          The losses will be calculated in the order of the list and summed up.
      ------------------------------------------------------------------
      stft_consistency: (deprecated, kept for compatibility) whether to compute
          the TF-domain loss while enforcing STFT consistency
          NOTE: STFT consistency is now always used for frequency-domain spectrum
          losses.
      loss_type: (deprecated, kept for compatibility) loss type
      mask_type: (deprecated, kept for compatibility) mask type in TF-domain model
      ------------------------------------------------------------------
      flexible_numspk: whether to allow the model to predict a variable number of
          speakers in its output.
          NOTE: This should be used when training a speech separation model for
          unknown number of speakers.
      ------------------------------------------------------------------
      extract_feats_in_collect_stats: used in espnet2/tasks/abs_task.py for
          determining whether or not to skip model building in collect_stats stage
          (stage 5 in egs2/*/enh1/enh.sh).
      normalize_variance: whether to normalize the signal variance before model
          forward, and revert it back after.
      normalize_variance_per_ch: whether to normalize the signal variance for each
          channel instead of the whole signal.
          NOTE: normalize_variance and normalize_variance_per_ch cannot be True
          at the same time.
      ------------------------------------------------------------------
      categories: list of all possible categories of minibatches (order matters!)
          (e.g. ["1ch_8k_reverb", "1ch_8k_both"] for multi-condition training)
          NOTE: this will be used to convert category index to the corresponding
          name for logging in forward_loss.
          Different categories will have different loss name suffixes.
      category_weights: list of weights for each category.
          Used to set loss weights for batches of different categories.
docstr_tok: ["Main","entry","of","speech","enhancement\/separation","model","training",".","Args",":","encoder",":","waveform","encoder","that","converts","waveforms","to","feature","representations","separator",":","separator","that","enhance","or","separate","the","feature","representations","decoder",":","waveform","decoder","that","converts","the","feature","back","to","waveforms","mask_module",":","mask","module","that","converts","the","feature","to","masks","NOTE",":","Only","used","for","compatibility","with","joint","speaker","diarization",".","See","test\/espnet2\/enh\/test_espnet_enh_s2t_model.py","for","details",".","loss_wrappers",":","list","of","loss","wrappers","Each","loss","wrapper","contains","a","criterion","for","loss","calculation","and","the","corresonding","loss","weight",".","The","losses","will","be","calculated","in","the","order","of","the","list","and","summed","up",".","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","stft_consistency",":","(","deprecated",",","kept","for","compatibility",")","whether","to","compute","the","TF-domain","loss","while","enforcing","STFT","consistency","NOTE",":","STFT","consistency","is","now","always","used","for","frequency-domain","spectrum","losses",".","loss_type",":","(","deprecated",",","kept","for","compatibility",")","loss","type","mask_type",":","(","deprecated",",","kept","for","compatibility",")","mask","type","in","TF-domain","model","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","flexible_numspk",":","whether","to","allow","the","model","to","predict","a","variable","number","of","speakers","in","its","output",".","NOTE",":","This","should","be","used","when","training","a","speech","separation","model","for","unknown","number","of","speakers",".","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","extract_feats_in_collect_stats",":","used","in","espnet2\/tasks\/abs_task.py","for","determining","whether","or","not","to","skip","model","building","in","collect_stats","stage","(","stage","5","in","egs2\/","*","\/enh1\/enh.sh",")",".","normalize_variance",":","whether","to","normalize","the","signal","variance","before","model","forward",",","and","revert","it","back","after",".","normalize_variance_per_ch",":","whether","to","normalize","the","signal","variance","for","each","channel","instead","of","the","whole","signal",".","NOTE",":","normalize_variance","and","normalize_variance_per_ch","can","not","be","True","at","the","same","time",".","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","categories",":","list","of","all","possible","categories","of","minibatches","(","order","matters","!",")","(","e.g",".","[","``","1ch_8k_reverb","''",",","``","1ch_8k_both","''","]","for","multi-condition","training",")","NOTE",":","this","will","be","used","to","convert","category","index","to","the","corresponding","name","for","logging","in","forward_loss",".","Different","categories","will","have","different","loss","name","suffixes",".","category_weights",":","list","of","weights","for","each","category",".","Used","to","set","loss","weights","for","batches","of","different","categories","."]
code: def __init__(
      self,
      encoder: AbsEncoder,
      separator: Optional[AbsSeparator],
      decoder: AbsDecoder,
      mask_module: Optional[AbsMask],
      loss_wrappers: Optional[List[AbsLossWrapper]],
      stft_consistency: bool = False,
      loss_type: str = "mask_mse",
      mask_type: Optional[str] = None,
      flexible_numspk: bool = False,
      extract_feats_in_collect_stats: bool = False,
      normalize_variance: bool = False,
      normalize_variance_per_ch: bool = False,
      categories: list = [],
      category_weights: list = [],
  ):
      """Main entry of speech enhancement/separation model training.
  
      Args:
          encoder: waveform encoder that converts waveforms to feature representations
          separator: separator that enhance or separate the feature representations
          decoder: waveform decoder that converts the feature back to waveforms
          mask_module: mask module that converts the feature to masks
              NOTE: Only used for compatibility with joint speaker diarization.
              See test/espnet2/enh/test_espnet_enh_s2t_model.py for details.
          loss_wrappers: list of loss wrappers
              Each loss wrapper contains a criterion for loss calculation and
              the corresonding loss weight.
              The losses will be calculated in the order of the list and summed up.
          ------------------------------------------------------------------
          stft_consistency: (deprecated, kept for compatibility) whether to compute
              the TF-domain loss while enforcing STFT consistency
              NOTE: STFT consistency is now always used for frequency-domain spectrum
              losses.
          loss_type: (deprecated, kept for compatibility) loss type
          mask_type: (deprecated, kept for compatibility) mask type in TF-domain model
          ------------------------------------------------------------------
          flexible_numspk: whether to allow the model to predict a variable number of
              speakers in its output.
              NOTE: This should be used when training a speech separation model for
              unknown number of speakers.
          ------------------------------------------------------------------
          extract_feats_in_collect_stats: used in espnet2/tasks/abs_task.py for
              determining whether or not to skip model building in collect_stats stage
              (stage 5 in egs2/*/enh1/enh.sh).
          normalize_variance: whether to normalize the signal variance before model
              forward, and revert it back after.
          normalize_variance_per_ch: whether to normalize the signal variance for each
              channel instead of the whole signal.
              NOTE: normalize_variance and normalize_variance_per_ch cannot be True
              at the same time.
          ------------------------------------------------------------------
          categories: list of all possible categories of minibatches (order matters!)
              (e.g. ["1ch_8k_reverb", "1ch_8k_both"] for multi-condition training)
              NOTE: this will be used to convert category index to the corresponding
              name for logging in forward_loss.
              Different categories will have different loss name suffixes.
          category_weights: list of weights for each category.
              Used to set loss weights for batches of different categories.
      """
      assert check_argument_types()
  
      super().__init__()
  
      self.encoder = encoder
      self.separator = separator
      self.decoder = decoder
      self.mask_module = mask_module
      # set num_spk to -1 if None for compatibility with `espnet2.enh.diffusion_enh`
      self.num_spk = separator.num_spk if separator is not None else -1
      # If True, self.num_spk is regarded as the MAXIMUM possible number of speakers
      self.flexible_numspk = flexible_numspk
      self.num_noise_type = getattr(self.separator, "num_noise_type", 1)
  
      self.loss_wrappers = loss_wrappers
      if self.loss_wrappers is not None:
          names = [w.criterion.name for w in self.loss_wrappers]
          if len(set(names)) != len(names):
              raise ValueError(
                  "Duplicated loss names are not allowed: {}".format(
                      names
                  )
              )
  
      # kept for compatibility
      self.mask_type = mask_type.upper() if mask_type else None
      self.loss_type = loss_type
      self.stft_consistency = stft_consistency
  
      # for multi-channel signal
      self.ref_channel = getattr(self.separator, "ref_channel", None)
      if self.ref_channel is None:
          self.ref_channel = 0
  
      self.extract_feats_in_collect_stats = (
          extract_feats_in_collect_stats
      )
  
      self.normalize_variance = normalize_variance
      self.normalize_variance_per_ch = normalize_variance_per_ch
      if normalize_variance and normalize_variance_per_ch:
          raise ValueError(
              "normalize_variance and normalize_variance_per_ch cannot be True "
              "at the same time."
          )
  
      # list all possible categories of the batch (order matters!)
      # (used to convert category index to the corresponding name for logging)
      self.categories = {}
      if categories:
          count = 0
          for c in categories:
              if c not in self.categories:
                  self.categories[count] = c
                  count += 1
      # used to set loss weights for batches of different categories
      if category_weights:
          assert len(category_weights) == len(self.categories)
          self.category_weights = tuple(category_weights)
      else:
          self.category_weights = tuple(1.0 for _ in self.categories)
  
imports: import contextlib
  from typing import Dict, List, Optional, OrderedDict, Tuple
  import numpy
  import torch
  from packaging.version import parse
  from typeguard import check_argument_types
  from espnet2.diar.layers.abs_mask import AbsMask
  from espnet2.enh.decoder.abs_decoder import AbsDecoder
  from espnet2.enh.encoder.abs_encoder import AbsEncoder
  from espnet2.enh.loss.criterions.tf_domain import FrequencyDomainLoss
  from espnet2.enh.loss.criterions.time_domain import TimeDomainLoss
  from espnet2.enh.loss.wrappers.abs_wrapper import AbsLossWrapper
  from espnet2.enh.separator.abs_separator import AbsSeparator
  from espnet2.enh.separator.dan_separator import DANSeparator
  from espnet2.enh.separator.uses_separator import USESSeparator
  from espnet2.torch_utils.device_funcs import force_gatherable
  from espnet2.train.abs_espnet_model import AbsESPnetModel
label_desc: Use image node_id 1 to create a new ESPnetEnhancementModel object from inherited base classes: AbsESPnetModel with example: obj = ESPnetEnhancementModel(encoder, separator, decoder, mask_module, loss_wrappers, stft_consistency, loss_type, mask_type, flexible_numspk, extract_feats_in_collect_stats, normalize_variance, normalize_variance_per_ch, categories, category_weights)