Update config.gin

config.gin

@@ -1,74 +1,56 @@
-MASK_AMOUNT = 150
-MASK_GAP_SIZE = 15
-MASKED_LOSS_WEIGHT = 0.9
-MODEL_DIM = 768
-NUM_DECODER_HEADS = 12
-NUM_DECODER_LAYERS = 2
-NUM_ENCODEC_TARGETS = 8
-NUM_ENCODER_HEADS = 12
-NUM_ENCODER_LAYERS = 10
-NUM_TARGET_TOKENS = 1024
-NUM_TOTAL_TARGETS = 8
-
-# Parameters for EncodecMAE:
-# ==============================================================================
-EncodecMAE.decoder = @decoder/tasks.models.transformers.TransformerEncoder
-EncodecMAE.head = @tasks.models.encodecmae.heads.FrameLevelClassificationHead
-EncodecMAE.lr_scheduler = None
-EncodecMAE.masked_weight = %MASKED_LOSS_WEIGHT
-EncodecMAE.masker = @tasks.models.encodecmae.masking.TimeGapMask
-EncodecMAE.optimizer = @torch.optim.AdamW
-EncodecMAE.positional_encoder = \
-    @tasks.models.transformers.SinusoidalPositionalEmbeddings
-EncodecMAE.quantizer_weights = \
-    [0.22407463,
-     0.1759858,
-     0.14499009,
-     0.12150037,
-     0.10315603,
-     0.08831368,
-     0.07608274,
-     0.06589669]
-EncodecMAE.target_encoder = @tasks.models.encodecmae.targets.EncodecQuantizer
-EncodecMAE.visible_encoder = @encoder/tasks.models.transformers.TransformerEncoder
-EncodecMAE.wav_encoder = @tasks.models.encodecmae.encoders.EncodecEncoder
-
-# Parameters for EncodecQuantizer:
-# ==============================================================================
-EncodecQuantizer.n = %NUM_ENCODEC_TARGETS
-
-# Parameters for decoder/MultiHeadAttention:
-# ==============================================================================
-decoder/MultiHeadAttention.model_dim = %MODEL_DIM
-decoder/MultiHeadAttention.num_heads = %NUM_DECODER_HEADS
-
-# Parameters for encoder/MultiHeadAttention:
-# ==============================================================================
-encoder/MultiHeadAttention.model_dim = %MODEL_DIM
-encoder/MultiHeadAttention.num_heads = %NUM_ENCODER_HEADS
-
-# Parameters for SinusoidalPositionalEmbeddings:
-# ==============================================================================
-SinusoidalPositionalEmbeddings.embedding_dim = %MODEL_DIM
-
-# Parameters for decoder/TransformerEncoder:
-# ==============================================================================
-decoder/TransformerEncoder.attention_layer = \
-    @decoder/tasks.models.transformers.MultiHeadAttention
-decoder/TransformerEncoder.compile = True
-decoder/TransformerEncoder.model_dim = %MODEL_DIM
-decoder/TransformerEncoder.num_layers = %NUM_DECODER_LAYERS
-
-# Parameters for encoder/TransformerEncoder:
-# ==============================================================================
-encoder/TransformerEncoder.attention_layer = \
-    @encoder/tasks.models.transformers.MultiHeadAttention
-encoder/TransformerEncoder.compile = True
-encoder/TransformerEncoder.model_dim = %MODEL_DIM
-encoder/TransformerEncoder.num_layers = %NUM_ENCODER_LAYERS
-
-# Parameters for FrameLevelClassificationHead:
-# ==============================================================================
-FrameLevelClassificationHead.model_dim = %MODEL_DIM
-FrameLevelClassificationHead.num_streams = %NUM_TOTAL_TARGETS
-FrameLevelClassificationHead.num_tokens = %NUM_TARGET_TOKENS
+NUM_ENCODEC_TARGETS=8
+NUM_TOTAL_TARGETS=8
+NUM_TARGET_TOKENS=1024
+MASK_AMOUNT=150
+MASK_GAP_SIZE=15
+MASK_PROP=0.5
+MODEL_DIM=768
+NUM_ENCODER_LAYERS=10
+NUM_ENCODER_HEADS=12
+NUM_DECODER_LAYERS=2
+NUM_DECODER_HEADS=12
+MASKED_LOSS_WEIGHT=0.9
+get_model.model=@models.EncodecMAE
+models.EncodecMAE:
+    wav_encoder = @models.encodecmae.encoders.EncodecEncoder
+    target_encoder = @models.encodecmae.targets.EncodecQuantizer
+    masker = @models.encodecmae.masking.TimeGapMask
+    visible_encoder = @encoder/models.transformers.TransformerEncoder
+    positional_encoder = @models.transformers.SinusoidalPositionalEmbeddings
+    decoder = @decoder/models.transformers.TransformerEncoder
+    head = @models.encodecmae.heads.FrameLevelClassificationHead
+    optimizer=@torch.optim.AdamW
+    lr_scheduler=None
+    masked_weight=%MASKED_LOSS_WEIGHT
+    quantizer_weights=[0.22407463, 0.1759858 , 0.14499009, 0.12150037, 0.10315603, 0.08831368, 0.07608274, 0.06589669]
+torch.optim.AdamW:
+    lr=%PRETRAIN_MAX_LR
+    betas=(0.9,0.95)
+    weight_decay=0.05
+models.encodecmae.targets.EncodecQuantizer:
+    n = %NUM_ENCODEC_TARGETS
+models.encodecmae.masking.TimeGapMask:
+    mask_amount = %MASK_AMOUNT
+    gap_size = %MASK_GAP_SIZE
+    mask_prop = %MASK_PROP
+encoder/models.transformers.TransformerEncoder:
+    model_dim=%MODEL_DIM
+    num_layers=%NUM_ENCODER_LAYERS
+    attention_layer=@encoder/models.transformers.MultiHeadAttention
+    compile=True
+encoder/models.transformers.MultiHeadAttention:
+    model_dim=%MODEL_DIM
+    num_heads=%NUM_ENCODER_HEADS
+decoder/models.transformers.TransformerEncoder:
+    model_dim=%MODEL_DIM
+    num_layers=%NUM_DECODER_LAYERS
+    attention_layer=@decoder/models.transformers.MultiHeadAttention
+    compile=True
+decoder/models.transformers.MultiHeadAttention:
+    model_dim=%MODEL_DIM
+    num_heads=%NUM_DECODER_HEADS
+models.transformers.SinusoidalPositionalEmbeddings.embedding_dim = %MODEL_DIM
+models.encodecmae.heads.FrameLevelClassificationHead:
+    model_dim=%MODEL_DIM
+    num_tokens=%NUM_TARGET_TOKENS
+    num_streams=%NUM_TOTAL_TARGETS