copied changes from soghomon-b:add-eval-step-limit-31561

src/transformers/trainer.py

@@ -3043,8 +3043,8 @@ def _issue_warnings_after_load(self, load_result):
                 f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}."
             )
 
-    def _evaluate(self, trial, ignore_keys_for_eval, skip_scheduler=False):
-        metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
+    def _evaluate(self, trial, ignore_keys_for_eval, skip_scheduler=False, limit_eval_sample_size=False):
+        metrics = self.evaluate(ignore_keys=ignore_keys_for_eval, limit_eval_sample_size=limit_eval_sample_size)
         self._report_to_hp_search(trial, self.state.global_step, metrics)
 
         # Run delayed LR scheduler now that metrics are populated
@@ -3090,7 +3090,7 @@ def _maybe_log_save_evaluate(self, tr_loss, grad_norm, model, trial, epoch, igno
 
         metrics = None
         if self.control.should_evaluate:
-            metrics = self._evaluate(trial, ignore_keys_for_eval)
+            metrics = self._evaluate(trial, ignore_keys_for_eval, limit_eval_sample_size=True)
             is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial)
 
             if self.args.save_strategy == SaveStrategy.BEST:
@@ -4070,6 +4070,7 @@ def evaluate(
         eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None,
         ignore_keys: Optional[list[str]] = None,
         metric_key_prefix: str = "eval",
+        limit_eval_sample_size: bool = False,
     ) -> dict[str, float]:
         """
         Run evaluation and returns metrics.
@@ -4141,6 +4142,7 @@ def evaluate(
             prediction_loss_only=True if self.compute_metrics is None else None,
             ignore_keys=ignore_keys,
             metric_key_prefix=metric_key_prefix,
+            limit_eval_sample_size=limit_eval_sample_size,
         )
 
         total_batch_size = self.args.eval_batch_size * self.args.world_size
@@ -4240,6 +4242,7 @@ def evaluation_loop(
         prediction_loss_only: Optional[bool] = None,
         ignore_keys: Optional[list[str]] = None,
         metric_key_prefix: str = "eval",
+        limit_eval_sample_size: bool = False,
     ) -> EvalLoopOutput:
         """
         Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.
@@ -4316,8 +4319,17 @@ def evaluation_loop(
         # Will be useful when we have an iterable dataset so don't know its length.
         observed_num_examples = 0
 
+        # setting the maximum number of samples the evaluation loop processes
+        max_eval_samples_checker = limit_eval_sample_size and self.args.max_eval_batches != -1
+        max_eval_samples = self.args.max_eval_batches * self.args.eval_batch_size if max_eval_samples_checker else -1
+
+
         # Main evaluation loop
         for step, inputs in enumerate(dataloader):
+
+            # check to see if maximum amount of eval samples reached
+            if max_eval_samples != -1 and observed_num_examples >= max_eval_samples:
+                break
             # Update the observed num examples
             observed_batch_size = find_batch_size(inputs)
             if observed_batch_size is not None:
@@ -4841,6 +4853,7 @@ def prediction_loop(
         prediction_loss_only: Optional[bool] = None,
         ignore_keys: Optional[list[str]] = None,
         metric_key_prefix: str = "eval",
+        limit_eval_sample_size: bool = False,
     ) -> EvalLoopOutput:
         """
         Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.

src/transformers/training_args.py

@@ -309,6 +309,9 @@ class TrainingArguments:
             If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`.
             For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until
             `max_steps` is reached.
+        max_eval_batches (`int`, *optional*, defaults to -1)
+            The maximum number of evaluation samples batches to be used per epoch during training only.
+           The whole evaluation dataset will be used during training by default
         lr_scheduler_type (`str` or [`SchedulerType`], *optional*, defaults to `"linear"`):
             The scheduler type to use. See the documentation of [`SchedulerType`] for all possible values.
         lr_scheduler_kwargs ('dict', *optional*, defaults to {}):
@@ -918,6 +921,12 @@ class TrainingArguments:
         default=-1,
         metadata={"help": "If > 0: set total number of training steps to perform. Override num_train_epochs."},
     )
+    max_eval_batches: int = field(
+        default=-1,
+        metadata={
+            "help": "If > 0: sets the number of evaluation samples batches used per epoch during training. If not set, evaluation during training will use all evaluation samples."
+        },
+    )
     lr_scheduler_type: Union[SchedulerType, str] = field(
         default="linear",
         metadata={"help": "The scheduler type to use."},

tests/trainer/test_trainer.py

@@ -237,6 +237,17 @@ def __getitem__(self, i):
         return result
 
 
+# A class to test Trainer with an evaluation dataset limit
+class TrainerEvalSampleLimit(Trainer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.observed_num_batches = 0
+
+    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+        self.observed_num_batches += 1
+        return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys)
+
+
 # Converting Bytes to Megabytes
 def bytes2megabytes(x):
     return int(x / 2**20)
@@ -1292,6 +1303,75 @@ def test_evaluation_with_keys_to_drop(self):
         result = trainer.predict(eval_dataset, ignore_keys=[])
         self.assertTrue(isinstance(result.predictions, tuple))
         self.assertEqual(len(result.predictions), 2)
+    def test_number_of_eval_samples_set(self):
+        config = LlamaConfig(vocab_size=100, hidden_size=32, num_hidden_layers=3, num_attention_heads=4)
+        tiny_llama = LlamaForCausalLM(config)
+
+        x = torch.randint(0, 100, (128,))
+        train_dataset = RepeatDataset(x)
+        eval_dataset = RepeatDataset(x)
+
+        args = TrainingArguments(
+            self.get_auto_remove_tmp_dir(),
+            max_eval_batches=4,
+            eval_strategy="epoch",  # Enable evaluation at the end of each epoch
+            num_train_epochs=1,  # leave 1 so that testing works. If > 1, test below will fail. Since it is
+            # prediction batch limit per epoch
+            per_device_train_batch_size=4,  # You'll likely want to set your batch size
+            per_device_eval_batch_size=4,  # You'll likely want to set your evaluation batch size
+        )
+
+        trainer = TrainerEvalSampleLimit(
+            model=tiny_llama,
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+        )
+
+        # when calling train, we will limit the eval sample size if max_eval_batches is defined as we have limit_eval_sample_size=True in _evaluate method
+        trainer.train()
+        self.assertEqual(args.max_eval_batches, trainer.observed_num_batches)
+
+        # reset number of batches seen
+        trainer.observed_num_batches = 0
+
+        # when calling evaluate directly, we won't limit the eval sample size if max_eval_batches is defined as we have limit_eval_sample_size=False
+        trainer.evaluate()
+        self.assertEqual(trainer.observed_num_batches, len(eval_dataset) / args.eval_batch_size)
+
+    def test_number_of_eval_samples_unset(self):
+        config = LlamaConfig(vocab_size=100, hidden_size=32, num_hidden_layers=3, num_attention_heads=4)
+        tiny_llama = LlamaForCausalLM(config)
+        x = torch.randint(0, 100, (128,))
+        train_dataset = RepeatDataset(x)
+        eval_dataset = RepeatDataset(x)
+
+        args = TrainingArguments(
+            self.get_auto_remove_tmp_dir(),
+            eval_strategy="epoch",  # Enable evaluation at the end of each epoch
+            num_train_epochs=1,  # leave 1 so that testing works. If > 1, test below will fail. Since it is
+            # prediction batch limit per epoch
+            per_device_train_batch_size=4,  # You'll likely want to set your batch size
+            per_device_eval_batch_size=4,  # You'll likely want to set your evaluation batch size
+        )
+
+        trainer = TrainerEvalSampleLimit(
+            model=tiny_llama,
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+        )
+
+        # when calling train, we will limit the eval sample size if max_eval_batches is defined as we have limit_eval_sample_size=True in _evaluate method
+        trainer.train()
+        self.assertEqual(trainer.observed_num_batches, len(eval_dataset) / args.eval_batch_size)
+
+        # reset number of batches seen
+        trainer.observed_num_batches = 0
+
+        # when calling evaluate directly, we won't limit the eval sample size if max_eval_batches is defined as we have limit_eval_sample_size=False
+        trainer.evaluate()
+        self.assertEqual(trainer.observed_num_batches, len(eval_dataset) / args.eval_batch_size)
 
     def test_training_arguments_are_left_untouched(self):
         tmp_dir = self.get_auto_remove_tmp_dir()