update

run_image_captioning_flax.py

@@ -112,7 +112,15 @@ class TrainingArguments:
     per_device_eval_batch_size: int = field(
         default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for evaluation."}
     )
-    block_size: int = field(default=None, metadata={"help": "???"})
+    _block_size_doc = \
+        """
+        Split a dataset into chunks of size `block_size`. On each block, images are transformed by the feature extractor
+        and are kept in memory, and the batches of size `batch_size` are yield before processing the next block.
+        """
+    block_size: int = field(
+        default=64,
+        metadata={"help": _block_size_doc}
+    )
     learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."})
     weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
     adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"})
@@ -351,7 +359,7 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
     Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
     Shuffle batches if `shuffle` is `True`.
     """
-    steps_per_epoch = len(dataset) // batch_size
+    steps = len(dataset) // batch_size
 
     if shuffle:
         batch_idx = jax.random.permutation(rng, len(dataset))
@@ -359,7 +367,7 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
     else:
         batch_idx = np.arange(len(dataset))
 
-    for idx in range(steps_per_epoch):
+    for idx in range(steps):
 
         start_idx = batch_size * idx
         end_idx = batch_size * (idx + 1)
@@ -661,37 +669,31 @@ def main():
                 "You can do it from another script, save it, and load it from here, using --feature_extractor_name."
             )
 
-    def get_tokenizer():
-
-        tokenizer = None
-        if model_args.tokenizer_name:
+    tokenizer = None
+    if model_args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+        )
+    elif model_args.model_name_or_path:
+        try:
             tokenizer = AutoTokenizer.from_pretrained(
-                model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+                model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
             )
-        elif model_args.model_name_or_path:
-            try:
-                tokenizer = AutoTokenizer.from_pretrained(
-                    model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
-                )
-            except ValueError as e:
-                logger.warning(e)
-
-        # Check decoder
-        if not tokenizer:
-            if model_args.decoder_model_name_or_path:
-                tokenizer = AutoTokenizer.from_pretrained(
-                    model_args.decoder_model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
-                )
-            else:
-                raise ValueError(
-                    "You are instantiating a new tokenizer from scratch. This is not supported by this script."
-                    "You can do it from another script, save it, and load it from here, using --tokenizer_name."
-                )
-        tokenizer.pad_token = tokenizer.convert_ids_to_tokens(config.pad_token_id)
-
-        return tokenizer
+        except ValueError as e:
+            logger.warning(e)
 
-    tokenizer = get_tokenizer()
+    # Check decoder
+    if not tokenizer:
+        if model_args.decoder_model_name_or_path:
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_args.decoder_model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+            )
+        else:
+            raise ValueError(
+                "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+                "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+            )
+    tokenizer.pad_token = tokenizer.convert_ids_to_tokens(config.pad_token_id)
 
     # Preprocessing the datasets.
     # We need to tokenize inputs and targets.
@@ -864,8 +866,6 @@ def main():
             fn_kwargs={"max_target_length": data_args.val_max_target_length},
         )
 
-    tokenizer = get_tokenizer()
-
     # Split the dataset into several chunks - each chunk is processed (.map) without cache to create a
     # data loader separately (in a sequential order).
     block_size = training_args.block_size
@@ -892,18 +892,26 @@ def main():
         num_test_examples = len(predict_dataset)
         test_steps = num_test_examples // eval_batch_size
 
-    def get_batch_iter(rng, ds, block_size, batch_size, shuffle=False, keep_in_memory=False, split=""):
+    def get_batch_iter(
+            rng: jax.random.PRNGKey,
+            ds: Dataset,
+            block_size: int,
+            batch_size: int,
+            shuffle: bool = False,
+            keep_in_memory: bool = False,
+            split: str = ""
+    ):
 
         if not block_size:
             block_size = len(ds)
 
-        steps_per_split = block_size // batch_size
+        steps_per_block = block_size // batch_size
         num_examples = len(ds)
         steps = num_examples // batch_size
-        num_splits = steps // steps_per_split + int(steps % steps_per_split > 0)
+        num_splits = steps // steps_per_block + int(steps % steps_per_block > 0)
 
         if shuffle:
-            indices = jax.random.permutation(rng, len(train_dataset))
+            indices = jax.random.permutation(rng, len(ds))
             indices = np.asarray(indices)
         else:
             indices = np.arange(len(ds))
@@ -1131,7 +1139,7 @@ def main():
     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, step):
+    def save_results(epoch: int, step: int):
 
         # save checkpoint after each epoch and push checkpoint to the hub
         if jax.process_index() == 0:
@@ -1143,7 +1151,7 @@ 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, dataset, split):
+    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.")
@@ -1239,10 +1247,10 @@ 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, dataset):
+    def evaluate(rng: jax.random.PRNGKey, dataset: Dataset):
         evaluation_loop(rng, dataset, split='valid')
 
-    def predict(rng, dataset):
+    def predict(rng: jax.random.PRNGKey, dataset: Dataset):
         evaluation_loop(rng, dataset, split='test')
 
     input_rng = None
@@ -1292,7 +1300,8 @@ 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 (inside 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)
                     save_results(epoch, cur_step)
@@ -1311,6 +1320,8 @@ def main():
                 if has_tensorboard and jax.process_index() == 0:
                     write_metric(summary_writer, "train", train_metrics, cur_step, train_time=train_time)
 
+            # ======================== 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)
                 save_results(epoch, cur_step)