improve code

run_image_captioning_flax.py

@@ -957,12 +957,6 @@ def main():
 
                 _ds = ds.select(selected_indices)
 
-                names = {
-                    "train": "train",
-                    "valid": "validation",
-                    "test": "prediction",
-                }
-
                 _ds = _ds.map(
                     feature_extraction_fn,
                     batched=True,
@@ -971,7 +965,7 @@ def main():
                     load_from_cache_file=not data_args.overwrite_cache,
                     features=features,
                     keep_in_memory=keep_in_memory,
-                    desc=f"Running feature extraction on {names[split]} dataset".replace("  ", " "),
+                    desc=f"Running feature extraction on {split} dataset".replace("  ", " "),
                 )
                 _ds = _ds.with_format("numpy")
 
@@ -1183,25 +1177,30 @@ def main():
                 commit_msg = f"Saving weights and logs of epoch {epoch + 1}- step {step}"
                 repo.push_to_hub(commit_message=commit_msg, blocking=False)
 
-    def evaluation_loop(rng: jax.random.PRNGKey, dataset: Dataset, split: str):
-
-        if split not in ["valid", "test"]:
-            raise ValueError(f"`name` must be either \"valid\" or \"test\". Got {split} instead.")
+    def evaluation_loop(rng: jax.random.PRNGKey, dataset: Dataset, split: str = "eval", ckpt_dir: str = "", is_prediction=False):
 
-        logger.info(f"*** {'Predict' if split == 'test' else 'Evaluate'} ***")
+        logger.info(f"*** {'Predict' if is_prediction else 'Evaluate'} ***")
 
         metrics = []
         preds = []
         labels = []
 
-        batches = blockwise_data_loader(rng, dataset, block_size=training_args.block_size, batch_size=eval_batch_size, keep_in_memory=False, shuffle=False, split=split)
+        batches = blockwise_data_loader(
+            rng,
+            dataset,
+            block_size=training_args.block_size,
+            batch_size=eval_batch_size,
+            keep_in_memory=False,
+            shuffle=False,
+            split="prediction" if is_prediction else "validation",
+        )
         steps = len(dataset) // eval_batch_size
-        for _ in tqdm(range(steps), desc=f"{'Predicting' if split == 'test' else 'Evaluating'}...", position=2, leave=False):
+        for _ in tqdm(range(steps), desc=f"{'Predicting' if is_prediction else 'Evaluating'}...", position=2, leave=False):
             # Model forward
             batch = next(batches)
             _labels = batch.get("labels", None)
-            if split == "valid" and _labels is None:
-                raise ValueError("Validation dataset requires `labels`")
+            if not is_prediction and _labels is None:
+                raise ValueError("Evaluation requires the validation dataset to have `labels`")
 
             if _labels is not None:
                 _metrics = p_eval_step(state.params, batch)
@@ -1226,7 +1225,7 @@ def main():
             if labels:
                 rouge_metrics, decoded_preds, decoded_labels = compute_metrics(preds, labels)
                 metrics.update(rouge_metrics)
-                rouge_desc = " ".join([f"{'Predict' if split == 'test' else 'Eval'} {key}: {value} |" for key, value in rouge_metrics.items()])
+                rouge_desc = " ".join([f"{'Predict' if is_prediction else 'Eval'} {key}: {value} |" for key, value in rouge_metrics.items()])
                 for pred, label in zip(decoded_preds, decoded_labels):
                     pred = pred.replace("\n", " ")
                     label = label.replace("\n", " ")
@@ -1243,8 +1242,8 @@ def main():
 
         if metrics:
             # Print metrics and update progress bar
-            desc = f"{'Predict' if split == 'test' else 'Eval'} Loss: {metrics['loss']} | {rouge_desc})"
-            if split == "valid":
+            desc = f"{'Predict' if is_prediction else 'Eval'} Loss: {metrics['loss']} | {rouge_desc})"
+            if not is_prediction:
                 desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | " + desc
                 epochs.write(desc)
                 epochs.desc = desc
@@ -1252,11 +1251,8 @@ def main():
 
         if jax.process_index() == 0:
 
-            ckpt_dir = ""
-            if split == "valid":
-                ckpt_dir = f'ckpt_epoch_{epoch + 1}_step_{cur_step}'
-                if not os.path.isdir(os.path.join(training_args.output_dir, ckpt_dir)):
-                    os.makedirs(os.path.join(training_args.output_dir, ckpt_dir), exist_ok=True)
+            if not os.path.isdir(os.path.join(training_args.output_dir, ckpt_dir)):
+                os.makedirs(os.path.join(training_args.output_dir, ckpt_dir), exist_ok=True)
 
             if metrics:
 
@@ -1271,7 +1267,7 @@ def main():
                     fp.write(desc + '\n')
 
                 # Save metrics
-                if has_tensorboard and split == "valid":
+                if has_tensorboard and is_prediction:
                     write_metric(summary_writer, split, metrics, cur_step)
 
             # Save generations
@@ -1279,11 +1275,11 @@ def main():
                 with open(os.path.join(training_args.output_dir, ckpt_dir, f'generation_{split}.json'), 'w', encoding='UTF-8') as fp:
                     json.dump(generations, fp, ensure_ascii=False, indent=4)
 
-    def evaluate(rng: jax.random.PRNGKey, dataset: Dataset):
-        evaluation_loop(rng, dataset, split='valid')
+    def evaluate(rng: jax.random.PRNGKey, dataset: Dataset, ckpt_dir: str = ""):
+        evaluation_loop(rng, dataset, split='eval', ckpt_dir=ckpt_dir)
 
     def predict(rng: jax.random.PRNGKey, dataset: Dataset):
-        evaluation_loop(rng, dataset, split='test')
+        evaluation_loop(rng, dataset, split='test', is_prediction=True)
 
     input_rng = None
 
@@ -1302,7 +1298,15 @@ def main():
 
             train_metrics = []
 
-            train_batches = blockwise_data_loader(input_rng, train_dataset, block_size=training_args.block_size, batch_size=train_batch_size, keep_in_memory=True, shuffle=True, split="train")
+            train_batches = blockwise_data_loader(
+                input_rng,
+                train_dataset,
+                block_size=training_args.block_size,
+                batch_size=train_batch_size,
+                keep_in_memory=True,
+                shuffle=True,
+                split="train"
+            )
 
             # train
             for (batch_idx, _) in enumerate(tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False)):
@@ -1332,10 +1336,10 @@ def main():
                     if has_tensorboard and jax.process_index() == 0:
                         write_metric(summary_writer, "train", train_metrics, cur_step, train_time=train_time)
 
-                # ======================== Evaluating (inside epoch) ==============================
+                # ======================== Evaluating (inside an epoch) ==============================
 
                 if training_args.do_eval and (training_args.eval_steps is not None and training_args.eval_steps > 0) and cur_step % training_args.eval_steps == 0:
-                    evaluate(input_rng, eval_dataset)
+                    evaluate(input_rng, eval_dataset, ckpt_dir=f"ckpt_epoch_{epoch + 1}_step_{cur_step}")
                     save_results(epoch, cur_step)
 
             # ======================== Epoch End ==============================
@@ -1355,16 +1359,20 @@ def main():
             # ======================== Evaluating (after each epoch) ==============================
 
             if training_args.do_eval and (training_args.eval_steps is None or training_args.eval_steps <= 0):
-                evaluate(input_rng, eval_dataset)
+                evaluate(input_rng, eval_dataset, ckpt_dir=f"ckpt_epoch_{epoch + 1}_step_{cur_step}")
                 save_results(epoch, cur_step)
 
-    # ======================== Prediction loop ==============================
+    # ======================== Evaluating | Predicting ==============================
 
     # Create sampling rng
     if input_rng is None:
         rng, input_rng = jax.random.split(rng)
 
-    # run prediction after each epoch (if not done during training)
+    # run evaluation without training
+    if training_args.do_eval and not training_args.do_train:
+        evaluate(input_rng, eval_dataset)
+
+    # run prediction after (or without) training
     if training_args.do_predict:
         predict(input_rng, predict_dataset)