From 8fcbe486e1592321e868f872545c8fd9d359a515 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Thu, 10 Sep 2020 15:52:55 -0700 Subject: [PATCH] these tests require non-multigpu env (#7059) * these tests require non-multigpu env * cleanup * clarify --- src/transformers/testing_utils.py | 14 ++++++++++++++ tests/test_trainer.py | 8 +++++++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index 9944aa9d30..cda9e0f889 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -122,6 +122,20 @@ def require_multigpu(test_case): return test_case +def require_non_multigpu(test_case): + """ + Decorator marking a test that requires 0 or 1 GPU setup (in PyTorch). + """ + if not _torch_available: + return unittest.skip("test requires PyTorch")(test_case) + + import torch + + if torch.cuda.device_count() > 1: + return unittest.skip("test requires 0 or 1 GPU")(test_case) + return test_case + + def require_torch_tpu(test_case): """ Decorator marking a test that requires a TPU (in PyTorch). diff --git a/tests/test_trainer.py b/tests/test_trainer.py index e676159957..cd0a2a938f 100755 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -4,7 +4,7 @@ import datasets import numpy as np from transformers import AutoTokenizer, TrainingArguments, is_torch_available -from transformers.testing_utils import get_tests_dir, require_torch +from transformers.testing_utils import get_tests_dir, require_non_multigpu, require_torch if is_torch_available(): @@ -111,6 +111,7 @@ class TrainerIntegrationTest(unittest.TestCase): self.n_epochs = args.num_train_epochs self.batch_size = args.per_device_train_batch_size + @require_non_multigpu def test_reproducible_training(self): # Checks that training worked, model trained and seed made a reproducible training. trainer = get_regression_trainer(learning_rate=0.1) @@ -122,6 +123,7 @@ class TrainerIntegrationTest(unittest.TestCase): trainer.train() self.check_trained_model(trainer.model, alternate_seed=True) + @require_non_multigpu def test_number_of_steps_in_training(self): # Regular training has n_epochs * len(train_dl) steps trainer = get_regression_trainer(learning_rate=0.1) @@ -138,6 +140,7 @@ class TrainerIntegrationTest(unittest.TestCase): train_output = trainer.train() self.assertEqual(train_output.global_step, 10) + @require_non_multigpu def test_train_and_eval_dataloaders(self): trainer = get_regression_trainer(learning_rate=0.1, per_device_train_batch_size=16) self.assertEqual(trainer.get_train_dataloader().batch_size, 16) @@ -200,6 +203,7 @@ class TrainerIntegrationTest(unittest.TestCase): x = trainer.eval_dataset.x self.assertTrue(np.allclose(preds, 1.5 * x + 2.5)) + @require_non_multigpu def test_trainer_with_datasets(self): np.random.seed(42) x = np.random.normal(size=(64,)).astype(np.float32) @@ -228,6 +232,7 @@ class TrainerIntegrationTest(unittest.TestCase): trainer.train() self.check_trained_model(trainer.model) + @require_non_multigpu def test_custom_optimizer(self): train_dataset = RegressionDataset() args = TrainingArguments("./regression") @@ -241,6 +246,7 @@ class TrainerIntegrationTest(unittest.TestCase): self.assertTrue(torch.abs(trainer.model.b - 2.5656) < 1e-4) self.assertEqual(trainer.optimizer.state_dict()["param_groups"][0]["lr"], 1.0) + @require_non_multigpu def test_model_init(self): train_dataset = RegressionDataset() args = TrainingArguments("./regression", learning_rate=0.1)