update

run_image_captioning_flax.py

@@ -359,13 +359,13 @@ def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuf
     else:
         batch_idx = np.arange(len(dataset))
 
-    batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
-    batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
+    for idx in range(steps_per_epoch):
 
-    for idx in batch_idx:
-        batch = dataset[idx]
-        batch = {k: jnp.array(v) for k, v in batch.items()}
+        start_idx = batch_size * idx
+        end_idx = batch_size * (idx + 1)
 
+        selected_indices = batch_idx[start_idx:end_idx]
+        batch = dataset[selected_indices]
         batch = shard(batch)
 
         yield batch
@@ -886,20 +886,20 @@ def main():
 
     if training_args.do_eval:
         num_eval_examples = len(eval_dataset)
-        eval_steps = num_eval_examples // eval_batch_size + int(num_eval_examples % eval_batch_size > 0)
+        eval_steps = num_eval_examples // eval_batch_size
 
     if training_args.do_predict:
         num_test_examples = len(predict_dataset)
-        test_steps = num_test_examples // eval_batch_size + int(num_test_examples % eval_batch_size > 0)
+        test_steps = num_test_examples // eval_batch_size
 
-    def get_batch_iter(rng, ds, block_size, batch_size, shuffle=False, drop_last_batch=False, keep_in_memory=False, split=""):
+    def get_batch_iter(rng, ds, block_size, batch_size, shuffle=False, keep_in_memory=False, split=""):
 
         if not block_size:
             block_size = len(ds)
 
         steps_per_split = block_size // batch_size
         num_examples = len(ds)
-        steps = num_examples // batch_size + int(num_examples % batch_size > 0 and not drop_last_batch)
+        steps = num_examples // batch_size
         num_splits = steps // steps_per_split + int(steps % steps_per_split > 0)
 
         if shuffle:
@@ -1155,7 +1155,7 @@ def main():
         labels = []
 
         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)
+        steps = len(dataset) // eval_batch_size
         for _ in tqdm(range(steps), desc=f"{'Predicting' if split == 'test' else 'Evaluating'}...", position=2, leave=False):
             # Model forward
             batch = next(batches)
@@ -1262,7 +1262,7 @@ def main():
 
             train_metrics = []
 
-            train_batches = get_batch_iter(input_rng, train_dataset, block_size=block_size, batch_size=train_batch_size, keep_in_memory=True, shuffle=True, drop_last_batch=training_args.dataloader_drop_last, split="train")
+            train_batches = get_batch_iter(input_rng, train_dataset, block_size=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)):