config.yaml

notes: "streaming model = online (Mamba) + 8k&16k"
# ============================================================================ #
config:
    # ========================== Key Variables ============================== #
    train_phase: &var_tr_phase "adversarial"  # pretrain / adversarial
    dataset_phase: "to48k"
    train_phase_list: ['pretrain_to48k', 'adversarial_to48k']

    # -- architecture --
    online: &var_online True           # False=TF-Locoformer (offline), True=Mamba (online)
    d_model_enc: &var_model_channels_1st 128
    d_model_dec: &var_model_channels_2nd 64
    d_hidden_enc: &var_model_hidden_1st 384
    d_hidden_dec: &var_model_hidden_2nd 192
    n_head: &var_n_head 4
    kernel_size: &var_kernel_size 7
    rope_theta: &var_rope_theta 10000
    num_frequency_bin_max: &num_frequency_bin_max 961

    # -- mamba (online) --
    d_state: &var_d_state 16
    d_conv_mamba: &var_kernel_mamba 3
    expand_mamba: &var_expand_mamba 4

    # -- global --
    fs_list: &var_fs_list ['8000', '16000', '22050', '24000', '32000', '44100', '48000']
    ssl_model_key: &var_SSL_model_key "microsoft/wavlm-large"
    win_48k: &var_win_48k [1920]
    win_48k_disc: &var_win_48k_disc [20, 40, 60, 80, 100]

    # ============================== Dataset =============================== #
    dataset:
        to48k:
            max_len: 3
            sample_rate_src: 48000
            sample_rate_in: 16000
            db_root: "/home/DB/VCTK"
            rir_dir: "/home/DB/DNS_RIR_48k"
            scp_dir: "data/scp/scp_VCTK"
            train:
                spk: "tr_s.scp"
                noise: "tr_n.scp"
            valid:
                spk: "cv_s.scp"
                noise: "tr_n.scp"
        synthesis_config:
            multi_spk_prob: 0.2
            rir:
                prob: 0.5
                rir_sidelobe: 1
            noise:
                SNR_range: [0, 20]
                c_SNR_range: [0, 20]
                c_beta_range: [0.5, 1.5]
            BPF:
                prob: 0.5
                fir_filter_beta: [0.25, 1.0]
                low_cutoff_freq_range: [1000, 3000]
            clipping:
                prob: 0.5
                clipping_level_range: [-15, 0]
            level:
                target_dB_FS: [-35, -15]
    dataset_test:
        testset_key: "VCTK_SR"
        tensorboard_logging: true
        input_eval: true
        output_eval: true
        VCTK_SR:
            clean_dir: "/home/DB/VCTK-Corpus-0.92/wav48_silence_trimmed/test"
            noisy_dir: "/home/DB/VCTK-Corpus-0.92/wav48_silence_trimmed/test"
        VoxCeleb:
            sample_rate_src: 48000
    # ============================== Dataloader ============================ #
    dataloader:
        batch_size: 2
        pin_memory: false
        num_workers: 0
        drop_last: false
    # ================================ STFT ================================ #
    stft:
        frame_length: 40
        frame_shift: 20
    # =============================== Model ================================ #
    model:
        online: *var_online
        input_embedding:
            online: *var_online
            d_model: *var_model_channels_1st
            d_freq: *num_frequency_bin_max
            freq_pe: True
        freq_linear:
            seq_len: *num_frequency_bin_max
            proj_len: 512
            n_heads: *var_n_head
            kv_shared: True
        encoder_stage:
            block_type: 'Encoder'
            RoPE:
                d_model: *var_model_channels_1st
                n_head: *var_n_head
                theta: *var_rope_theta
            TF_block_Stage:
                online: *var_online
                time_module:
                    offline:
                        d_model: *var_model_channels_1st
                        d_hidden: *var_model_hidden_1st
                        n_head: *var_n_head
                        kernel_size: *var_kernel_size
                        dropout_rate: 0.00
                    online:
                        d_model: *var_model_channels_1st
                        d_state: *var_d_state
                        d_conv: *var_kernel_mamba
                        expand: *var_expand_mamba
                        dropout_rate: 0.00
                freq_module:
                    d_model: *var_model_channels_1st
                    d_hidden: *var_model_hidden_1st
                    n_head: *var_n_head
                    kernel_size: *var_kernel_size
                    dropout_rate: 0.00
            num_repeat: 6
        freq_upsampler:
            d_model: *var_model_channels_1st
            d_model_out: *var_model_channels_2nd
            d_freq_min: 161
            d_freq_max: *num_frequency_bin_max
        decoder_stage:
            block_type: 'Decoder'
            RoPE:
                d_model: *var_model_channels_2nd
                n_head: *var_n_head
                theta: *var_rope_theta
            TF_block_Stage:
                online: *var_online
                time_module:
                    offline:
                        d_model: *var_model_channels_2nd
                        d_hidden: *var_model_hidden_2nd
                        n_head: *var_n_head
                        kernel_size: *var_kernel_size
                        dropout_rate: 0.00
                    online:
                        d_model: *var_model_channels_2nd
                        d_state: *var_d_state
                        d_conv: *var_kernel_mamba
                        expand: *var_expand_mamba
                        dropout_rate: 0.00
                freq_module:
                    d_model: *var_model_channels_2nd
                    d_model_kv: *var_model_channels_1st
                    d_hidden: *var_model_hidden_2nd
                    n_head: *var_n_head
                    kernel_size: *var_kernel_size
                    dropout_rate: 0.00
            num_repeat: 3
        output_spec:
            online: *var_online
            d_model: *var_model_channels_2nd
    # ============================== Training ============================== #
    engine:
        prob_effect:
            downsample_8k: 0.25
            codec: 0.3
            crystalizer: 0.15
            flanger: 0.05
            crusher: 0.1
        subset:
            train:
                subset: true
                num_per_epoch: 20000
            valid:
                subset: true
                num_per_epoch: 2000
        pretrain_to48k:
            downsample_src:
                prob: 0.6
                fs_list_src: [16000, 24000, 44100]
            loss_enhance:
                tau: 1.0e-4
                window_size: *var_win_48k
            loss_time:
                beta: 1.0e-3
            loss_rep:
                model_key: *var_SSL_model_key
                resampler:
                    orig_freq: 48000
                    new_freq: 16000
            loss_weight:
                se: 1.0
                time: 0.0
                ssl: 1.0e+2
            RandSpecMasking:
                t_len: [0, 10]
                f_len: [0, 10]
                t_num: [0, 2]
                f_num: [0, 3]
        adversarial_to48k:
            downsample_src:
                prob: 0.6
                fs_list_src: [16000, 24000, 44100]
            loss_enhance:
                tau: 1.0e-4
                window_size: *var_win_48k
            loss_time:
                beta: 1.0e-3
            loss_rep:
                model_key: *var_SSL_model_key
                resampler:
                    orig_freq: 48000
                    new_freq: 16000
            loss_weight:
                se: 1.0
                time: 0.0
                ssl: 1.0e+2
                gan: 1.0e-3
                fm: 0.1
                pesq: 1.0e-4
            RandSpecMasking:
                t_len: [0, 10]
                f_len: [0, 10]
                t_num: [0, 1]
                f_num: [0, 1]
            msstftd:
                filters: 32
                n_ffts_ms: *var_win_48k_disc
                fs_list: *var_fs_list
        sample_validation:
            - 'data/valid_sample/UNIVERSE_sample/0015.wav'
            - 'data/valid_sample/UNIVERSE_sample/0034.wav'
            - 'data/valid_sample/UNIVERSE_sample/0061.wav'
            - 'data/valid_sample/UNIVERSE_sample/0090.wav'
            - 'data/valid_sample/UNIVERSE_sample/0096.wav'
            - 'data/valid_sample/UNIVERSE_sample/0098.wav'
            - 'data/valid_sample/Real_sample/0.wav'
            - 'data/valid_sample/Real_sample/1.wav'
            - 'data/valid_sample/Real_sample/2.wav'
            - 'data/valid_sample/VoxCeleb_sample/f77-id10281-ni6gO5jDLJE-00010.wav'
            - 'data/valid_sample/VoxCeleb_sample/f83-id10282-hgB5ziAudzU-00001.wav'
            - 'data/valid_sample/VoxCeleb_sample/m87-id10271-PfcJLmkhGbk-00007.wav'
            - 'data/valid_sample/VoxCeleb_sample/m89-id10304-jUSC4i_eGHs-00002.wav'
            - 'data/valid_sample/VoxCeleb_sample/m99-id10297-FvbLoirHpx0-00006.wav'
        optimizer:  # alternatives: Adam, SGD
            name: "AdamW"
            AdamW:
                lr: 2.0e-4
                betas: [0.9, 0.995]
                weight_decay: 1.0e-2
        optimizer_D:  # discriminator optimizer
            name: "AdamW"
            AdamW:
                lr: 2.0e-4
                betas: [0.8, 0.999]
                weight_decay: 1.0e-2
        scheduler:  # alternatives: ReduceLROnPlateau, CosineAnnealingLR
            name: "StepLR"
            WarmupConstantSchedule:
                warmup_steps: 10000
            StepLR:
                step_size: 1
                gamma: 0.9
        max_epoch:
            pretrain_to48k: 20
            adversarial_to48k: 20
        gpuid: "1"
        clip_norm: 10
        start_scheduling:
            pretrain_to48k: 10
            adversarial_to48k: 1
    # ============================= Inference ============================== #
    inference:
        alpha: 0.2
        max_iter: 3