Add test to check reported training loss (#15096)

* Add test

* Add tests for the reported train loss
This commit is contained in:
Sylvain Gugger 2022-01-11 03:14:11 -05:00 committed by GitHub
parent b67fd797be
commit 9dc8fb2fc7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -444,6 +444,27 @@ class TrainerIntegrationPrerunTest(TestCasePlus, TrainerIntegrationCommon):
trainer.train()
self.check_trained_model(trainer.model)
def test_training_loss(self):
n_gpus = max(1, get_gpu_count())
# With even logs
trainer = get_regression_trainer(logging_steps=64 / (8 * n_gpus))
trainer.train()
log_history = trainer.state.log_history
losses = [log["loss"] for log in log_history if "loss" in log]
train_loss = log_history[-1]["train_loss"]
self.assertAlmostEqual(sum(losses) / len(losses), train_loss, places=4)
# With uneven logs
trainer = get_regression_trainer(logging_steps=5)
trainer.train()
log_history = trainer.state.log_history
# Training loss should be the same as before
new_train_loss = log_history[-1]["train_loss"]
self.assertAlmostEqual(train_loss, new_train_loss, places=4)
def test_custom_optimizer(self):
train_dataset = RegressionDataset()
args = TrainingArguments("./regression")