diff --git a/src/transformers/data/data_collator.py b/src/transformers/data/data_collator.py index 193545eea3..11b8535096 100644 --- a/src/transformers/data/data_collator.py +++ b/src/transformers/data/data_collator.py @@ -62,7 +62,7 @@ def default_data_collator(features: List[InputDataClass]) -> Dict[str, torch.Ten if isinstance(v, torch.Tensor): batch[k] = torch.stack([f[k] for f in features]) else: - batch[k] = torch.tensor([f[k] for f in features], dtype=torch.long) + batch[k] = torch.tensor([f[k] for f in features]) return batch diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index efc29467f1..271f778b30 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -449,6 +449,7 @@ class Trainer: else: t_total = int(len(train_dataloader) // self.args.gradient_accumulation_steps * self.args.num_train_epochs) num_train_epochs = self.args.num_train_epochs + self.args.max_steps = t_total self.create_optimizer_and_scheduler(num_training_steps=t_total) @@ -530,7 +531,7 @@ class Trainer: logging_loss = 0.0 model.zero_grad() train_iterator = trange( - epochs_trained, int(num_train_epochs), desc="Epoch", disable=not self.is_local_process_zero() + epochs_trained, int(np.ceil(num_train_epochs)), desc="Epoch", disable=not self.is_local_process_zero() ) for epoch in train_iterator: if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler): @@ -626,10 +627,10 @@ class Trainer: torch.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt")) torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt")) - if self.args.max_steps > 0 and self.global_step > self.args.max_steps: + if self.args.max_steps > 0 and self.global_step >= self.args.max_steps: epoch_iterator.close() break - if self.args.max_steps > 0 and self.global_step > self.args.max_steps: + if self.args.max_steps > 0 and self.global_step >= self.args.max_steps: train_iterator.close() break if self.args.tpu_metrics_debug or self.args.debug: @@ -986,10 +987,13 @@ class Trainer: if self.args.past_index >= 0: self._past = None + samples_count = 0 for inputs in tqdm(dataloader, desc=description): loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only) + batch_size = inputs[list(inputs.keys())[0]].shape[0] + samples_count += batch_size if loss is not None: - eval_losses.append(loss) + eval_losses.append(loss * batch_size) if logits is not None: preds = logits if preds is None else torch.cat((preds, logits), dim=0) if labels is not None: @@ -1023,7 +1027,7 @@ class Trainer: else: metrics = {} if len(eval_losses) > 0: - metrics["eval_loss"] = np.mean(eval_losses) + metrics["eval_loss"] = np.sum(eval_losses) / samples_count # Prefix all keys with eval_ for key in list(metrics.keys()): diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py index a9d2ea99b2..649001362c 100644 --- a/src/transformers/training_args.py +++ b/src/transformers/training_args.py @@ -69,7 +69,8 @@ class TrainingArguments: max_grad_norm (:obj:`float`, `optional`, defaults to 1.0): Maximum gradient norm (for gradient clipping). num_train_epochs(:obj:`float`, `optional`, defaults to 3.0): - Total number of training epochs to perform. + Total number of training epochs to perform (if not an integer, will perform the decimal part percents of + the last epoch before stopping training). max_steps (:obj:`int`, `optional`, defaults to -1): If set to a positive number, the total number of training steps to perform. Overrides :obj:`num_train_epochs`. diff --git a/tests/test_data_collator.py b/tests/test_data_collator.py new file mode 100644 index 0000000000..41b3b371b9 --- /dev/null +++ b/tests/test_data_collator.py @@ -0,0 +1,152 @@ +import unittest + +from transformers import AutoTokenizer, is_torch_available +from transformers.testing_utils import require_torch + + +if is_torch_available(): + import torch + + from transformers import ( + DataCollatorForLanguageModeling, + DataCollatorForPermutationLanguageModeling, + GlueDataset, + GlueDataTrainingArguments, + LineByLineTextDataset, + TextDataset, + default_data_collator, + ) + + +PATH_SAMPLE_TEXT = "./tests/fixtures/sample_text.txt" + + +@require_torch +class DataCollatorIntegrationTest(unittest.TestCase): + def test_default_with_dict(self): + features = [{"label": i, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)] + batch = default_data_collator(features) + self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8))))) + self.assertEqual(batch["labels"].dtype, torch.long) + self.assertEqual(batch["inputs"].shape, torch.Size([8, 6])) + + # With label_ids + features = [{"label_ids": [0, 1, 2], "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)] + batch = default_data_collator(features) + self.assertTrue(batch["labels"].equal(torch.tensor([[0, 1, 2]] * 8))) + self.assertEqual(batch["labels"].dtype, torch.long) + self.assertEqual(batch["inputs"].shape, torch.Size([8, 6])) + + # Features can already be tensors + features = [{"label": i, "inputs": torch.randint(10, [10])} for i in range(8)] + batch = default_data_collator(features) + self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8))))) + self.assertEqual(batch["labels"].dtype, torch.long) + self.assertEqual(batch["inputs"].shape, torch.Size([8, 10])) + + # Labels can already be tensors + features = [{"label": torch.tensor(i), "inputs": torch.randint(10, [10])} for i in range(8)] + batch = default_data_collator(features) + self.assertEqual(batch["labels"].dtype, torch.long) + self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8))))) + self.assertEqual(batch["labels"].dtype, torch.long) + self.assertEqual(batch["inputs"].shape, torch.Size([8, 10])) + + def test_default_with_no_labels(self): + features = [{"label": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)] + batch = default_data_collator(features) + self.assertTrue("labels" not in batch) + self.assertEqual(batch["inputs"].shape, torch.Size([8, 6])) + + # With label_ids + features = [{"label_ids": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)] + batch = default_data_collator(features) + self.assertTrue("labels" not in batch) + self.assertEqual(batch["inputs"].shape, torch.Size([8, 6])) + + def test_default_classification(self): + MODEL_ID = "bert-base-cased-finetuned-mrpc" + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) + data_args = GlueDataTrainingArguments( + task_name="mrpc", data_dir="./tests/fixtures/tests_samples/MRPC", overwrite_cache=True + ) + dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev") + data_collator = default_data_collator + batch = data_collator(dataset.features) + self.assertEqual(batch["labels"].dtype, torch.long) + + def test_default_regression(self): + MODEL_ID = "distilroberta-base" + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) + data_args = GlueDataTrainingArguments( + task_name="sts-b", data_dir="./tests/fixtures/tests_samples/STS-B", overwrite_cache=True + ) + dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev") + data_collator = default_data_collator + batch = data_collator(dataset.features) + self.assertEqual(batch["labels"].dtype, torch.float) + + def test_lm_tokenizer_without_padding(self): + tokenizer = AutoTokenizer.from_pretrained("gpt2") + data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False) + # ^ causal lm + + dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512) + examples = [dataset[i] for i in range(len(dataset))] + with self.assertRaises(ValueError): + # Expect error due to padding token missing on gpt2: + data_collator(examples) + + dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True) + examples = [dataset[i] for i in range(len(dataset))] + batch = data_collator(examples) + self.assertIsInstance(batch, dict) + self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512))) + self.assertEqual(batch["labels"].shape, torch.Size((2, 512))) + + def test_lm_tokenizer_with_padding(self): + tokenizer = AutoTokenizer.from_pretrained("distilroberta-base") + data_collator = DataCollatorForLanguageModeling(tokenizer) + # ^ masked lm + + dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512) + examples = [dataset[i] for i in range(len(dataset))] + batch = data_collator(examples) + self.assertIsInstance(batch, dict) + self.assertEqual(batch["input_ids"].shape, torch.Size((31, 107))) + self.assertEqual(batch["labels"].shape, torch.Size((31, 107))) + + dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True) + examples = [dataset[i] for i in range(len(dataset))] + batch = data_collator(examples) + self.assertIsInstance(batch, dict) + self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512))) + self.assertEqual(batch["labels"].shape, torch.Size((2, 512))) + + def test_plm(self): + tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased") + data_collator = DataCollatorForPermutationLanguageModeling(tokenizer) + # ^ permutation lm + + dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512) + examples = [dataset[i] for i in range(len(dataset))] + batch = data_collator(examples) + self.assertIsInstance(batch, dict) + self.assertEqual(batch["input_ids"].shape, torch.Size((31, 112))) + self.assertEqual(batch["perm_mask"].shape, torch.Size((31, 112, 112))) + self.assertEqual(batch["target_mapping"].shape, torch.Size((31, 112, 112))) + self.assertEqual(batch["labels"].shape, torch.Size((31, 112))) + + dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True) + examples = [dataset[i] for i in range(len(dataset))] + batch = data_collator(examples) + self.assertIsInstance(batch, dict) + self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512))) + self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 512, 512))) + self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 512, 512))) + self.assertEqual(batch["labels"].shape, torch.Size((2, 512))) + + example = [torch.randint(5, [5])] + with self.assertRaises(ValueError): + # Expect error due to odd sequence length + data_collator(example) diff --git a/tests/test_trainer.py b/tests/test_trainer.py index f5ad0bbeec..072ae78302 100755 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1,5 +1,7 @@ import unittest +import numpy as np + from transformers import AutoTokenizer, TrainingArguments, is_torch_available from transformers.testing_utils import require_torch @@ -10,149 +12,38 @@ if is_torch_available(): from transformers import ( AutoModelForSequenceClassification, - DataCollatorForLanguageModeling, - DataCollatorForPermutationLanguageModeling, GlueDataset, GlueDataTrainingArguments, LineByLineTextDataset, - TextDataset, Trainer, - default_data_collator, ) PATH_SAMPLE_TEXT = "./tests/fixtures/sample_text.txt" -@require_torch -class DataCollatorIntegrationTest(unittest.TestCase): - def test_default_with_dict(self): - features = [{"label": i, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)] - batch = default_data_collator(features) - self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8))))) - self.assertEqual(batch["labels"].dtype, torch.long) - self.assertEqual(batch["inputs"].shape, torch.Size([8, 6])) +class RegressionDataset: + def __init__(self, a=2, b=3, length=64, seed=42): + np.random.seed(seed) + self.length = length + self.x = np.random.normal(size=(length,)).astype(np.float32) + self.y = a * self.x + b + np.random.normal(scale=0.1, size=(length,)) - # With label_ids - features = [{"label_ids": [0, 1, 2], "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)] - batch = default_data_collator(features) - self.assertTrue(batch["labels"].equal(torch.tensor([[0, 1, 2]] * 8))) - self.assertEqual(batch["labels"].dtype, torch.long) - self.assertEqual(batch["inputs"].shape, torch.Size([8, 6])) + def __len__(self): + return self.length - # Features can already be tensors - features = [{"label": i, "inputs": torch.randint(10, [10])} for i in range(8)] - batch = default_data_collator(features) - self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8))))) - self.assertEqual(batch["labels"].dtype, torch.long) - self.assertEqual(batch["inputs"].shape, torch.Size([8, 10])) + def __getitem__(self, i): + return {"input_x": self.x[i], "label": self.y[i]} - # Labels can already be tensors - features = [{"label": torch.tensor(i), "inputs": torch.randint(10, [10])} for i in range(8)] - batch = default_data_collator(features) - self.assertEqual(batch["labels"].dtype, torch.long) - self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8))))) - self.assertEqual(batch["labels"].dtype, torch.long) - self.assertEqual(batch["inputs"].shape, torch.Size([8, 10])) - def test_default_with_no_labels(self): - features = [{"label": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)] - batch = default_data_collator(features) - self.assertTrue("labels" not in batch) - self.assertEqual(batch["inputs"].shape, torch.Size([8, 6])) +class AlmostAccuracy: + def __init__(self, thresh=0.25): + self.thresh = thresh - # With label_ids - features = [{"label_ids": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)] - batch = default_data_collator(features) - self.assertTrue("labels" not in batch) - self.assertEqual(batch["inputs"].shape, torch.Size([8, 6])) - - def test_default_classification(self): - MODEL_ID = "bert-base-cased-finetuned-mrpc" - tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) - data_args = GlueDataTrainingArguments( - task_name="mrpc", data_dir="./tests/fixtures/tests_samples/MRPC", overwrite_cache=True - ) - dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev") - data_collator = default_data_collator - batch = data_collator(dataset.features) - self.assertEqual(batch["labels"].dtype, torch.long) - - def test_default_regression(self): - MODEL_ID = "distilroberta-base" - tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) - data_args = GlueDataTrainingArguments( - task_name="sts-b", data_dir="./tests/fixtures/tests_samples/STS-B", overwrite_cache=True - ) - dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev") - data_collator = default_data_collator - batch = data_collator(dataset.features) - self.assertEqual(batch["labels"].dtype, torch.float) - - def test_lm_tokenizer_without_padding(self): - tokenizer = AutoTokenizer.from_pretrained("gpt2") - data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False) - # ^ causal lm - - dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512) - examples = [dataset[i] for i in range(len(dataset))] - with self.assertRaises(ValueError): - # Expect error due to padding token missing on gpt2: - data_collator(examples) - - dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True) - examples = [dataset[i] for i in range(len(dataset))] - batch = data_collator(examples) - self.assertIsInstance(batch, dict) - self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512))) - self.assertEqual(batch["labels"].shape, torch.Size((2, 512))) - - def test_lm_tokenizer_with_padding(self): - tokenizer = AutoTokenizer.from_pretrained("distilroberta-base") - data_collator = DataCollatorForLanguageModeling(tokenizer) - # ^ masked lm - - dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512) - examples = [dataset[i] for i in range(len(dataset))] - batch = data_collator(examples) - self.assertIsInstance(batch, dict) - self.assertEqual(batch["input_ids"].shape, torch.Size((31, 107))) - self.assertEqual(batch["labels"].shape, torch.Size((31, 107))) - - dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True) - examples = [dataset[i] for i in range(len(dataset))] - batch = data_collator(examples) - self.assertIsInstance(batch, dict) - self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512))) - self.assertEqual(batch["labels"].shape, torch.Size((2, 512))) - - def test_plm(self): - tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased") - data_collator = DataCollatorForPermutationLanguageModeling(tokenizer) - # ^ permutation lm - - dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512) - examples = [dataset[i] for i in range(len(dataset))] - batch = data_collator(examples) - self.assertIsInstance(batch, dict) - self.assertEqual(batch["input_ids"].shape, torch.Size((31, 112))) - self.assertEqual(batch["perm_mask"].shape, torch.Size((31, 112, 112))) - self.assertEqual(batch["target_mapping"].shape, torch.Size((31, 112, 112))) - self.assertEqual(batch["labels"].shape, torch.Size((31, 112))) - - dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True) - examples = [dataset[i] for i in range(len(dataset))] - batch = data_collator(examples) - self.assertIsInstance(batch, dict) - self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512))) - self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 512, 512))) - self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 512, 512))) - self.assertEqual(batch["labels"].shape, torch.Size((2, 512))) - - example = [torch.randint(5, [5])] - with self.assertRaises(ValueError): - # Expect error due to odd sequence length - data_collator(example) + def __call__(self, eval_pred): + predictions, labels = eval_pred + true = np.abs(predictions - labels) <= self.thresh + return {"accuracy": true.astype(np.float32).mean().item()} if is_torch_available(): @@ -168,9 +59,137 @@ if is_torch_available(): def __iter__(self): return iter(self.parse_file()) + class RegressionModel(torch.nn.Module): + def __init__(self, a=0, b=0): + super().__init__() + self.a = torch.nn.Parameter(torch.tensor(a).float()) + self.b = torch.nn.Parameter(torch.tensor(b).float()) + + def forward(self, input_x=None, labels=None): + y = input_x * self.a + self.b + if labels is None: + return (y,) + loss = torch.nn.functional.mse_loss(y, labels) + return (loss, y) + + def get_regression_trainer(a=0, b=0, train_len=64, eval_len=64, **kwargs): + train_dataset = RegressionDataset(length=train_len) + eval_dataset = RegressionDataset(length=eval_len) + model = RegressionModel(a, b) + compute_metrics = kwargs.pop("compute_metrics", None) + data_collator = kwargs.pop("data_collator", None) + optimizers = kwargs.pop("optimizers", (None, None)) + args = TrainingArguments("./regression", **kwargs) + return Trainer( + model, + args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + compute_metrics=compute_metrics, + optimizers=optimizers, + ) + @require_torch class TrainerIntegrationTest(unittest.TestCase): + def setUp(self): + # Get the default values (in case they change): + args = TrainingArguments(".") + self.n_epochs = args.num_train_epochs + self.batch_size = args.per_device_train_batch_size + + 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) + trainer.train() + self.assertTrue(torch.abs(trainer.model.a - 0.6975) < 1e-4) + self.assertTrue(torch.abs(trainer.model.b - 1.2415) < 1e-4) + + # Checks that a different seed gets different (reproducible) results. + trainer = get_regression_trainer(learning_rate=0.1, seed=314) + trainer.train() + self.assertTrue(torch.abs(trainer.model.a - 1.0171) < 1e-4) + self.assertTrue(torch.abs(trainer.model.b - 1.2494) < 1e-4) + + 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) + train_output = trainer.train() + self.assertEqual(train_output.global_step, self.n_epochs * 64 / self.batch_size) + + # Check passing num_train_epochs works (and a float version too): + trainer = get_regression_trainer(learning_rate=0.1, num_train_epochs=1.5) + train_output = trainer.train() + self.assertEqual(train_output.global_step, int(1.5 * 64 / self.batch_size)) + + # If we pass a max_steps, num_train_epochs is ignored + trainer = get_regression_trainer(learning_rate=0.1, max_steps=10) + train_output = trainer.train() + self.assertEqual(train_output.global_step, 10) + + 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) + trainer = get_regression_trainer(learning_rate=0.1, per_device_eval_batch_size=16) + self.assertEqual(trainer.get_eval_dataloader().batch_size, 16) + + # Check drop_last works + trainer = get_regression_trainer( + train_len=66, eval_len=74, learning_rate=0.1, per_device_train_batch_size=16, per_device_eval_batch_size=32 + ) + self.assertEqual(len(trainer.get_train_dataloader()), 66 // 16 + 1) + self.assertEqual(len(trainer.get_eval_dataloader()), 74 // 32 + 1) + + trainer = get_regression_trainer( + train_len=66, + eval_len=74, + learning_rate=0.1, + per_device_train_batch_size=16, + per_device_eval_batch_size=32, + dataloader_drop_last=True, + ) + self.assertEqual(len(trainer.get_train_dataloader()), 66 // 16) + self.assertEqual(len(trainer.get_eval_dataloader()), 74 // 32) + + # Check passing a new dataset fpr evaluation wors + new_eval_dataset = RegressionDataset(length=128) + self.assertEqual(len(trainer.get_eval_dataloader(new_eval_dataset)), 128 // 32) + + def test_evaluate(self): + trainer = get_regression_trainer(a=1.5, b=2.5, compute_metrics=AlmostAccuracy()) + results = trainer.evaluate() + + x, y = trainer.eval_dataset.x, trainer.eval_dataset.y + pred = 1.5 * x + 2.5 + expected_loss = ((pred - y) ** 2).mean() + self.assertAlmostEqual(results["eval_loss"], expected_loss) + expected_acc = AlmostAccuracy()((pred, y))["accuracy"] + self.assertAlmostEqual(results["eval_accuracy"], expected_acc) + + # With a number of elements not a round multiple of the batch size + trainer = get_regression_trainer(a=1.5, b=2.5, eval_len=66, compute_metrics=AlmostAccuracy()) + results = trainer.evaluate() + + x, y = trainer.eval_dataset.x, trainer.eval_dataset.y + pred = 1.5 * x + 2.5 + expected_loss = ((pred - y) ** 2).mean() + self.assertAlmostEqual(results["eval_loss"], expected_loss) + expected_acc = AlmostAccuracy()((pred, y))["accuracy"] + self.assertAlmostEqual(results["eval_accuracy"], expected_acc) + + def test_predict(self): + trainer = get_regression_trainer(a=1.5, b=2.5) + preds = trainer.predict(trainer.eval_dataset).predictions + x = trainer.eval_dataset.x + self.assertTrue(np.allclose(preds, 1.5 * x + 2.5)) + + # With a number of elements not a round multiple of the batch size + trainer = get_regression_trainer(a=1.5, b=2.5, eval_len=66) + preds = trainer.predict(trainer.eval_dataset).predictions + x = trainer.eval_dataset.x + self.assertTrue(np.allclose(preds, 1.5 * x + 2.5)) + def test_trainer_eval_mrpc(self): MODEL_ID = "bert-base-cased-finetuned-mrpc" tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) diff --git a/tests/test_trainer_distributed.py b/tests/test_trainer_distributed.py index 0a31e2eaad..cdc88f9d57 100644 --- a/tests/test_trainer_distributed.py +++ b/tests/test_trainer_distributed.py @@ -60,7 +60,8 @@ if is_torch_available(): if __name__ == "__main__": parser = HfArgumentParser((TrainingArguments,)) - training_args = parser.parse_args_into_dataclasses(sys.argv + ["--output_dir", "./examples"])[0] + sys.argv += ["--output_dir", "./examples"] + training_args = parser.parse_args_into_dataclasses()[0] logger.warning( "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s",