diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py old mode 100755 new mode 100644 index e456e78cfed8..13ee0a4dca64 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2996,7 +2996,8 @@ def evaluation_loop( # Update containers if losses is not None: - losses = self.gather_function(losses.repeat(batch_size)) + loss_batch_size = observed_batch_size if observed_batch_size is not None else batch_size + losses = self.gather_function(losses.repeat(loss_batch_size)) all_losses.add(losses) if inputs_decode is not None: inputs_decode = self.accelerator.pad_across_processes(inputs_decode, dim=1, pad_index=-100) diff --git a/tests/trainer/test_trainer_data.py b/tests/trainer/test_trainer_data.py index d8e573db7797..f67de6dff1c0 100644 --- a/tests/trainer/test_trainer_data.py +++ b/tests/trainer/test_trainer_data.py @@ -1159,6 +1159,37 @@ def test_evaluation_iterable_dataset(self): expected_acc = AlmostAccuracy()((pred, y))["accuracy"] self.assertAlmostEqual(results["eval_accuracy"], expected_acc) + def test_iterable_evaluation_loss_uses_actual_batch_size(self): + for dispatch in (False, True): + for accumulation in (None, 1): + for batch_size in (1, 3, 8): + with self.subTest(dispatch=dispatch, accumulation=accumulation, batch_size=batch_size): + dataset = SampleIterableDataset(length=7) + trainer = get_regression_trainer( + a=1.5, + b=2.5, + output_dir=self.get_auto_remove_tmp_dir(), + use_cpu=True, + per_device_eval_batch_size=batch_size, + eval_accumulation_steps=accumulation, + accelerator_config={"dispatch_batches": dispatch}, + report_to=[], + ) + x, y = dataset.dataset.x, dataset.dataset.ys[0] + expected_predictions = 1.5 * x + 2.5 + expected_loss = ((expected_predictions - y) ** 2).mean() + map_result = trainer.evaluate(dataset.dataset) + iterable_result = trainer.evaluate(dataset) + prediction = trainer.predict(dataset) + for result, prefix in ( + (map_result, "eval"), + (iterable_result, "eval"), + (prediction.metrics, "test"), + ): + self.assertAlmostEqual(result[f"{prefix}_loss"], expected_loss, places=6) + np.testing.assert_allclose(prediction.predictions, expected_predictions) + np.testing.assert_array_equal(prediction.label_ids, y) + def test_predict_iterable_dataset(self): config = RegressionModelConfig(a=1.5, b=2.5) model = RegressionPreTrainedModel(config)