Adds config and weights of UNIVERSE++

config.yaml

@@ -36,49 +36,45 @@ datamodule:
       fs: 16000
       split: train
       audio_len: 2.0
-      augmentation: false
     vb-val-16k:
       _target_: open_universe.datasets.NoisyDataset
       audio_path: ${..vb-train-16k.audio_path}
       fs: ${..vb-train-16k.fs}
       split: val
       audio_len: null
-      augmentation: false
     vb-test-16k:
       _target_: open_universe.datasets.NoisyDataset
       audio_path: ${..vb-train-16k.audio_path}
       fs: ${..vb-train-16k.fs}
       split: test
       audio_len: null
-      augmentation: false
     vb-train-24k:
       _target_: open_universe.datasets.NoisyDataset
       audio_path: data/voicebank_demand/24k
       fs: 24000
       split: train
       audio_len: 2.0
-      augmentation: false
     vb-val-24k:
       _target_: open_universe.datasets.NoisyDataset
       audio_path: ${..vb-train-24k.audio_path}
       fs: ${..vb-train-24k.fs}
       split: val
       audio_len: null
-      augmentation: false
     vb-test-24k:
       _target_: open_universe.datasets.NoisyDataset
       audio_path: ${..vb-train-24k.audio_path}
       fs: ${..vb-train-24k.fs}
       split: test
       audio_len: null
-      augmentation: false
 model:
-  _target_: open_universe.networks.universe.Universe
+  _target_: open_universe.networks.universe.UniverseGAN
   fs: 16000
   normalization_norm: 2
   normalization_kwargs:
     ref: both
     level_db: -26.0
+  edm:
+    noise: 0.25
   score_model:
     _target_: open_universe.networks.universe.ScoreNetwork
     fb_kernel_size: 3
@@ -93,9 +89,9 @@ model:
     encoder_gru_conv_sandwich: false
     extra_conv_block: true
     decoder_act_type: prelu
-    use_weight_norm: false
-    seq_model: gru
-    use_antialiasing: false
+    use_weight_norm: true
+    use_antialiasing: true
+    time_embedding: simple
   condition_model:
     _target_: open_universe.networks.universe.ConditionerNetwork
     fb_kernel_size: ${model.score_model.fb_kernel_size}
@@ -107,7 +103,6 @@ model:
     extra_conv_block: ${model.score_model.extra_conv_block}
     decoder_act_type: prelu
     use_weight_norm: ${model.score_model.use_weight_norm}
-    seq_model: ${model.score_model.seq_model}
     use_antialiasing: false
   diffusion:
     schedule: geometric
@@ -116,17 +111,39 @@ model:
     n_steps: 8
     epsilon: 1.3
   losses:
+    multi_period_discriminator:
+      mpd_reshapes:
+      - 2
+      - 3
+      - 5
+      - 7
+      - 11
+      use_spectral_norm: false
+      discriminator_channel_mult: 1
+    multi_resolution_discriminator:
+      resolutions:
+      - - 1024
+        - 120
+        - 600
+      - - 2048
+        - 240
+        - 1200
+      - - 512
+        - 50
+        - 240
+      use_spectral_norm: false
+      discriminator_channel_mult: 1
+    disc_freeze_step: 0
     weights:
+      mel_l1: 45.0
       score: 1.0
-      signal: 1.0
-      latent: 1.0
-    mdn_n_comp: 3
-    mdn_alpha_per_sample: true
+    use_signal_decoupling: true
+    signal_decoupling_act: snake
     score_loss:
       _target_: torch.nn.MSELoss
   training:
     audio_len: ${datamodule.datasets.vb-train-16k.audio_len}
-    time_sampling: time_uniform
+    time_sampling: time_normal_0.95
     dynamic_mixing: false
     ema_decay: 0.999
   validation:
@@ -134,31 +151,52 @@ model:
     main_loss_mode: max
     n_bins: 5
     max_enh_batches: 4
-    num_tb_samples: 0
     enh_losses:
       val/:
         _target_: open_universe.metrics.EvalMetrics
         audio_fs: ${model.fs}
   optimizer:
-    _target_: torch.optim.AdamW
-    lr: 0.0002
-    weight_decay: 0.01
-    weight_decay_exclude:
-    - prelu
-    - bias
-    lr_warmup: null
-    betas:
-    - 0.8
-    - 0.99
+    accumulate_grad_batches: 1
+    generator:
+      _target_: torch.optim.AdamW
+      lr: 0.0002
+      weight_decay: 0.01
+      betas:
+      - 0.8
+      - 0.99
+      weight_decay_exclude:
+      - prelu
+      - bias
+    discriminator:
+      _target_: torch.optim.AdamW
+      lr: 0.0002
+      betas:
+      - 0.8
+      - 0.99
+    grad_clip_vals:
+      mrd: 1000.0
+      mpd: 1000.0
+      score: 1000.0
+      cond: 1000.0
   scheduler:
-    scheduler:
-      _target_: open_universe.utils.schedulers.LinearWarmupCosineAnnealingLR
-      T_warmup: 50000
-      T_cosine: 50001
-      eta_min: 1.6e-06
-      T_max: ${trainer.max_steps}
-    interval: step
-    frequency: 1
+    generator:
+      scheduler:
+        _target_: open_universe.utils.schedulers.LinearWarmupCosineAnnealingLR
+        T_warmup: 20000
+        T_cosine: 400000
+        eta_min: 1.6e-06
+        T_max: ${trainer.max_steps}
+      interval: step
+      frequency: 1
+    discriminator:
+      scheduler:
+        _target_: open_universe.utils.schedulers.LinearWarmupCosineAnnealingLR
+        T_warmup: 20000
+        T_cosine: 400000
+        eta_min: 1.6e-06
+        T_max: ${trainer.max_steps}
+      interval: step
+      frequency: 1
   grad_clipper:
     _target_: open_universe.utils.FixedClipper
     max_norm: 1000.0
@@ -167,7 +205,7 @@ trainer:
   accumulate_grad_batches: 1
   min_epochs: 1
   max_epochs: -1
-  max_steps: 300000
+  max_steps: 600000
   deterministic: warn
   accelerator: gpu
   devices: -1

weights.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ee7b031d055bd65d1e849426ba7867bf1416b53adf46e32c4a69312768361222
-size 901069356
+oid sha256:0d90ab343c86501a23d5dd0011242d1129ad2f54d8cebec68c55dd387037879c
+size 1025936580