Handling states-steps

src/run.sh

@@ -12,16 +12,19 @@ export VALIDATION_FILE=/to/../dev.csv
 export TEST_FILE=/to/../test.csv
 export TEXT_COLUMN=inputs
 export TARGET_COLUMN=targets
-export MAX_SOURCE_LENGTH=128
+export MAX_SOURCE_LENGTH=256
 export MAX_TARGET_LENGTH=1024
 export SOURCE_PREFIX=ingredients
 
 export PER_DEVICE_TRAIN_BATCH_SIZE=8
 export PER_DEVICE_EVAL_BATCH_SIZE=8
 export GRADIENT_ACCUMULATION_STEPS=2
-export NUM_TRAIN_EPOCHS=3.0
-export LEARNING_RATE=5e-4
+export NUM_TRAIN_EPOCHS=5.0
+export LEARNING_RATE=1e-4
 export WARMUP_STEPS=5000
+export LOGGING_STEPS=500
+export EVAL_STEPS=2500
+export SAVE_STEPS=2500
 
 python run_ed_recipe_nlg.py \
     --output_dir="$OUTPUT_DIR"  \
@@ -42,10 +45,11 @@ python run_ed_recipe_nlg.py \
     --num_train_epochs=$NUM_TRAIN_EPOCHS \
     --learning_rate=$LEARNING_RATE \
     --warmup_steps=$WARMUP_STEPS \
-    --preprocessing_num_workers=4 \
+    --logging_step=$LOGGING_STEPS \
+    --eval_steps=$EVAL_STEPS \
+    --save_steps=$SAVE_STEPS \
     --prediction_debug \
     --do_train \
     --do_eval \
-    --do_predict \
     --overwrite_output_dir \
     --predict_with_generate

src/run_ed_recipe_nlg.py

@@ -258,7 +258,20 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
         yield batch
 
 
-def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
+# def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
+#     summary_writer.scalar("train_time", train_time, step)
+#
+#     train_metrics = get_metrics(train_metrics)
+#     for key, vals in train_metrics.items():
+#         tag = f"train_{key}"
+#         for i, val in enumerate(vals):
+#             summary_writer.scalar(tag, val, step - len(vals) + i + 1)
+#
+#     for metric_name, value in eval_metrics.items():
+#         summary_writer.scalar(f"eval_{metric_name}", value, step)
+#
+
+def write_train_metric(summary_writer, train_metrics, train_time, step):
     summary_writer.scalar("train_time", train_time, step)
 
     train_metrics = get_metrics(train_metrics)
@@ -267,6 +280,8 @@ def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
         for i, val in enumerate(vals):
             summary_writer.scalar(tag, val, step - len(vals) + i + 1)
 
+
+def write_eval_metric(summary_writer, eval_metrics, step):
     for metric_name, value in eval_metrics.items():
         summary_writer.scalar(f"eval_{metric_name}", value, step)
 
@@ -553,7 +568,7 @@ def main():
         result = {}
 
         try:
-            result_blue = bleu.compute(predictions=decoded_preds, references=decoded_labels_wer)
+            result_blue = bleu.compute(predictions=decoded_preds, references=decoded_labels_bleu)
             result_blue = result_blue["score"]
         except Exception as e:
             logger.info(f'Error occurred during bleu {e}')
@@ -734,6 +749,7 @@ def main():
     logger.info(f"  Total optimization steps = {total_train_steps}")
 
     train_time = 0
+    train_metrics = []
     epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
     for epoch in epochs:
         # ======================== Training ================================
@@ -741,115 +757,84 @@ def main():
 
         # Create sampling rng
         rng, input_rng = jax.random.split(rng)
-        train_metrics = []
 
         # Generate an epoch by shuffling sampling indices from the train dataset
         train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
         steps_per_epoch = len(train_dataset) // train_batch_size
         # train
-        for _ in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
+        for step in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
             batch = next(train_loader)
             state, train_metric = p_train_step(state, batch)
             train_metrics.append(train_metric)
 
-        train_time += time.time() - train_start
-
-        train_metric = unreplicate(train_metric)
-
-        epochs.write(
-            f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']})"
-        )
-
-        # ======================== Evaluating ==============================
-        eval_metrics = []
-        eval_preds = []
-        eval_labels = []
-
-        eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
-        eval_steps = len(eval_dataset) // eval_batch_size
-        for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
-            # Model forward
-            batch = next(eval_loader)
-            labels = batch["labels"]
-
-            metrics = p_eval_step(state.params, batch)
-            eval_metrics.append(metrics)
-
-            # generation
-            if data_args.predict_with_generate:
-                generated_ids = p_generate_step(state.params, batch)
-                eval_preds.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"])))
-                eval_labels.extend(jax.device_get(labels.reshape(-1, labels.shape[-1])))
-
-        # normalize eval metrics
-        eval_metrics = get_metrics(eval_metrics)
-        eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
-
-        # compute ROUGE metrics
-        rouge_desc = ""
-        if data_args.predict_with_generate:
-            rouge_metrics = compute_metrics(eval_preds, eval_labels)
-            eval_metrics.update(rouge_metrics)
-            rouge_desc = " ".join([f"Eval {key}: {value} |" for key, value in rouge_metrics.items()])
-
-        # Print metrics and update progress bar
-        desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']} | {rouge_desc})"
-        epochs.write(desc)
-        epochs.desc = desc
-
-        # Save metrics
-        if has_tensorboard and jax.process_index() == 0:
-            cur_step = epoch * (len(train_dataset) // train_batch_size)
-            write_metric(summary_writer, train_metrics, eval_metrics, train_time, cur_step)
-
-    # ======================== Prediction loop ==============================
-    if training_args.do_predict:
-        logger.info("*** Predict ***")
-
-        pred_metrics = []
-        pred_generations = []
-        pred_labels = []
-
-        pred_loader = data_loader(input_rng, predict_dataset, eval_batch_size)
-        pred_steps = len(predict_dataset) // eval_batch_size
-        for _ in tqdm(range(pred_steps), desc="Predicting...", position=2, leave=False):
-            # Model forward
-            batch = next(pred_loader)
-            labels = batch["labels"]
-
-            metrics = p_eval_step(state.params, batch)
-            pred_metrics.append(metrics)
-
-            # generation
-            if data_args.predict_with_generate:
-                generated_ids = p_generate_step(state.params, batch)
-                pred_generations.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"])))
-                pred_labels.extend(jax.device_get(labels.reshape(-1, labels.shape[-1])))
-
-        # normalize prediction metrics
-        pred_metrics = get_metrics(pred_metrics)
-        pred_metrics = jax.tree_map(jnp.mean, pred_metrics)
-
-        # compute ROUGE metrics
-        mix_desc = ""
-        if data_args.predict_with_generate:
-            mix_metrics = compute_metrics(pred_generations, pred_labels)
-            pred_metrics.update(mix_metrics)
-            mix_desc = " ".join([f"Predict {key}: {value} |" for key, value in mix_metrics.items()])
-
-        # Print metrics
-        desc = f"Predict Loss: {pred_metrics['loss']} | {mix_desc})"
-        logger.info(desc)
-
-        # save checkpoint after each epoch and push checkpoint to the hub
-        if jax.process_index() == 0:
-            params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
-            model.save_pretrained(
-                training_args.output_dir,
-                params=params,
-                push_to_hub=training_args.push_to_hub,
-                commit_message=f"Saving weights and logs of epoch {epoch + 1}",
-            )
+            cur_step = epoch * (len(train_dataset) // train_batch_size) + step
+
+            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
+                # Save metrics
+                train_metric = unreplicate(train_metric)
+                train_time += time.time() - train_start
+                if has_tensorboard and jax.process_index() == 0:
+                    write_train_metric(summary_writer, train_metrics, train_time, cur_step)
+
+                epochs.write(
+                    f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
+                )
+
+                train_metrics = []
+
+            if cur_step % training_args.eval_steps == 0 and cur_step > 0 and training_args.do_eval:
+                eval_metrics = []
+                eval_preds = []
+                eval_labels = []
+
+                eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
+                eval_steps = len(eval_dataset) // eval_batch_size
+                for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
+                    # Model forward
+                    batch = next(eval_loader)
+                    labels = batch["labels"]
+
+                    metrics = p_eval_step(state.params, batch)
+                    eval_metrics.append(metrics)
+
+                    # generation
+                    if data_args.predict_with_generate:
+                        generated_ids = p_generate_step(state.params, batch)
+                        eval_preds.extend(jax.device_get(generated_ids.reshape(-1, gen_kwargs["max_length"])))
+                        eval_labels.extend(jax.device_get(labels.reshape(-1, labels.shape[-1])))
+
+                # normalize eval metrics
+                eval_metrics = get_metrics(eval_metrics)
+                eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
+
+                # compute MIX metrics
+                mix_desc = ""
+                if data_args.predict_with_generate:
+                    mix_metrics = compute_metrics(eval_preds, eval_labels)
+                    eval_metrics.update(mix_metrics)
+                    mix_desc = " ".join([f"Eval {key}: {value} |" for key, value in mix_metrics.items()])
+
+                # Print metrics and update progress bar
+                desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']} | {mix_desc})"
+                epochs.write(desc)
+                epochs.desc = desc
+
+                # Save metrics
+                if has_tensorboard and jax.process_index() == 0:
+                    cur_step = epoch * (len(train_dataset) // train_batch_size)
+                    write_eval_metric(summary_writer, eval_metrics, cur_step)
+
+            if cur_step % training_args.save_steps == 0 and cur_step > 0:
+                # save checkpoint after each epoch and push checkpoint to the hub
+                if jax.process_index() == 0:
+                    # params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
+                    params = jax.device_get(unreplicate(state.params))
+                    model.save_pretrained(
+                        training_args.output_dir,
+                        params=params,
+                        push_to_hub=training_args.push_to_hub,
+                        commit_message=f"Saving weights and logs of step {cur_step}",
+                    )
 
 
 if __name__ == "__main__":