feat(mixln): add config switch to disable speaker-shuffle during training

configs/acoustic.yaml

@@ -56,6 +56,7 @@ use_spk_id: false
 num_spk: 1
 use_mix_ln: false
 mix_ln_layer: [0, 2]
+mixln_shuffle_speakers: false
 use_energy_embed: false
 use_breathiness_embed: false
 use_voicing_embed: false

configs/templates/config_acoustic.yaml

@@ -45,6 +45,7 @@ num_spk: 1
 
 use_mix_ln: false
 mix_ln_layer: [0, 2]
+mixln_shuffle_speakers: false
 
 # NOTICE: before enabling variance embeddings, please read the docs at
 # https://github.com/openvpi/DiffSinger/tree/main/docs/BestPractices.md#choosing-variance-parameters

modules/commons/common_layers.py

@@ -245,11 +245,14 @@ def __init__(
  condition_channels: int,
  beta_distribution_concentration: float = 0.2,
  eps: float = 1e-5,
- bias: bool = True
+ bias: bool = True,
+ *,
+ shuffle_speakers: bool = False
  ):
  super().__init__()
  self.channels = channels
  self.eps = eps
+ self.shuffle_speakers = shuffle_speakers
 
  self.beta_distribution = torch.distributions.Beta(
  beta_distribution_concentration,
@@ -275,6 +278,8 @@ def forward(
 
  if not self.training or x.size(0) == 1:
  return gammas * x + betas
+ if not self.shuffle_speakers:
+ return gammas * x + betas
 
  shuffle_indices = torch.randperm(x.size(0), device=x.device)
  shuffled_betas = betas[shuffle_indices]
@@ -396,7 +401,7 @@ def forward(self, x, key_padding_mask=None):
 class EncSALayer(nn.Module):
  def __init__(self, c, num_heads, dropout, attention_dropout=0.1,
  relu_dropout=0.1, kernel_size=9, act='gelu', rotary_embed=None,
- layer_idx=None, mix_ln_layer=None
+ layer_idx=None, mix_ln_layer=None, mixln_shuffle_speakers=False
  ):
  super().__init__()
  self.dropout = dropout
@@ -406,7 +411,7 @@ def __init__(self, c, num_heads, dropout, attention_dropout=0.1,
  and layer_idx in mix_ln_layer
  )
  if self.use_mix_ln:
- self.layer_norm1 = Mixed_LayerNorm(c, c)
+ self.layer_norm1 = Mixed_LayerNorm(c, c, shuffle_speakers=mixln_shuffle_speakers)
  else:
  self.layer_norm1 = LayerNorm(c)
  # Always use the in-house manual attention. With rotary_embed=None this
@@ -419,7 +424,7 @@ def __init__(self, c, num_heads, dropout, attention_dropout=0.1,
  c, num_heads, dropout=attention_dropout, bias=False, rotary_embed=rotary_embed
  )
  if self.use_mix_ln:
- self.layer_norm2 = Mixed_LayerNorm(c, c)
+ self.layer_norm2 = Mixed_LayerNorm(c, c, shuffle_speakers=mixln_shuffle_speakers)
  else:
  self.layer_norm2 = LayerNorm(c)
  self.ffn = TransformerFFNLayer(

modules/fastspeech/acoustic_encoder.py

@@ -58,7 +58,8 @@ def __init__(self, vocab_size):
  dropout=hparams['dropout'], num_heads=hparams['num_heads'],
  use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False),
  use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True),
- mix_ln_layer=self.mix_ln_layer
+ mix_ln_layer=self.mix_ln_layer,
+ mixln_shuffle_speakers=hparams.get('mixln_shuffle_speakers', False),
  )
 
  self.pitch_embed = AdamWLinear(1, hparams['hidden_size'])

modules/fastspeech/tts_modules.py

@@ -13,14 +13,15 @@
 
 class TransformerEncoderLayer(nn.Module):
  def __init__(self, hidden_size, dropout, kernel_size=None, act='gelu', num_heads=2, rotary_embed=None,
- layer_idx=None, mix_ln_layer=None):
+ layer_idx=None, mix_ln_layer=None, mixln_shuffle_speakers=False):
  super().__init__()
  self.op = EncSALayer(
  hidden_size, num_heads, dropout=dropout,
  attention_dropout=0.0, relu_dropout=dropout,
  kernel_size=kernel_size,
  act=act, rotary_embed=rotary_embed,
- layer_idx=layer_idx, mix_ln_layer=mix_ln_layer
+ layer_idx=layer_idx, mix_ln_layer=mix_ln_layer,
+ mixln_shuffle_speakers=mixln_shuffle_speakers
  )
 
  def forward(self, x, **kwargs):
@@ -373,7 +374,8 @@ def __init__(
  self, hidden_size, num_layers,
  ffn_kernel_size=9, ffn_act='gelu',
  dropout=None, num_heads=2, use_pos_embed=True, rel_pos=True,
- use_rope=False, rope_interleaved=True, mix_ln_layer=None
+ use_rope=False, rope_interleaved=True, mix_ln_layer=None,
+ mixln_shuffle_speakers=False
  ):
  super().__init__()
  self.num_layers = num_layers
@@ -394,7 +396,8 @@ def __init__(
  self.hidden_size, self.dropout,
  kernel_size=ffn_kernel_size, act=ffn_act,
  num_heads=num_heads, rotary_embed=rotary_embed,
- layer_idx=i, mix_ln_layer=mix_ln_layer
+ layer_idx=i, mix_ln_layer=mix_ln_layer,
+ mixln_shuffle_speakers=mixln_shuffle_speakers
  )
  for i in range(self.num_layers)
  ])