amt/src/config/config.py

"""config.py"""
import numpy as np
import torch

if torch.cuda.is_available():
    TEST_BSZ = 128
else:
    TEST_BSZ = 16
# yapf: disable
"""
audio_cfg:
- Used by 'ymt3' to create a spectrogram layer.
- Input shape of model is determined by audio_cfg.
- 'train.py' arguments can override these defaults.
"""
audio_cfg = {
    # Overwrittable by args in train.py
    "codec": "melspec",  # {melspec, spec} melspec for MT3, spec for PerceiverTF
    "hop_length": 128,  # {128, 300} 128 for MT3, 300 for PerceiverTF
    # Shared audio parameters
    "audio_backend": "torchaudio",  # {torchaudio, nnAudio}
    "sample_rate": 16000,
    "input_frames": 32767, # number of input frames (~=2.048 s), determining in-/output shape of front layers. 
    "n_fft": 2048,
    "n_mels": 512,  # only for melspec
    "f_min": 50.0,
    "f_max": 8000.0,
} # TODO: currently dataloader is not updated by "input_frames"

"""
model_cfg:
- Encoder type dictates use of T5_CFG or PERCEIVER_TF_CFG.
- 'train.py' arguments can override these defaults.
"""
model_cfg = {
    "encoder_type": "t5",  # {"t5", "perceiver-tf", "conformer"}
    "decoder_type": "t5", # {"t5", "multi-t5"}
    "pre_encoder_type": "default",  # {None, "default", "conv", "conv1d", "conv2d_avpt"} by default, t5:None, perceiver:conv.
    "pre_encoder_type_default": {"t5": None, "perceiver-tf": "conv", "conformer": None},
    "pre_decoder_type": "default", # {None, 'linear', 'conv1', 'mlp', 'group_linear'} see model/projection_layer.py
    "pre_decoder_type_default": { # [enc_type][dec_type]
        "t5": {"t5": None,},
        "perceiver-tf": {"t5": "linear", "multi-t5": "mc_shared_linear"},
        "conformer": {"t5": None,},
    },
    "conv_out_channels": 128, # number of filters for 'conv' pre_encoder. Otherwise ignored.
    "t5_basename": "google/t5-v1_1-small",
    "pretrained": False, # bool, if True, load pretrained weights from t5_basename. Mismatched layers are ignored.
    "use_task_conditional_encoder": True, # True by default, but default task is None. So not activated by default. 
    "use_task_conditional_decoder": True, # True by default, but default task is None. So not activated by default.  
    "d_feat": "auto", # Input audio feature dimension for encoder. Automatically inferred by audio_cfg and existence of pre_encoders.
    "tie_word_embeddings": True, # If True, weights of embed_tokens and lm_head are tied for stabilizing gradients. 
    "vocab_size": "auto", # int or "auto", automatically inferred by task manager.
    "num_max_positions": "auto", # int or "auto". Length of positional encoding. Automatically inferred by "feat_length", "event_length" and task_manager.max_task_token_length.
    # 'vocab_size', 'tie_word_embeddings' and 'num_max_positions' are auto-copied to encoder and decoder configs in the below.
    "encoder": {
        "t5": {
            "d_model": 512, # Hidden size of T5 encoder. 
            "num_heads": 6,
            "num_layers": 8,
            "dropout_rate": 0.05,
            "position_encoding_type": "sinusoidal", # {'sinusoidal', 'trainable'}.
            "ff_widening_factor": 2, # wideening factor for MLP/MoE layers. Default is 2 in T5.
            "ff_layer_type": "t5_gmlp", # {'t5_gmlp', 'moe', 'mlp', 'gmlp'}. 'moe' for mixture of experts, 'mlp' for standard transformer dense layer, 'gmlp' for simple gated MLP.
        },
        "perceiver-tf": {
            "num_latents": 24, # number of latents in Perceiver. 24 in perceiver-tf paper.
            "d_latent": 128, # latent dimension of Perceiver. 128 in perceiver-tf paper.
            "d_model": "q", # int or "q" or "kv". Inner-dim of sca and local/temporal self-att.
                # "q" follows "latent_dim". "kv" follows  "d_feat". Best practice is to inc-/decrease 'd_latent', instead of 'd_model'.
            "num_blocks": 3, # number of Perceiver-TF blocks in encoder. L in the paper.
            "num_local_transformers_per_block": 2, # N in the paper.
            "num_temporal_transformers_per_block": 2,  # M in the paper.
            "sca_use_query_residual": False,
            "dropout_rate": 0.1,
            "position_encoding_type": "trainable", # {'trainable', 'rotary', 'alibi', 'alibit', None, 'tkd','td', 'tk', 'kdt'}. alibit is alibi with trainable slopes.
            "attention_to_channel": True, # Whether to use channel attention in sca.
            "layer_norm_type": "layer_norm", # {'layer_norm', 'rms_norm'}
            "ff_layer_type": "mlp", # {'moe', 'mlp', gmlp}. 'moe' for mixture of experts, 'mlp' for standard transformer dense layer, 'gmlp' for simple gated MLP.
            "ff_widening_factor": 1, # wideening factor for MLP/MoE layers. Default is 1.
            "moe_num_experts": 4, # number of experts in MoE layer. Default is 4. Disabled if ff_layer_type is not 'moe'.
            "moe_topk": 2, # top-k routing in MoE layer. Default is 2. Disabled if ff_layer_type is not 'moe'.
            "hidden_act": 'gelu', # activation function in MLP/MoE layer. Default is 'gelu'. {'gelu', 'silu', 'relu'}
            "rotary_type_sca": "pixel", # {'l'|'lang', 'p'|'pixel'}. Default is 'pixel'.
            "rotary_type_latent": "pixel", # {'l'|'lang', 'p'|'pixel'}. Default is 'pixel'.
            "rotary_type_temporal": "lang", # {'l'|'lang', 'p'|'pixel'}. Default is 'lang'.
            "rotary_apply_to_keys": False, # Whether to apply rotary to keys. Default is False.
            "rotary_partial_pe": False, # Whether to use partial positional encoding. Default is False.
        },
        "conformer": {
            "d_model": 512, # Hidden size of T5 encoder. 
            "intermediate_size": 512, # or 2048. size of the intermediate feed forward layer in each T5Block
            "num_heads": 8,
            "num_layers": 8,
            "dropout_rate": 0.1,
            "layerdrop": 0.1, # see https://arxiv.org/abs/1909.11556
            "position_encoding_type": "rotary", # {'rotary', 'relative'}. 
            "conv_dim": (512, 512, 512, 512, 512, 512, 512),
            "conv_stride": (5, 2, 2, 2, 2, 2, 2),
            "conv_kernel": (10, 3, 3, 3, 3, 3, 3),
            "conv_depthwise_kernel_size": 31,
        },

    },
    "decoder": {
        "t5": {
            "d_model": 512, # Hidden size of T5 encoder. If encoder has lower dim, it is projected to this dim for enc-dec cross att.
            "num_heads": 6,
            "num_layers": 8,
            "dropout_rate": 0.05,
            "position_encoding_type": "sinusoidal", # {'sinusoidal', 'trainable'}.
            "ff_widening_factor": 2, # wideening factor for MLP/MoE layers. Default is 2 in T5.
            "ff_layer_type": "t5_gmlp", # {'t5_gmlp', 'moe', 'mlp', 'gmlp'}. 'moe' for mixture of experts, 'mlp' for standard transformer dense layer, 'gmlp' for simple gated MLP.
        },
        "multi-t5": {
            "d_model": 512, # Hidden size of T5 encoder. Recommended: {256 or 512}
            "num_heads": 6,
            "num_layers": 8,
            "dropout_rate": 0.05,
            "position_encoding_type": "sinusoidal", # {'sinusoidal', 'trainable'}.
            "ff_widening_factor": 2, # wideening factor for MLP/MoE layers. Default is 2 in T5.
            "ff_layer_type": "t5_gmlp", # {'t5_gmlp', 'moe', 'mlp', 'gmlp'}. 'moe' for mixture of experts, 'mlp' for standard transformer dense layer, 'gmlp' for simple gated MLP.
            "num_channels": 13,
        },
    },
    "feat_length": "auto", # Input audio feature length for encoder. Automatically inferred by audio_cfg.
        # mt3: 256 time steps
    "event_length": 1024,  # max length of event tokens excluding task tokens <-- 128 for multi-t5
    "init_factor": 1.0, # initialization factor for embedding layers
}

# yapf: enable
shared_cfg = {
    "PATH": {
        "data_home": "../../data", # path to the data directory. If using relative path, it is relative to /src directory.
    },
    "BSZ": { # global batch size is local_bsz * n_GPUs in DDP mode
        "train_sub": 12, #20, # sub-batch size is per CPU worker
        "train_local": 24, #40, # local batch size is per GPU in DDP mode
        "validation": 64, # validation batch size is per GPU in DDP mode
        "test": TEST_BSZ,
    },
    "AUGMENTATION": {
        "train_random_amp_range": [0.8, 1.1], # min and max amplitude scaling factor
        "train_stem_iaug_prob": 0.7, # probability of stem activation in intra-stem augmentation
        "train_stem_xaug_policy": {
            "max_k": 3,
            "tau": 0.3,
            "alpha": 1.0,
            "max_subunit_stems": 12, # the number of subunit stems to be reduced to this number of stems
            "p_include_singing": None,  # NOT IMPLEMENTED; probability of including singing for cross augmented examples. if None, use base probaility.
            "no_instr_overlap": True,
            "no_drum_overlap": True,
            "uhat_intra_stem_augment": True,
        },
        "train_pitch_shift_range": [-2, 2], # [min, max] in semitones. None or [0, 0] for no pitch shift.
    },
    "DATAIO": { # do not set `shuffle` here. 
        "num_workers": 4, # num_worker is per GPU in DDP mode
        "prefetch_factor": 2, #2,
        "pin_memory": True,
        "persistent_workers": False,
    },
    "CHECKPOINT": {
        "save_top_k": 4, # max top k checkpoints to save
        "monitor": 'validation/macro_onset_f',
        "mode": 'max',
        # "every_n_epochs": 20, # only working when check_val_every_n_epoch is 0
        "save_last": True, # save last model
        "filename": "{epoch}-{step}",
    },
    "TRAINER": { # do not coverwrite args in this section
        "limit_train_batches": 1.0, # How much of training dataset to check (float = fraction, int = num_batches)
        "limit_val_batches": 1.0,
        "limit_test_batches": 1.0,
        "gradient_clip_val": 1.0, # {0 or None} means don't clip.
        "accumulate_grad_batches": 1, #1, # Accumulates grads every k batches. If set to 1, no effect.
        "check_val_every_n_epoch": 1, #5, 1 for very large dataset such as EGMD
        "num_sanity_val_steps": 0,
    },
    "WANDB": {
        # "save_dir": "../logs",
        "save_dir": "amt/logs", # modified for huggingface spaces...
        "cache_dir": "../logs/.wandb_cache",
        "resume": "allow",
        "anonymous": "allow", # {never, allow, must}
        "mode": "online", # {online, offline, disabled}
    },
    "LR_SCHEDULE": {
        # "scheduler_type": "cosine", # {legacy, cosine, constant}
        "warmup_steps": 1000, # only for cosine scheduler, legacy scheduler follows T5's legacy schedule
        "total_steps": 100000, # argparser of train.py can overwrite this
        "final_cosine": 1e-5, # only for cosine scheduler
    },
    "TOKENIZER": {
        "max_shift_steps": "auto", # max number of shift steps in the model. (int) or "auto". If "auto", it is set by audio_cfg["input_frames"] and shift_steps_ms. 206 with default setup.
        "shift_step_ms": 10, # shift step in ms
    },
}

T5_BASE_CFG = {
    "google/t5-v1_1-small": {
        "architectures": ["T5ForConditionalGeneration"],
        "d_ff":
            1024,  # size of the intermediate feed forward layer in each T5Block. Can be overwrten by ff_widening_factor in model_cfg.
        "d_kv": 64,  # d_kv has to be equal to d_model // num_heads.
        # "d_model": 512,  # encoder hiddnen size, defined by model_cfg
        "decoder_start_token_id": 0,
        "dense_act_fn": "gelu_new",
        # "dropout_rate": 0.05,  # can be overwritten by args in ymt3
        "eos_token_id": 1,
        "feed_forward_proj": "gated-gelu",
        "initializer_factor": 1.0,
        "is_encoder_decoder": True,
        "is_gated_act": True,
        "layer_norm_epsilon": 1e-06,
        "model_type": "t5",
        # "num_decoder_layers": 8, # defined by model_cfg
        # "num_heads": 6,  # defined by model_cfg
        # "num_layers": 8,  # defined by model_cfg
        "output_past": True,
        "pad_token_id": 0,
        "relative_attention_num_buckets": 32,
        # "tie_word_embeddings": True,
        "use_cache": True,
        # "vocab_size": 1391 # vocab_size is automatically set by the task manager...
    },
    "google/t5-efficient-small": {
        "architectures": ["T5ForConditionalGeneration"],
        "d_ff": 2048,
        "d_kv": 64,
        "d_model": 512,
        "decoder_start_token_id": 0,
        "dropout_rate": 0.1,
        "eos_token_id": 1,
        "feed_forward_proj": "relu",
        "initializer_factor": 1.0,
        "is_encoder_decoder": True,
        "layer_norm_epsilon": 1e-06,
        "model_type": "t5",
        "num_decoder_layers": 6,
        "num_heads": 8,
        "num_layers": 6,
        "pad_token_id": 0,
        "relative_attention_num_buckets": 32,
        "torch_dtype": "float32",
        "transformers_version": "4.17.0.dev0",
        "use_cache": True,
    },
}

# yapf: enable
DEEPSPEED_CFG = {
    "zero_allow_untested_optimizer": True,
    "optimizer": {
        "type": "adam",
        "params": {
            "lr": 1e-4,
            "betas": [0.998, 0.999],
            "eps": 1e-3,
            "weight_decay": 0.001,
            "adam_w_mode": True,
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "last_batch_iteration": -1,
            "warmup_min_lr": 0,
            "warmup_max_lr": 3e-5,
            "warmup_num_steps": 100,
        },
    },
    "zero_optimization": {
        "stage": 0,  #0,1,2,3
        # "offload_optimizer":
        #     False,  # Enable Offloading optimizer state/calculation to the host CPU
    },
}