Adding proper data_loader, shuffling, evaluation, lr_schedulers. Refactor.

run_flax_speech_recognition_seq2seq_streaming_v3.py

@@ -18,6 +18,7 @@ Fine-tuning the Flax library models for sequence to sequence speech recognition.
 """
 # You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
 
+import itertools
 import logging
 import math
 import os
@@ -26,7 +27,7 @@ import time
 from dataclasses import field
 from functools import partial
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Union
 
 import datasets
 import flax
@@ -264,7 +265,6 @@ class DataTrainingArguments:
     )
     num_train_steps: int = field(default=50000, metadata={
                                  "help": "The number of training steps."})
-    # num_eval_samples: int = field(default=50000, metadata={"help": "The number of samples to be used for evaluation"})
     shuffle_buffer_size: Optional[int] = field(
         default=500,
         metadata={
@@ -402,29 +402,28 @@ def load_maybe_streaming_dataset(dataset_name, dataset_config_name, split="train
         return dataset
 
 
-def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False, drop_last=True):
+def collate_batch(samples):
+    return {key: [feature[key] for feature in samples] for key in samples[0]}
+
+
+def data_loader(
+    dataset: Dataset,
+    batch_size: int,
+    drop_last: bool=True,
+    num_workers: int=0,
+) -> Generator:
     """
     Returns batches of size `batch_size` from `dataset`. If `drop_last` is set to `False`, the final batch may be incomplete,
     and range in size from 1 to `batch_size`. Shuffle batches if `shuffle` is `True`.
     """
-    if shuffle:
-        batch_idx = jax.random.permutation(rng, len(dataset))
-        batch_idx = np.asarray(batch_idx)
-    else:
-        batch_idx = np.arange(len(dataset))
-
-    if drop_last:
-        steps_per_epoch = len(dataset) // batch_size
-        # Skip incomplete batch.
-        batch_idx = batch_idx[: steps_per_epoch * batch_size]
-        batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
-    else:
-        steps_per_epoch = math.ceil(len(dataset) / batch_size)
-        batch_idx = np.array_split(batch_idx, steps_per_epoch)
-
-    for idx in batch_idx:
-        batch = dataset[idx]
-        yield batch
+    data_loader_iterator = iter(torch.utils.data.DataLoader(
+        batch_size=batch_size,
+        dataset=dataset.with_format("torch"),
+        num_workers=num_workers,
+        collate_fn=collate_batch,
+        drop_last=drop_last,
+    ))
+    return data_loader_iterator
 
 
 class TrainState(train_state.TrainState):
@@ -448,13 +447,13 @@ def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
 
 
 def create_learning_rate_fn(
-    num_train_steps: int, num_warmup_steps: int, learning_rate: float
+    num_train_steps: int, num_warmup_steps: int, learning_rate: float, warmup_init_value: float=0.0, decay_end_value: float=0.0,
 ) -> Callable[[int], jnp.array]:
     """Returns a linear warmup, linear_decay learning rate function."""
     warmup_fn = optax.linear_schedule(
-        init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
+        init_value=warmup_init_value, end_value=learning_rate, transition_steps=num_warmup_steps)
     decay_fn = optax.linear_schedule(
-        init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
+        init_value=learning_rate, end_value=decay_end_value, transition_steps=num_train_steps - num_warmup_steps
     )
     schedule_fn = optax.join_schedules(
         schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
@@ -662,7 +661,7 @@ def main():
         vectorized_datasets = raw_datasets.map(
             prepare_dataset,
             remove_columns=raw_datasets_features,
-        ).with_format("torch")
+        )
 
     # filter training data with inputs longer than max_input_length
     def is_audio_in_length_range(length):
@@ -681,14 +680,14 @@ def main():
         )
 
     # 8. Load Metric
-    metric = evaluate.load("wer")
+    metric_wer = evaluate.load("wer")
+    metric_cer = evaluate.load("cer")
     do_normalize_eval = data_args.do_normalize_eval
 
     def compute_metrics(pred_ids, label_ids):
         # replace padded labels by the padding token
         for idx in range(len(label_ids)):
             label_ids[idx][label_ids[idx] == -100] = tokenizer.pad_token_id
-        #label_ids[label_ids == -100] = tokenizer.pad_token_id
 
         pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
         # we do not want to group tokens when computing the metrics
@@ -703,9 +702,10 @@ def main():
             label_str = [label_str[i]
                          for i in range(len(label_str)) if len(label_str[i]) > 0]
 
-        wer = 100 * metric.compute(predictions=pred_str, references=label_str)
+        wer = 100 * metric_wer.compute(predictions=pred_str, references=label_str)
+        cer = 100 * metric_cer.compute(predictions=pred_str, references=label_str)
 
-        return {"wer": wer}
+        return {"wer": wer, "cer": cer}
 
     # 9. Save feature extractor, tokenizer and config
     feature_extractor.save_pretrained(training_args.output_dir)
@@ -746,6 +746,7 @@ def main():
     # Initialize our training
     rng = jax.random.PRNGKey(training_args.seed)
     rng, dropout_rng = jax.random.split(rng)
+    # rng, input_rng = jax.random.split(rng)
 
     # Store some constant
     #num_epochs = int(training_args.num_train_epochs)
@@ -755,10 +756,26 @@ def main():
         training_args.per_device_eval_batch_size) * jax.device_count()
 
     # Create learning rate schedule
+    lr_scheduler_types = {"linear", "constant", "constant_with_warmup"}
+    if training_args.lr_scheduler_type not in lr_scheduler_types:
+        raise ValueError(
+            f"lr_scheduler_type of type {training_args.lr_scheduler_type} not supported, choose from {lr_scheduler_types}."
+        )
+    elif training_args.lr_scheduler_type == "constant":
+        warmup_init_value = training_args.learning_rate
+        decay_end_value = training_args.learning_rate
+    elif training_args.lr_scheduler_type == "constant_with_warmup":
+        warmup_init_value = 0.0
+        decay_end_value = training_args.learning_rate
+    else:
+        warmup_init_value = 0.0
+        decay_end_value = 0.0
     linear_decay_lr_schedule_fn = create_learning_rate_fn(
-        data_args.num_train_steps*train_batch_size,
+        data_args.num_train_steps * train_batch_size,
         training_args.warmup_steps,
         training_args.learning_rate,
+        warmup_init_value=warmup_init_value,
+        decay_end_value=decay_end_value,
     )
 
     # We use Optax's "masking" functionality to not apply weight decay
@@ -877,7 +894,7 @@ def main():
 
     # Create parallel version of the train and eval step
     p_train_step = jax.pmap(
-        partial(train_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch", donate_argnums=(0,)
+        partial(train_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch", donate_argnums=(0, )
     )
     p_eval_step = jax.pmap(partial(
         eval_step, label_smoothing_factor=training_args.label_smoothing_factor), "batch")
@@ -888,7 +905,7 @@ def main():
 
     logger.info("***** Running training *****")
     logger.info(
-        f"  Num examples = {data_args.num_train_steps*train_batch_size}")
+        f"  Num examples = {data_args.num_train_steps * train_batch_size}")
     logger.info(
         f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
     logger.info(
@@ -900,44 +917,25 @@ def main():
     # ======================== Training ================================
     train_start = time.time()
 
-    # Create sampling rng
-    #rng, input_rng = jax.random.split(rng)
     train_metrics = []
     epoch = 0
-
-    def collate_batch(samples):
-        return {key: [feature[key] for feature in samples] for key in samples[0].keys()}
-
-    # Create a batched data iterator
-    num_workers = 0
-    # This is not working
-    # vectorized_datasets["train"] = vectorized_datasets["train"].shuffle()
-    train_data_loader = torch.utils.data.DataLoader(
-        batch_size=train_batch_size, dataset=vectorized_datasets["train"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=True)
-    train_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(
-        train_data_loader)
-
+    train_dataset = vectorized_datasets["train"].shuffle(seed=training_args.seed)
+    eval_dataset = vectorized_datasets["eval"]
+    train_loader = data_loader(train_dataset, train_batch_size)
     # train
     for step in tqdm(range(data_args.num_train_steps), desc="Training...", position=1, leave=False):
-
         try:
-            samples = next(train_data_iterator)
-
+            samples = next(train_loader)
         except StopIteration:
             epoch += 1
-            train_data_loader = torch.utils.data.DataLoader(
-                batch_size=train_batch_size, dataset=vectorized_datasets["train"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=True)
-            train_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(
-                train_data_loader)
-            samples = next(train_data_iterator)
-
+            train_dataset.set_epoch(epoch)
+            train_loader = data_loader(train_dataset, train_batch_size)
+            samples = next(train_loader)
             logger.info(
                 f"Completed epoch ({epoch} | Loss: {train_metric['loss']}, Learning Rate:"
                 f" {train_metric['learning_rate']})"
             )
 
-        # reshaped_samples = {key: [feature[key] for feature in samples] for key in samples[0].keys()}
-        # breakpoint()
         batch = data_collator(samples)
         batch = shard(batch.data)
         state, train_metric = p_train_step(state, batch)
@@ -950,16 +948,17 @@ def main():
             eval_metrics = []
             eval_preds = []
             eval_labels = []
-
-            #eval_loader = data_loader(input_rng, vectorized_datasets["eval"], eval_batch_size, drop_last=False)
-            eval_data_loader = torch.utils.data.DataLoader(
-                batch_size=eval_batch_size, dataset=vectorized_datasets["eval"],  num_workers=num_workers, collate_fn=collate_batch, drop_last=False)
-            eval_data_iterator = torch.utils.data.dataloader._SingleProcessDataLoaderIter(
-                eval_data_loader)
-
-            for _ in tqdm(range(training_args.eval_steps), desc="Evaluating...", position=2, leave=False):
+            eval_loader = data_loader(eval_dataset, eval_batch_size, drop_last=False)
+            if data_args.max_eval_samples:
+                max_eval_steps_iter = range(1 + data_args.max_eval_samples // eval_batch_size)
+            else:
+                max_eval_steps_iter = itertools.repeat(None)
+            for _ in tqdm(max_eval_steps_iter, desc="Evaluating...", position=2, leave=False):
                 # Model forward
-                samples = next(eval_data_iterator)
+                try:
+                    samples = next(eval_loader)
+                except StopIteration:
+                    break
                 batch = data_collator(samples)
                 labels = batch["labels"]
 
@@ -968,29 +967,28 @@ def main():
                 )
                 eval_metrics.append(metrics)
 
-            # generation
-            if training_args.predict_with_generate:
-                generated_ids = pad_shard_unpad(
-                    p_generate_step)(state.params, batch.data)
-                eval_preds.extend(jax.device_get(
-                    generated_ids.reshape(-1, gen_kwargs["max_length"])))
-                eval_labels.extend(labels)
-           
-        
+                # generation
+                if training_args.predict_with_generate:
+                    generated_ids = pad_shard_unpad(
+                        p_generate_step)(state.params, batch.data)
+                    eval_preds.extend(jax.device_get(
+                        generated_ids.reshape(-1, gen_kwargs["max_length"])))
+                    eval_labels.extend(labels)
+
             # normalize eval metrics
             eval_metrics = get_metrics(eval_metrics)
             eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
 
-            # compute WER metric
-            wer_desc = ""
+            # compute metrics
+            metric_desc = ""
             if training_args.predict_with_generate:
-                wer_metric = compute_metrics(eval_preds, eval_labels)
-                eval_metrics.update(wer_metric)
-                wer_desc = " ".join(
-                    [f"Eval {key}: {value} |" for key, value in wer_metric.items()])
+                metric_values = compute_metrics(eval_preds, eval_labels)
+                eval_metrics.update(metric_values)
+                metric_desc = " ".join(
+                    [f"Eval {key}: {value} |" for key, value in metric_values.items()])
 
             # Print metrics
-            desc = f"Epoch... ({epoch} | Eval Loss: {eval_metrics['loss']} | {wer_desc})"
+            desc = f"Epoch... ({epoch} | Eval Loss: {eval_metrics['loss']} | {metric_desc})"
             logger.info(desc)
 
             # Save metrics

run_streaming_v2.sh → run_streaming_v3.sh

@@ -1,4 +1,4 @@
-python run_flax_speech_recognition_seq2seq_streaming_v2.py \
+python run_flax_speech_recognition_seq2seq_streaming_v3.py \
             --model_name_or_path openai/whisper-tiny.en \
             --dataset_name mozilla-foundation/common_voice_11_0 \
             --dataset_config nn-NO \
@@ -7,11 +7,11 @@ python run_flax_speech_recognition_seq2seq_streaming_v2.py \
             --eval_split_name test\
             --output_dir whisper-tiny-ft-dummy\
             --overwrite_output_dir\
-            --num_train_epochs=1\
+            --max_eval_samples=16\
             --warmup_steps=8 \
             --do_train \
             --do_eval \
-            --num_train_steps 1000 \
+            --num_train_steps 30 \
             --eval_steps 10 \
             --learning_rate=2e-4 \
             --per_device_train_batch_size=2 \