clean up

run_image_captioning_flax.py

@@ -123,11 +123,9 @@ class TrainingArguments:
     label_smoothing_factor: float = field(
         default=0.0, metadata={"help": "The label smoothing epsilon to apply (zero means no label smoothing)."}
     )
-    adafactor: bool = field(default=False, metadata={"help": "Whether or not to replace AdamW by Adafactor."})
     num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."})
     warmup_steps: int = field(default=0, metadata={"help": "Linear warmup over warmup_steps."})
     logging_steps: int = field(default=500, metadata={"help": "Log every X updates steps."})
-    save_steps: int = field(default=500, metadata={"help": "Save checkpoint every X updates steps."})
     eval_steps: int = field(default=None, metadata={"help": "Run an evaluation every X steps."})
     dataloader_drop_last: bool = field(
         default=False, metadata={"help": "Drop the last incomplete batch if it is not divisible by the batch size."}
@@ -1155,24 +1153,24 @@ 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 run_eval_or_test(rng, dataset, name, is_inside_training=True):
+    def evaluation_loop(rng, dataset, split):
 
-        if name not in ["valid", "test"]:
-            raise ValueError(f"`name` must be either \"valid\" or \"test\". Got {name} instead.")
+        if split not in ["valid", "test"]:
+            raise ValueError(f"`name` must be either \"valid\" or \"test\". Got {split} instead.")
 
-        logger.info(f"*** {'Predict' if name == 'test' else 'Evaluate'} ***")
+        logger.info(f"*** {'Predict' if split == 'test' else 'Evaluate'} ***")
 
         metrics = []
         preds = []
         labels = []
 
-        batches = get_batch_iter(rng, dataset, block_size=block_size, batch_size=eval_batch_size, keep_in_memory=False, shuffle=False, split=name)
+        batches = get_batch_iter(rng, dataset, block_size=block_size, batch_size=eval_batch_size, keep_in_memory=False, shuffle=False, split=split)
         steps = len(dataset) // eval_batch_size + int(len(dataset) % eval_batch_size > 0)
-        for _ in tqdm(range(steps), desc=f"{'Predicting' if name == 'test' else 'Evaluating'}...", position=2, leave=False):
+        for _ in tqdm(range(steps), desc=f"{'Predicting' if split == 'test' else 'Evaluating'}...", position=2, leave=False):
             # Model forward
             batch = next(batches)
             _labels = batch.get("labels", None)
-            if name == "valid" and _labels is None:
+            if split == "valid" and _labels is None:
                 raise ValueError("Validation dataset requires `labels`")
 
             if _labels is not None:
@@ -1198,7 +1196,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 name == 'test' else 'Eval'} {key}: {value} |" for key, value in rouge_metrics.items()])
+                rouge_desc = " ".join([f"{'Predict' if split == 'test' 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", " ")
@@ -1215,8 +1213,8 @@ def main():
 
         if metrics:
             # Print metrics and update progress bar
-            desc = f"{'Predict' if name == 'test' else 'Eval'} Loss: {metrics['loss']} | {rouge_desc})"
-            if is_inside_training:
+            desc = f"{'Predict' if split == 'test' else 'Eval'} Loss: {metrics['loss']} | {rouge_desc})"
+            if split == "valid":
                 desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | " + desc
                 epochs.write(desc)
                 epochs.desc = desc
@@ -1225,7 +1223,7 @@ def main():
         if jax.process_index() == 0:
 
             ckpt_dir = ""
-            if is_inside_training:
+            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)
@@ -1233,9 +1231,9 @@ def main():
             if metrics:
 
                 # save final metrics in json
-                metrics = {f"{name}_{metric_name}": round(value.item(), 6) for metric_name, value in metrics.items()}
-                path = os.path.join(training_args.output_dir, ckpt_dir, f"{name}_results.json")
-                with open(path, "w") as f:
+                metrics = {f"{split}_{metric_name}": round(value.item(), 6) for metric_name, value in metrics.items()}
+                _path = os.path.join(training_args.output_dir, ckpt_dir, f"{split}_results.json")
+                with open(_path, "w") as f:
                     json.dump(metrics, f, indent=4, sort_keys=True)
 
                 # Update report
@@ -1243,14 +1241,20 @@ def main():
                     fp.write(desc + '\n')
 
                 # Save metrics
-                if has_tensorboard and is_inside_training:
-                    write_metric(summary_writer, name, metrics, cur_step)
+                if has_tensorboard and split == "valid":
+                    write_metric(summary_writer, split, metrics, cur_step)
 
             # Save generations
             if generations:
-                with open(os.path.join(training_args.output_dir, ckpt_dir, f'generation_{name}.json'), 'w', encoding='UTF-8') as fp:
+                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, dataset):
+        evaluation_loop(rng, dataset, split='eval')
+
+    def predict(rng, dataset):
+        evaluation_loop(rng, dataset, split='test')
+
     input_rng = None
 
     if training_args.do_train:
@@ -1280,15 +1284,17 @@ def main():
                 train_metrics.append(train_metric)
                 train_time += time.time() - batch_start
 
-                if cur_step % training_args.logging_steps == 0 or (training_args.eval_steps is not None and cur_step % training_args.eval_steps == 0) or cur_step % steps_per_epoch == 0:
+                time_per_step = train_time / cur_step
+                _train_metric = unreplicate(train_metric)
+                desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | Loss: {_train_metric['loss']} | Learning Rate: {_train_metric['learning_rate']} | Time per step: {time_per_step})"
+                epochs.desc = desc
+                epochs.write(desc)
 
-                    time_per_step = train_time / cur_step
+                # log and save info
+                if training_args.logging_steps > 0 and cur_step % training_args.logging_steps == 0:
 
-                    _train_metric = unreplicate(train_metric)
-                    desc = f"Epoch... ({epoch + 1}/{num_epochs} | Step: {cur_step} | Loss: {_train_metric['loss']} | Learning Rate: {_train_metric['learning_rate']} | Time per step: {time_per_step})"
-                    epochs.desc = desc
-                    epochs.write(desc)
                     logger.info(desc)
+
                     with open(os.path.join(training_args.output_dir, 'report.txt'), 'a', encoding='UTF-8') as fp:
                         fp.write(desc + '\n')
 
@@ -1296,34 +1302,38 @@ def main():
                     if has_tensorboard and jax.process_index() == 0:
                         write_metric(summary_writer, "train", train_metrics, cur_step, train_time=train_time)
 
-                    # ======================== Evaluating ==============================
+                # ======================== Evaluating ==============================
+                if 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)
+                    save_results(epoch + 1, cur_step)
 
-                    if training_args.do_eval and ((training_args.eval_steps is not None and cur_step % training_args.eval_steps == 0) or cur_step % steps_per_epoch == 0):
-                        run_eval_or_test(input_rng, eval_dataset, name="valid", is_inside_training=True)
+            # ======================== Epoch End ==============================
 
-                        # ======================== Prediction loop ==============================
+            # log and save info
+            if training_args.logging_steps <= 0:
 
-                        # run prediction after evaluation if specified, otherwise only after each epoch
-                        if training_args.do_predict and training_args.do_predict_during_training and training_args.do_predict_after_evaluation:
-                            run_eval_or_test(input_rng, predict_dataset, name='test', is_inside_training=True)
+                logger.info(desc)
 
-                # ======================== Save ==============================
+                with open(os.path.join(training_args.output_dir, 'report.txt'), 'a', encoding='UTF-8') as fp:
+                    fp.write(desc + '\n')
 
-                if cur_step % training_args.save_steps == 0:
-                    save_results(epoch + 1, cur_step)
+                # Save metrics
+                if has_tensorboard and jax.process_index() == 0:
+                    write_metric(summary_writer, "train", train_metrics, cur_step, train_time=train_time)
 
-            # run prediction after each epoch (if not done during training)
-            if training_args.do_predict and training_args.do_predict_during_training and not training_args.do_predict_after_evaluation:
-                run_eval_or_test(input_rng, predict_dataset, name='test', is_inside_training=True)
+            if training_args.eval_steps is None or training_args.eval_steps <= 0:
+                evaluate(input_rng, eval_dataset)
                 save_results(epoch + 1, cur_step)
 
+    # ======================== Prediction loop ==============================
+
     # 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)
-    if training_args.do_predict and not (training_args.do_train and training_args.do_predict_during_training):
-        run_eval_or_test(input_rng, predict_dataset, name='test', is_inside_training=False)
+    if training_args.do_predict:
+        predict(input_rng, predict_dataset)
 
 
 if __name__ == "__main__":