improve code

run_image_captioning_flax.py

@@ -1157,23 +1157,20 @@ def main():
         logger.info(f"  Num test examples = {num_test_examples}")
         logger.info(f"  Instantaneous test batch size per device = {training_args.per_device_eval_batch_size}")
         logger.info(f"  Total test batch size (w. parallel & distributed) = {eval_batch_size}")
-        logger.info(f"  Total train batch size (w. parallel & distributed) = {eval_batch_size}")
         logger.info(f"  Test steps = {test_steps}")
 
     # create output directory
     if not os.path.isdir(os.path.join(training_args.output_dir)):
         os.makedirs(os.path.join(training_args.output_dir), exist_ok=True)
 
-    def save_results(epoch: int, step: int):
+    def save_ckpt(ckpt_dir: str, commit_msg: str =""):
 
         # 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))
-            dir_name = f'ckpt_epoch_{epoch + 1}_step_{step}'
-            model.save_pretrained(os.path.join(training_args.output_dir, dir_name), params=params)
-            tokenizer.save_pretrained(os.path.join(training_args.output_dir, dir_name))
+            model.save_pretrained(os.path.join(training_args.output_dir, ckpt_dir), params=params)
+            tokenizer.save_pretrained(os.path.join(training_args.output_dir, ckpt_dir))
             if training_args.push_to_hub:
-                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, metric_key_prefix: str = "eval", ckpt_dir: str = "", is_prediction=False):
@@ -1242,7 +1239,7 @@ def main():
         if metrics:
             # Print metrics and update progress bar
             desc = f"{'Predict' if is_prediction else 'Eval'} Loss: {metrics['loss']} | {rouge_desc})"
-            if training_args.do_train:
+            if training_args.do_train and not is_prediction:
                 desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | " + desc
                 epochs.write(desc)
                 epochs.desc = desc
@@ -1338,8 +1335,10 @@ def main():
                 # ======================== 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, ckpt_dir=f"ckpt_epoch_{epoch + 1}_step_{cur_step}")
-                    save_results(epoch, cur_step)
+                    ckpt_dir = f"ckpt_epoch_{epoch + 1}_step_{cur_step}"
+                    commit_msg = f"Saving weights and logs of epoch {epoch + 1} - step {cur_step}"
+                    evaluate(input_rng, eval_dataset, ckpt_dir)
+                    save_ckpt(ckpt_dir=ckpt_dir, commit_msg=commit_msg)
 
             # ======================== Epoch End ==============================
 
@@ -1358,8 +1357,10 @@ 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, ckpt_dir=f"ckpt_epoch_{epoch + 1}_step_{cur_step}")
-                save_results(epoch, cur_step)
+                ckpt_dir = f"ckpt_epoch_{epoch + 1}_step_{cur_step}"
+                commit_msg = f"Saving weights and logs of epoch {epoch + 1} - step {cur_step}"
+                evaluate(input_rng, eval_dataset, ckpt_dir)
+                save_ckpt(ckpt_dir=ckpt_dir, commit_msg=commit_msg)
 
     # ======================== Evaluating | Predicting ==============================