init

german-1st/events.out.tfevents.1626192156.instance-3.21320.3.v2

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:eb69ed2d3a263b3e1161b25ce2cb012898003f1f56e7d8b6114511f85cc161b8
-size 40

german-1st/{events.out.tfevents.1625643205.t1v-n-3abeb69a-w-0.838585.3.v2 → events.out.tfevents.1626262981.t1v-n-3abeb69a-w-0.12272.3.v2}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6160222312f88827d6178c3853ae198991de9b0a3cab6bdd6e9b257c77909892
+oid sha256:73239178ccdb9423fbc79388ce2a835996b379734f7c7d160bf4ee2822d34e0a
 size 40

run_german.sh

@@ -16,7 +16,6 @@
     --max_duration_in_seconds="10.0" \
     --adam_beta1="0.9" \
     --adam_beta2="0.98" \
-		--dtype="bfloat16" \
-		--cache_dir="./wav2vec2/" \
+		--adafactor \
 		--pad_to_multiple_of="16384" \
     --push_to_hub

run_wav2vec2_pretrain_flax.py

@@ -275,7 +275,7 @@ def main():
     )
 
     # save vectorized dataset once
-    vectorized_datasets = datasets.load_from_disk("/home/german-common-voice-processed/normalized")
+    vectorized_datasets = datasets.load_from_disk("/home/wav2vec2-experiments/datasets/german-common-voice-processed/normalized/")
 
     # pretraining is only supported for "newer" stable layer norm architecture
     # apply_spec_augment has to be True, mask_feature_prob has to be 0.0
@@ -353,17 +353,24 @@ def main():
         return traverse_util.unflatten_dict(flat_mask)
 
     # create adam optimizer
-    adamw = optax.adamw(
-        learning_rate=linear_decay_lr_schedule_fn,
-        b1=training_args.adam_beta1,
-        b2=training_args.adam_beta2,
-        eps=training_args.adam_epsilon,
-        weight_decay=training_args.weight_decay,
-        mask=decay_mask_fn,
-    )
+    if training_args.adafactor:
+        # We use the default parameters here to initialize adafactor,
+        # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
+        optimizer = optax.adafactor(
+            learning_rate=linear_decay_lr_schedule_fn,
+        )
+    else:
+        optimizer = optax.adamw(
+            learning_rate=linear_decay_lr_schedule_fn,
+            b1=training_args.adam_beta1,
+            b2=training_args.adam_beta2,
+            eps=training_args.adam_epsilon,
+            weight_decay=training_args.weight_decay,
+            mask=decay_mask_fn,
+        )
 
     # Setup train state and define training hyper-parameters
-    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw)
+    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer)
     num_negatives = model.config.num_negatives
     contrastive_logits_temperature = model.config.contrastive_logits_temperature
     num_codevectors = model.config.num_codevectors_per_group * model.config.num_codevector_groups