[Test refactor 1/5] Per-folder tests reorganization (#15725)
* Per-folder tests reorganization Co-authored-by: sgugger <sylvain.gugger@gmail.com> Co-authored-by: Stas Bekman <stas@stason.org>
This commit is contained in:
Lysandre Debut
0
tests/trainer/__init__.py
Normal file
0
tests/trainer/__init__.py
Normal file
864
tests/trainer/test_data_collator.py
Normal file
864
tests/trainer/test_data_collator.py
Normal file
@@ -0,0 +1,864 @@
|
||||
# Copyright 2020 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from transformers import (
|
||||
BertTokenizer,
|
||||
DataCollatorForLanguageModeling,
|
||||
DataCollatorForPermutationLanguageModeling,
|
||||
DataCollatorForTokenClassification,
|
||||
DataCollatorForWholeWordMask,
|
||||
DataCollatorWithPadding,
|
||||
default_data_collator,
|
||||
is_tf_available,
|
||||
is_torch_available,
|
||||
set_seed,
|
||||
)
|
||||
from transformers.testing_utils import require_tf, require_torch
|
||||
|
||||
|
||||
if is_torch_available():
|
||||
import torch
|
||||
|
||||
if is_tf_available():
|
||||
import tensorflow as tf
|
||||
|
||||
|
||||
@require_torch
|
||||
class DataCollatorIntegrationTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdirname = tempfile.mkdtemp()
|
||||
|
||||
vocab_tokens = ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
|
||||
self.vocab_file = os.path.join(self.tmpdirname, "vocab.txt")
|
||||
with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
|
||||
vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmpdirname)
|
||||
|
||||
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": np.random.randint(0, 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": np.random.randint(0, 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_classification_and_regression(self):
|
||||
data_collator = default_data_collator
|
||||
|
||||
features = [{"input_ids": [0, 1, 2, 3, 4], "label": i} for i in range(4)]
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["labels"].dtype, torch.long)
|
||||
|
||||
features = [{"input_ids": [0, 1, 2, 3, 4], "label": float(i)} for i in range(4)]
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["labels"].dtype, torch.float)
|
||||
|
||||
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_data_collator_with_padding(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [{"input_ids": [0, 1, 2]}, {"input_ids": [0, 1, 2, 3, 4, 5]}]
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size([2, 6]))
|
||||
self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer, padding="max_length", max_length=10)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size([2, 10]))
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size([2, 8]))
|
||||
|
||||
def test_data_collator_for_token_classification(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{"input_ids": [0, 1, 2], "labels": [0, 1, 2]},
|
||||
{"input_ids": [0, 1, 2, 3, 4, 5], "labels": [0, 1, 2, 3, 4, 5]},
|
||||
]
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size([2, 6]))
|
||||
self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
self.assertEqual(batch["labels"].shape, torch.Size([2, 6]))
|
||||
self.assertEqual(batch["labels"][0].tolist(), [0, 1, 2] + [-100] * 3)
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, padding="max_length", max_length=10)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size([2, 10]))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size([2, 10]))
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=8)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size([2, 8]))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size([2, 8]))
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, label_pad_token_id=-1)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size([2, 6]))
|
||||
self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
self.assertEqual(batch["labels"].shape, torch.Size([2, 6]))
|
||||
self.assertEqual(batch["labels"][0].tolist(), [0, 1, 2] + [-1] * 3)
|
||||
|
||||
def _test_no_pad_and_pad(self, no_pad_features, pad_features):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False, pad_to_multiple_of=8)
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 16)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 16)))
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 16)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 16)))
|
||||
|
||||
tokenizer._pad_token = None
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
|
||||
with self.assertRaises(ValueError):
|
||||
# Expect error due to padding token missing
|
||||
data_collator(pad_features)
|
||||
|
||||
set_seed(42) # For reproducibility
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer)
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(torch.any(masked_tokens))
|
||||
self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(torch.any(masked_tokens))
|
||||
self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8)
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 16)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 16)))
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(torch.any(masked_tokens))
|
||||
self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 16)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 16)))
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(torch.any(masked_tokens))
|
||||
self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
|
||||
|
||||
def test_data_collator_for_language_modeling(self):
|
||||
no_pad_features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
pad_features = [{"input_ids": list(range(5))}, {"input_ids": list(range(10))}]
|
||||
self._test_no_pad_and_pad(no_pad_features, pad_features)
|
||||
|
||||
no_pad_features = [list(range(10)), list(range(10))]
|
||||
pad_features = [list(range(5)), list(range(10))]
|
||||
self._test_no_pad_and_pad(no_pad_features, pad_features)
|
||||
|
||||
def test_data_collator_for_whole_word_mask(self):
|
||||
features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForWholeWordMask(tokenizer, return_tensors="pt")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
|
||||
|
||||
def test_plm(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
no_pad_features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
pad_features = [{"input_ids": list(range(5))}, {"input_ids": list(range(10))}]
|
||||
|
||||
data_collator = DataCollatorForPermutationLanguageModeling(tokenizer)
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertIsInstance(batch, dict)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
|
||||
self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 10, 10)))
|
||||
self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 10, 10)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
|
||||
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertIsInstance(batch, dict)
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 10)))
|
||||
self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 10, 10)))
|
||||
self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 10, 10)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 10)))
|
||||
|
||||
example = [np.random.randint(0, 5, [5])]
|
||||
with self.assertRaises(ValueError):
|
||||
# Expect error due to odd sequence length
|
||||
data_collator(example)
|
||||
|
||||
def test_nsp(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{"input_ids": [0, 1, 2, 3, 4], "token_type_ids": [0, 1, 2, 3, 4], "next_sentence_label": i}
|
||||
for i in range(2)
|
||||
]
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer)
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 5)))
|
||||
self.assertEqual(batch["token_type_ids"].shape, torch.Size((2, 5)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 5)))
|
||||
self.assertEqual(batch["next_sentence_label"].shape, torch.Size((2,)))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8)
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 8)))
|
||||
self.assertEqual(batch["token_type_ids"].shape, torch.Size((2, 8)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 8)))
|
||||
self.assertEqual(batch["next_sentence_label"].shape, torch.Size((2,)))
|
||||
|
||||
def test_sop(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{
|
||||
"input_ids": torch.tensor([0, 1, 2, 3, 4]),
|
||||
"token_type_ids": torch.tensor([0, 1, 2, 3, 4]),
|
||||
"sentence_order_label": i,
|
||||
}
|
||||
for i in range(2)
|
||||
]
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer)
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 5)))
|
||||
self.assertEqual(batch["token_type_ids"].shape, torch.Size((2, 5)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 5)))
|
||||
self.assertEqual(batch["sentence_order_label"].shape, torch.Size((2,)))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8)
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, torch.Size((2, 8)))
|
||||
self.assertEqual(batch["token_type_ids"].shape, torch.Size((2, 8)))
|
||||
self.assertEqual(batch["labels"].shape, torch.Size((2, 8)))
|
||||
self.assertEqual(batch["sentence_order_label"].shape, torch.Size((2,)))
|
||||
|
||||
|
||||
@require_tf
|
||||
class TFDataCollatorIntegrationTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.tmpdirname = tempfile.mkdtemp()
|
||||
|
||||
vocab_tokens = ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
|
||||
self.vocab_file = os.path.join(self.tmpdirname, "vocab.txt")
|
||||
with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
|
||||
vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmpdirname)
|
||||
|
||||
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, return_tensors="tf")
|
||||
self.assertEqual(batch["labels"].numpy().tolist(), list(range(8)))
|
||||
self.assertEqual(batch["labels"].dtype, tf.int64)
|
||||
self.assertEqual(batch["inputs"].shape.as_list(), [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, return_tensors="tf")
|
||||
self.assertEqual(batch["labels"].numpy().tolist(), ([[0, 1, 2]] * 8))
|
||||
self.assertEqual(batch["labels"].dtype, tf.int64)
|
||||
self.assertEqual(batch["inputs"].shape.as_list(), [8, 6])
|
||||
|
||||
# Features can already be tensors
|
||||
features = [{"label": i, "inputs": np.random.randint(0, 10, [10])} for i in range(8)]
|
||||
batch = default_data_collator(features, return_tensors="tf")
|
||||
self.assertEqual(batch["labels"].numpy().tolist(), (list(range(8))))
|
||||
self.assertEqual(batch["labels"].dtype, tf.int64)
|
||||
self.assertEqual(batch["inputs"].shape.as_list(), [8, 10])
|
||||
|
||||
# Labels can already be tensors
|
||||
features = [{"label": np.array(i), "inputs": np.random.randint(0, 10, [10])} for i in range(8)]
|
||||
batch = default_data_collator(features, return_tensors="tf")
|
||||
self.assertEqual(batch["labels"].dtype, tf.int64)
|
||||
self.assertEqual(batch["labels"].numpy().tolist(), list(range(8)))
|
||||
self.assertEqual(batch["labels"].dtype, tf.int64)
|
||||
self.assertEqual(batch["inputs"].shape.as_list(), [8, 10])
|
||||
|
||||
def test_numpy_dtype_preservation(self):
|
||||
data_collator = default_data_collator
|
||||
|
||||
# Confirms that numpy inputs are handled correctly even when scalars
|
||||
features = [{"input_ids": np.array([0, 1, 2, 3, 4]), "label": np.int64(i)} for i in range(4)]
|
||||
batch = data_collator(features, return_tensors="tf")
|
||||
self.assertEqual(batch["labels"].dtype, tf.int64)
|
||||
|
||||
def test_default_classification_and_regression(self):
|
||||
data_collator = default_data_collator
|
||||
|
||||
features = [{"input_ids": [0, 1, 2, 3, 4], "label": i} for i in range(4)]
|
||||
batch = data_collator(features, return_tensors="tf")
|
||||
self.assertEqual(batch["labels"].dtype, tf.int64)
|
||||
|
||||
features = [{"input_ids": [0, 1, 2, 3, 4], "label": float(i)} for i in range(4)]
|
||||
batch = data_collator(features, return_tensors="tf")
|
||||
self.assertEqual(batch["labels"].dtype, tf.float32)
|
||||
|
||||
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, return_tensors="tf")
|
||||
self.assertTrue("labels" not in batch)
|
||||
self.assertEqual(batch["inputs"].shape.as_list(), [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, return_tensors="tf")
|
||||
self.assertTrue("labels" not in batch)
|
||||
self.assertEqual(batch["inputs"].shape.as_list(), [8, 6])
|
||||
|
||||
def test_data_collator_with_padding(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [{"input_ids": [0, 1, 2]}, {"input_ids": [0, 1, 2, 3, 4, 5]}]
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 6])
|
||||
self.assertEqual(batch["input_ids"][0].numpy().tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer, padding="max_length", max_length=10, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, [2, 8])
|
||||
|
||||
def test_data_collator_for_token_classification(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{"input_ids": [0, 1, 2], "labels": [0, 1, 2]},
|
||||
{"input_ids": [0, 1, 2, 3, 4, 5], "labels": [0, 1, 2, 3, 4, 5]},
|
||||
]
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 6])
|
||||
self.assertEqual(batch["input_ids"][0].numpy().tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 6])
|
||||
self.assertEqual(batch["labels"][0].numpy().tolist(), [0, 1, 2] + [-100] * 3)
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(
|
||||
tokenizer, padding="max_length", max_length=10, return_tensors="tf"
|
||||
)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 10])
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=8, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 8])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 8])
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, label_pad_token_id=-1, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 6])
|
||||
self.assertEqual(batch["input_ids"][0].numpy().tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 6])
|
||||
self.assertEqual(batch["labels"][0].numpy().tolist(), [0, 1, 2] + [-1] * 3)
|
||||
|
||||
def _test_no_pad_and_pad(self, no_pad_features, pad_features):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False, return_tensors="tf")
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 10])
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 10])
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(
|
||||
tokenizer, mlm=False, pad_to_multiple_of=8, return_tensors="tf"
|
||||
)
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 16])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 16])
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 16])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 16])
|
||||
|
||||
tokenizer._pad_token = None
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False, return_tensors="tf")
|
||||
with self.assertRaises(ValueError):
|
||||
# Expect error due to padding token missing
|
||||
data_collator(pad_features)
|
||||
|
||||
set_seed(42) # For reproducibility
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, return_tensors="tf")
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 10])
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(tf.reduce_any(masked_tokens))
|
||||
# self.assertTrue(all(x == -100 for x in batch["labels"].numpy()[~masked_tokens.numpy()].tolist()))
|
||||
|
||||
batch = data_collator(pad_features, return_tensors="tf")
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 10])
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(tf.reduce_any(masked_tokens))
|
||||
# self.assertTrue(all(x == -100 for x in batch["labels"].numpy()[~masked_tokens.numpy()].tolist()))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors="tf")
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 16])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 16])
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(tf.reduce_any(masked_tokens))
|
||||
# self.assertTrue(all(x == -100 for x in batch["labels"].numpy()[~masked_tokens.numpy()].tolist()))
|
||||
|
||||
batch = data_collator(pad_features, return_tensors="tf")
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 16])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 16])
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(tf.reduce_any(masked_tokens))
|
||||
# self.assertTrue(all(x == -100 for x in batch["labels"].numpy()[~masked_tokens.numpy()].tolist()))
|
||||
|
||||
def test_data_collator_for_language_modeling(self):
|
||||
no_pad_features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
pad_features = [{"input_ids": list(range(5))}, {"input_ids": list(range(10))}]
|
||||
self._test_no_pad_and_pad(no_pad_features, pad_features)
|
||||
|
||||
no_pad_features = [list(range(10)), list(range(10))]
|
||||
pad_features = [list(range(5)), list(range(10))]
|
||||
self._test_no_pad_and_pad(no_pad_features, pad_features)
|
||||
|
||||
def test_data_collator_for_whole_word_mask(self):
|
||||
features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForWholeWordMask(tokenizer, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 10])
|
||||
|
||||
def test_plm(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
no_pad_features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
pad_features = [{"input_ids": list(range(5))}, {"input_ids": list(range(10))}]
|
||||
|
||||
data_collator = DataCollatorForPermutationLanguageModeling(tokenizer, return_tensors="tf")
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertIsInstance(batch, dict)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
self.assertEqual(batch["perm_mask"].shape.as_list(), [2, 10, 10])
|
||||
self.assertEqual(batch["target_mapping"].shape.as_list(), [2, 10, 10])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 10])
|
||||
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertIsInstance(batch, dict)
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 10])
|
||||
self.assertEqual(batch["perm_mask"].shape.as_list(), [2, 10, 10])
|
||||
self.assertEqual(batch["target_mapping"].shape.as_list(), [2, 10, 10])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 10])
|
||||
|
||||
example = [np.random.randint(0, 5, [5])]
|
||||
with self.assertRaises(ValueError):
|
||||
# Expect error due to odd sequence length
|
||||
data_collator(example)
|
||||
|
||||
def test_nsp(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{"input_ids": [0, 1, 2, 3, 4], "token_type_ids": [0, 1, 2, 3, 4], "next_sentence_label": i}
|
||||
for i in range(2)
|
||||
]
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 5])
|
||||
self.assertEqual(batch["token_type_ids"].shape.as_list(), [2, 5])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 5])
|
||||
self.assertEqual(batch["next_sentence_label"].shape.as_list(), [2])
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 8])
|
||||
self.assertEqual(batch["token_type_ids"].shape.as_list(), [2, 8])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 8])
|
||||
self.assertEqual(batch["next_sentence_label"].shape.as_list(), [2])
|
||||
|
||||
def test_sop(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{
|
||||
"input_ids": tf.convert_to_tensor([0, 1, 2, 3, 4]),
|
||||
"token_type_ids": tf.convert_to_tensor([0, 1, 2, 3, 4]),
|
||||
"sentence_order_label": i,
|
||||
}
|
||||
for i in range(2)
|
||||
]
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 5])
|
||||
self.assertEqual(batch["token_type_ids"].shape.as_list(), [2, 5])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 5])
|
||||
self.assertEqual(batch["sentence_order_label"].shape.as_list(), [2])
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors="tf")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape.as_list(), [2, 8])
|
||||
self.assertEqual(batch["token_type_ids"].shape.as_list(), [2, 8])
|
||||
self.assertEqual(batch["labels"].shape.as_list(), [2, 8])
|
||||
self.assertEqual(batch["sentence_order_label"].shape.as_list(), [2])
|
||||
|
||||
|
||||
class NumpyDataCollatorIntegrationTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdirname = tempfile.mkdtemp()
|
||||
|
||||
vocab_tokens = ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
|
||||
self.vocab_file = os.path.join(self.tmpdirname, "vocab.txt")
|
||||
with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
|
||||
vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmpdirname)
|
||||
|
||||
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, return_tensors="np")
|
||||
self.assertEqual(batch["labels"].tolist(), list(range(8)))
|
||||
self.assertEqual(batch["labels"].dtype, np.int64)
|
||||
self.assertEqual(batch["inputs"].shape, (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, return_tensors="np")
|
||||
self.assertEqual(batch["labels"].tolist(), [[0, 1, 2]] * 8)
|
||||
self.assertEqual(batch["labels"].dtype, np.int64)
|
||||
self.assertEqual(batch["inputs"].shape, (8, 6))
|
||||
|
||||
# Features can already be tensors
|
||||
features = [{"label": i, "inputs": np.random.randint(0, 10, [10])} for i in range(8)]
|
||||
batch = default_data_collator(features, return_tensors="np")
|
||||
self.assertEqual(batch["labels"].tolist(), list(range(8)))
|
||||
self.assertEqual(batch["labels"].dtype, np.int64)
|
||||
self.assertEqual(batch["inputs"].shape, (8, 10))
|
||||
|
||||
# Labels can already be tensors
|
||||
features = [{"label": np.array(i), "inputs": np.random.randint(0, 10, [10])} for i in range(8)]
|
||||
batch = default_data_collator(features, return_tensors="np")
|
||||
self.assertEqual(batch["labels"].dtype, np.int64)
|
||||
self.assertEqual(batch["labels"].tolist(), (list(range(8))))
|
||||
self.assertEqual(batch["labels"].dtype, np.int64)
|
||||
self.assertEqual(batch["inputs"].shape, (8, 10))
|
||||
|
||||
def test_default_classification_and_regression(self):
|
||||
data_collator = default_data_collator
|
||||
|
||||
features = [{"input_ids": [0, 1, 2, 3, 4], "label": i} for i in range(4)]
|
||||
batch = data_collator(features, return_tensors="np")
|
||||
self.assertEqual(batch["labels"].dtype, np.int64)
|
||||
|
||||
features = [{"input_ids": [0, 1, 2, 3, 4], "label": float(i)} for i in range(4)]
|
||||
batch = data_collator(features, return_tensors="np")
|
||||
self.assertEqual(batch["labels"].dtype, np.float32)
|
||||
|
||||
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, return_tensors="np")
|
||||
self.assertTrue("labels" not in batch)
|
||||
self.assertEqual(batch["inputs"].shape, (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, return_tensors="np")
|
||||
self.assertTrue("labels" not in batch)
|
||||
self.assertEqual(batch["inputs"].shape, (8, 6))
|
||||
|
||||
def test_data_collator_with_padding(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [{"input_ids": [0, 1, 2]}, {"input_ids": [0, 1, 2, 3, 4, 5]}]
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 6))
|
||||
self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer, padding="max_length", max_length=10, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
|
||||
data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 8))
|
||||
|
||||
def test_data_collator_for_token_classification(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{"input_ids": [0, 1, 2], "labels": [0, 1, 2]},
|
||||
{"input_ids": [0, 1, 2, 3, 4, 5], "labels": [0, 1, 2, 3, 4, 5]},
|
||||
]
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 6))
|
||||
self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
self.assertEqual(batch["labels"].shape, (2, 6))
|
||||
self.assertEqual(batch["labels"][0].tolist(), [0, 1, 2] + [-100] * 3)
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(
|
||||
tokenizer, padding="max_length", max_length=10, return_tensors="np"
|
||||
)
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
self.assertEqual(batch["labels"].shape, (2, 10))
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=8, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 8))
|
||||
self.assertEqual(batch["labels"].shape, (2, 8))
|
||||
|
||||
data_collator = DataCollatorForTokenClassification(tokenizer, label_pad_token_id=-1, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 6))
|
||||
self.assertEqual(batch["input_ids"][0].tolist(), [0, 1, 2] + [tokenizer.pad_token_id] * 3)
|
||||
self.assertEqual(batch["labels"].shape, (2, 6))
|
||||
self.assertEqual(batch["labels"][0].tolist(), [0, 1, 2] + [-1] * 3)
|
||||
|
||||
def _test_no_pad_and_pad(self, no_pad_features, pad_features):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False, return_tensors="np")
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
self.assertEqual(batch["labels"].shape, (2, 10))
|
||||
|
||||
batch = data_collator(pad_features, return_tensors="np")
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
self.assertEqual(batch["labels"].shape, (2, 10))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(
|
||||
tokenizer, mlm=False, pad_to_multiple_of=8, return_tensors="np"
|
||||
)
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 16))
|
||||
self.assertEqual(batch["labels"].shape, (2, 16))
|
||||
|
||||
batch = data_collator(pad_features, return_tensors="np")
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 16))
|
||||
self.assertEqual(batch["labels"].shape, (2, 16))
|
||||
|
||||
tokenizer._pad_token = None
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False, return_tensors="np")
|
||||
with self.assertRaises(ValueError):
|
||||
# Expect error due to padding token missing
|
||||
data_collator(pad_features)
|
||||
|
||||
set_seed(42) # For reproducibility
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, return_tensors="np")
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
self.assertEqual(batch["labels"].shape, (2, 10))
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(np.any(masked_tokens))
|
||||
# self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
self.assertEqual(batch["labels"].shape, (2, 10))
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(np.any(masked_tokens))
|
||||
# self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors="np")
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 16))
|
||||
self.assertEqual(batch["labels"].shape, (2, 16))
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(np.any(masked_tokens))
|
||||
# self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 16))
|
||||
self.assertEqual(batch["labels"].shape, (2, 16))
|
||||
|
||||
masked_tokens = batch["input_ids"] == tokenizer.mask_token_id
|
||||
self.assertTrue(np.any(masked_tokens))
|
||||
# self.assertTrue(all(x == -100 for x in batch["labels"][~masked_tokens].tolist()))
|
||||
|
||||
def test_data_collator_for_language_modeling(self):
|
||||
no_pad_features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
pad_features = [{"input_ids": list(range(5))}, {"input_ids": list(range(10))}]
|
||||
self._test_no_pad_and_pad(no_pad_features, pad_features)
|
||||
|
||||
no_pad_features = [list(range(10)), list(range(10))]
|
||||
pad_features = [list(range(5)), list(range(10))]
|
||||
self._test_no_pad_and_pad(no_pad_features, pad_features)
|
||||
|
||||
def test_data_collator_for_whole_word_mask(self):
|
||||
features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
data_collator = DataCollatorForWholeWordMask(tokenizer, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
self.assertEqual(batch["labels"].shape, (2, 10))
|
||||
|
||||
def test_plm(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
no_pad_features = [{"input_ids": list(range(10))}, {"input_ids": list(range(10))}]
|
||||
pad_features = [{"input_ids": list(range(5))}, {"input_ids": list(range(10))}]
|
||||
|
||||
data_collator = DataCollatorForPermutationLanguageModeling(tokenizer, return_tensors="np")
|
||||
|
||||
batch = data_collator(pad_features)
|
||||
self.assertIsInstance(batch, dict)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
self.assertEqual(batch["perm_mask"].shape, (2, 10, 10))
|
||||
self.assertEqual(batch["target_mapping"].shape, (2, 10, 10))
|
||||
self.assertEqual(batch["labels"].shape, (2, 10))
|
||||
|
||||
batch = data_collator(no_pad_features)
|
||||
self.assertIsInstance(batch, dict)
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 10))
|
||||
self.assertEqual(batch["perm_mask"].shape, (2, 10, 10))
|
||||
self.assertEqual(batch["target_mapping"].shape, (2, 10, 10))
|
||||
self.assertEqual(batch["labels"].shape, (2, 10))
|
||||
|
||||
example = [np.random.randint(0, 5, [5])]
|
||||
with self.assertRaises(ValueError):
|
||||
# Expect error due to odd sequence length
|
||||
data_collator(example)
|
||||
|
||||
def test_nsp(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{"input_ids": [0, 1, 2, 3, 4], "token_type_ids": [0, 1, 2, 3, 4], "next_sentence_label": i}
|
||||
for i in range(2)
|
||||
]
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 5))
|
||||
self.assertEqual(batch["token_type_ids"].shape, (2, 5))
|
||||
self.assertEqual(batch["labels"].shape, (2, 5))
|
||||
self.assertEqual(batch["next_sentence_label"].shape, (2,))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 8))
|
||||
self.assertEqual(batch["token_type_ids"].shape, (2, 8))
|
||||
self.assertEqual(batch["labels"].shape, (2, 8))
|
||||
self.assertEqual(batch["next_sentence_label"].shape, (2,))
|
||||
|
||||
def test_sop(self):
|
||||
tokenizer = BertTokenizer(self.vocab_file)
|
||||
features = [
|
||||
{
|
||||
"input_ids": np.array([0, 1, 2, 3, 4]),
|
||||
"token_type_ids": np.array([0, 1, 2, 3, 4]),
|
||||
"sentence_order_label": i,
|
||||
}
|
||||
for i in range(2)
|
||||
]
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 5))
|
||||
self.assertEqual(batch["token_type_ids"].shape, (2, 5))
|
||||
self.assertEqual(batch["labels"].shape, (2, 5))
|
||||
self.assertEqual(batch["sentence_order_label"].shape, (2,))
|
||||
|
||||
data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors="np")
|
||||
batch = data_collator(features)
|
||||
|
||||
self.assertEqual(batch["input_ids"].shape, (2, 8))
|
||||
self.assertEqual(batch["token_type_ids"].shape, (2, 8))
|
||||
self.assertEqual(batch["labels"].shape, (2, 8))
|
||||
self.assertEqual(batch["sentence_order_label"].shape, (2,))
|
||||
1897
tests/trainer/test_trainer.py
Normal file
1897
tests/trainer/test_trainer.py
Normal file
File diff suppressed because it is too large
Load Diff
242
tests/trainer/test_trainer_callback.py
Normal file
242
tests/trainer/test_trainer_callback.py
Normal file
@@ -0,0 +1,242 @@
|
||||
# Copyright 2020 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from transformers import (
|
||||
DefaultFlowCallback,
|
||||
IntervalStrategy,
|
||||
PrinterCallback,
|
||||
ProgressCallback,
|
||||
Trainer,
|
||||
TrainerCallback,
|
||||
TrainingArguments,
|
||||
is_torch_available,
|
||||
)
|
||||
from transformers.testing_utils import require_torch
|
||||
|
||||
|
||||
if is_torch_available():
|
||||
from transformers.trainer import DEFAULT_CALLBACKS
|
||||
|
||||
from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel
|
||||
|
||||
|
||||
class MyTestTrainerCallback(TrainerCallback):
|
||||
"A callback that registers the events that goes through."
|
||||
|
||||
def __init__(self):
|
||||
self.events = []
|
||||
|
||||
def on_init_end(self, args, state, control, **kwargs):
|
||||
self.events.append("on_init_end")
|
||||
|
||||
def on_train_begin(self, args, state, control, **kwargs):
|
||||
self.events.append("on_train_begin")
|
||||
|
||||
def on_train_end(self, args, state, control, **kwargs):
|
||||
self.events.append("on_train_end")
|
||||
|
||||
def on_epoch_begin(self, args, state, control, **kwargs):
|
||||
self.events.append("on_epoch_begin")
|
||||
|
||||
def on_epoch_end(self, args, state, control, **kwargs):
|
||||
self.events.append("on_epoch_end")
|
||||
|
||||
def on_step_begin(self, args, state, control, **kwargs):
|
||||
self.events.append("on_step_begin")
|
||||
|
||||
def on_step_end(self, args, state, control, **kwargs):
|
||||
self.events.append("on_step_end")
|
||||
|
||||
def on_evaluate(self, args, state, control, **kwargs):
|
||||
self.events.append("on_evaluate")
|
||||
|
||||
def on_save(self, args, state, control, **kwargs):
|
||||
self.events.append("on_save")
|
||||
|
||||
def on_log(self, args, state, control, **kwargs):
|
||||
self.events.append("on_log")
|
||||
|
||||
def on_prediction_step(self, args, state, control, **kwargs):
|
||||
self.events.append("on_prediction_step")
|
||||
|
||||
|
||||
@require_torch
|
||||
class TrainerCallbackTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.output_dir = tempfile.mkdtemp()
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.output_dir)
|
||||
|
||||
def get_trainer(self, a=0, b=0, train_len=64, eval_len=64, callbacks=None, disable_tqdm=False, **kwargs):
|
||||
# disable_tqdm in TrainingArguments has a flaky default since it depends on the level of logging. We make sure
|
||||
# its set to False since the tests later on depend on its value.
|
||||
train_dataset = RegressionDataset(length=train_len)
|
||||
eval_dataset = RegressionDataset(length=eval_len)
|
||||
config = RegressionModelConfig(a=a, b=b)
|
||||
model = RegressionPreTrainedModel(config)
|
||||
|
||||
args = TrainingArguments(self.output_dir, disable_tqdm=disable_tqdm, report_to=[], **kwargs)
|
||||
return Trainer(
|
||||
model,
|
||||
args,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=eval_dataset,
|
||||
callbacks=callbacks,
|
||||
)
|
||||
|
||||
def check_callbacks_equality(self, cbs1, cbs2):
|
||||
self.assertEqual(len(cbs1), len(cbs2))
|
||||
|
||||
# Order doesn't matter
|
||||
cbs1 = list(sorted(cbs1, key=lambda cb: cb.__name__ if isinstance(cb, type) else cb.__class__.__name__))
|
||||
cbs2 = list(sorted(cbs2, key=lambda cb: cb.__name__ if isinstance(cb, type) else cb.__class__.__name__))
|
||||
|
||||
for cb1, cb2 in zip(cbs1, cbs2):
|
||||
if isinstance(cb1, type) and isinstance(cb2, type):
|
||||
self.assertEqual(cb1, cb2)
|
||||
elif isinstance(cb1, type) and not isinstance(cb2, type):
|
||||
self.assertEqual(cb1, cb2.__class__)
|
||||
elif not isinstance(cb1, type) and isinstance(cb2, type):
|
||||
self.assertEqual(cb1.__class__, cb2)
|
||||
else:
|
||||
self.assertEqual(cb1, cb2)
|
||||
|
||||
def get_expected_events(self, trainer):
|
||||
expected_events = ["on_init_end", "on_train_begin"]
|
||||
step = 0
|
||||
train_dl_len = len(trainer.get_eval_dataloader())
|
||||
evaluation_events = ["on_prediction_step"] * len(trainer.get_eval_dataloader()) + ["on_log", "on_evaluate"]
|
||||
for _ in range(trainer.state.num_train_epochs):
|
||||
expected_events.append("on_epoch_begin")
|
||||
for _ in range(train_dl_len):
|
||||
step += 1
|
||||
expected_events += ["on_step_begin", "on_step_end"]
|
||||
if step % trainer.args.logging_steps == 0:
|
||||
expected_events.append("on_log")
|
||||
if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
|
||||
expected_events += evaluation_events.copy()
|
||||
if step % trainer.args.save_steps == 0:
|
||||
expected_events.append("on_save")
|
||||
expected_events.append("on_epoch_end")
|
||||
if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
|
||||
expected_events += evaluation_events.copy()
|
||||
expected_events += ["on_log", "on_train_end"]
|
||||
return expected_events
|
||||
|
||||
def test_init_callback(self):
|
||||
trainer = self.get_trainer()
|
||||
expected_callbacks = DEFAULT_CALLBACKS.copy() + [ProgressCallback]
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
# Callbacks passed at init are added to the default callbacks
|
||||
trainer = self.get_trainer(callbacks=[MyTestTrainerCallback])
|
||||
expected_callbacks.append(MyTestTrainerCallback)
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
# TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
|
||||
trainer = self.get_trainer(disable_tqdm=True)
|
||||
expected_callbacks = DEFAULT_CALLBACKS.copy() + [PrinterCallback]
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
def test_add_remove_callback(self):
|
||||
expected_callbacks = DEFAULT_CALLBACKS.copy() + [ProgressCallback]
|
||||
trainer = self.get_trainer()
|
||||
|
||||
# We can add, pop, or remove by class name
|
||||
trainer.remove_callback(DefaultFlowCallback)
|
||||
expected_callbacks.remove(DefaultFlowCallback)
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
trainer = self.get_trainer()
|
||||
cb = trainer.pop_callback(DefaultFlowCallback)
|
||||
self.assertEqual(cb.__class__, DefaultFlowCallback)
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
trainer.add_callback(DefaultFlowCallback)
|
||||
expected_callbacks.insert(0, DefaultFlowCallback)
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
# We can also add, pop, or remove by instance
|
||||
trainer = self.get_trainer()
|
||||
cb = trainer.callback_handler.callbacks[0]
|
||||
trainer.remove_callback(cb)
|
||||
expected_callbacks.remove(DefaultFlowCallback)
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
trainer = self.get_trainer()
|
||||
cb1 = trainer.callback_handler.callbacks[0]
|
||||
cb2 = trainer.pop_callback(cb1)
|
||||
self.assertEqual(cb1, cb2)
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
trainer.add_callback(cb1)
|
||||
expected_callbacks.insert(0, DefaultFlowCallback)
|
||||
self.check_callbacks_equality(trainer.callback_handler.callbacks, expected_callbacks)
|
||||
|
||||
def test_event_flow(self):
|
||||
import warnings
|
||||
|
||||
# XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
|
||||
warnings.simplefilter(action="ignore", category=UserWarning)
|
||||
|
||||
trainer = self.get_trainer(callbacks=[MyTestTrainerCallback])
|
||||
trainer.train()
|
||||
events = trainer.callback_handler.callbacks[-2].events
|
||||
self.assertEqual(events, self.get_expected_events(trainer))
|
||||
|
||||
# Independent log/save/eval
|
||||
trainer = self.get_trainer(callbacks=[MyTestTrainerCallback], logging_steps=5)
|
||||
trainer.train()
|
||||
events = trainer.callback_handler.callbacks[-2].events
|
||||
self.assertEqual(events, self.get_expected_events(trainer))
|
||||
|
||||
trainer = self.get_trainer(callbacks=[MyTestTrainerCallback], save_steps=5)
|
||||
trainer.train()
|
||||
events = trainer.callback_handler.callbacks[-2].events
|
||||
self.assertEqual(events, self.get_expected_events(trainer))
|
||||
|
||||
trainer = self.get_trainer(callbacks=[MyTestTrainerCallback], eval_steps=5, evaluation_strategy="steps")
|
||||
trainer.train()
|
||||
events = trainer.callback_handler.callbacks[-2].events
|
||||
self.assertEqual(events, self.get_expected_events(trainer))
|
||||
|
||||
trainer = self.get_trainer(callbacks=[MyTestTrainerCallback], evaluation_strategy="epoch")
|
||||
trainer.train()
|
||||
events = trainer.callback_handler.callbacks[-2].events
|
||||
self.assertEqual(events, self.get_expected_events(trainer))
|
||||
|
||||
# A bit of everything
|
||||
trainer = self.get_trainer(
|
||||
callbacks=[MyTestTrainerCallback],
|
||||
logging_steps=3,
|
||||
save_steps=10,
|
||||
eval_steps=5,
|
||||
evaluation_strategy="steps",
|
||||
)
|
||||
trainer.train()
|
||||
events = trainer.callback_handler.callbacks[-2].events
|
||||
self.assertEqual(events, self.get_expected_events(trainer))
|
||||
|
||||
# warning should be emitted for duplicated callbacks
|
||||
with patch("transformers.trainer_callback.logger.warning") as warn_mock:
|
||||
trainer = self.get_trainer(
|
||||
callbacks=[MyTestTrainerCallback, MyTestTrainerCallback],
|
||||
)
|
||||
assert str(MyTestTrainerCallback) in warn_mock.call_args[0][0]
|
||||
143
tests/trainer/test_trainer_distributed.py
Normal file
143
tests/trainer/test_trainer_distributed.py
Normal file
@@ -0,0 +1,143 @@
|
||||
# Copyright 2020 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import sys
|
||||
from typing import Dict
|
||||
|
||||
from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
|
||||
from transformers.testing_utils import (
|
||||
TestCasePlus,
|
||||
execute_subprocess_async,
|
||||
get_torch_dist_unique_port,
|
||||
require_torch_multi_gpu,
|
||||
)
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
if is_torch_available():
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from transformers import Trainer
|
||||
|
||||
class DummyDataset(Dataset):
|
||||
def __init__(self, length: int = 101):
|
||||
self.length = length
|
||||
|
||||
def __len__(self):
|
||||
return self.length
|
||||
|
||||
def __getitem__(self, i) -> int:
|
||||
return i
|
||||
|
||||
class DummyDataCollator:
|
||||
def __call__(self, features):
|
||||
return {"input_ids": torch.tensor(features), "labels": torch.tensor(features)}
|
||||
|
||||
class DummyModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Add some (unused) params otherwise DDP will complain.
|
||||
self.fc = nn.Linear(120, 80)
|
||||
|
||||
def forward(self, input_ids, labels=None):
|
||||
if labels is not None:
|
||||
return torch.tensor(0.0, device=input_ids.device), input_ids
|
||||
else:
|
||||
return input_ids
|
||||
|
||||
|
||||
class TestTrainerDistributed(TestCasePlus):
|
||||
@require_torch_multi_gpu
|
||||
def test_trainer(self):
|
||||
|
||||
distributed_args = f"""
|
||||
-m torch.distributed.launch
|
||||
--nproc_per_node={torch.cuda.device_count()}
|
||||
--master_port={get_torch_dist_unique_port()}
|
||||
{self.test_file_dir}/test_trainer_distributed.py
|
||||
""".split()
|
||||
output_dir = self.get_auto_remove_tmp_dir()
|
||||
args = f"--output_dir {output_dir}".split()
|
||||
cmd = [sys.executable] + distributed_args + args
|
||||
execute_subprocess_async(cmd, env=self.get_env())
|
||||
# successful return here == success - any errors would have caused an error in the sub-call
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
|
||||
#
|
||||
# PYTHONPATH="src" python -m torch.distributed.launch --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py
|
||||
|
||||
parser = HfArgumentParser((TrainingArguments,))
|
||||
training_args = parser.parse_args_into_dataclasses()[0]
|
||||
|
||||
logger.warning(
|
||||
f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
|
||||
f"distributed training: {training_args.local_rank != -1}"
|
||||
)
|
||||
|
||||
# Essentially, what we want to verify in the distributed case is that we get all samples back,
|
||||
# in the right order. (this is crucial for prediction for instance)
|
||||
for dataset_length in [101, 40, 7]:
|
||||
dataset = DummyDataset(dataset_length)
|
||||
|
||||
def compute_metrics(p: EvalPrediction) -> Dict:
|
||||
sequential = list(range(len(dataset)))
|
||||
success = p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
|
||||
if not success and training_args.local_rank == 0:
|
||||
logger.warning(
|
||||
"Predictions and/or labels do not match expected results:\n - predictions: "
|
||||
f"{p.predictions.tolist()}\n - labels: {p.label_ids.tolist()}\n - expected: {sequential}"
|
||||
)
|
||||
return {"success": success}
|
||||
|
||||
trainer = Trainer(
|
||||
model=DummyModel(),
|
||||
args=training_args,
|
||||
data_collator=DummyDataCollator(),
|
||||
eval_dataset=dataset,
|
||||
compute_metrics=compute_metrics,
|
||||
)
|
||||
metrics = trainer.evaluate()
|
||||
logger.info(metrics)
|
||||
if metrics["eval_success"] is not True:
|
||||
logger.error(metrics)
|
||||
exit(1)
|
||||
|
||||
p = trainer.predict(dataset)
|
||||
logger.info(p.metrics)
|
||||
if p.metrics["test_success"] is not True:
|
||||
logger.error(p.metrics)
|
||||
exit(1)
|
||||
|
||||
trainer.args.eval_accumulation_steps = 2
|
||||
|
||||
metrics = trainer.evaluate()
|
||||
logger.info(metrics)
|
||||
if metrics["eval_success"] is not True:
|
||||
logger.error(metrics)
|
||||
exit(1)
|
||||
|
||||
p = trainer.predict(dataset)
|
||||
logger.info(p.metrics)
|
||||
if p.metrics["test_success"] is not True:
|
||||
logger.error(p.metrics)
|
||||
exit(1)
|
||||
|
||||
trainer.args.eval_accumulation_steps = None
|
||||
126
tests/trainer/test_trainer_seq2seq.py
Normal file
126
tests/trainer/test_trainer_seq2seq.py
Normal file
@@ -0,0 +1,126 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2020 the HuggingFace Inc. team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from transformers import BertTokenizer, EncoderDecoderModel, Seq2SeqTrainer, Seq2SeqTrainingArguments
|
||||
from transformers.file_utils import is_datasets_available
|
||||
from transformers.testing_utils import TestCasePlus, require_torch, slow
|
||||
|
||||
|
||||
if is_datasets_available():
|
||||
import datasets
|
||||
|
||||
|
||||
class Seq2seqTrainerTester(TestCasePlus):
|
||||
@slow
|
||||
@require_torch
|
||||
def test_finetune_bert2bert(self):
|
||||
bert2bert = EncoderDecoderModel.from_encoder_decoder_pretrained("prajjwal1/bert-tiny", "prajjwal1/bert-tiny")
|
||||
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
|
||||
|
||||
bert2bert.config.vocab_size = bert2bert.config.encoder.vocab_size
|
||||
bert2bert.config.eos_token_id = tokenizer.sep_token_id
|
||||
bert2bert.config.decoder_start_token_id = tokenizer.cls_token_id
|
||||
bert2bert.config.max_length = 128
|
||||
|
||||
train_dataset = datasets.load_dataset("cnn_dailymail", "3.0.0", split="train[:1%]")
|
||||
val_dataset = datasets.load_dataset("cnn_dailymail", "3.0.0", split="validation[:1%]")
|
||||
|
||||
train_dataset = train_dataset.select(range(32))
|
||||
val_dataset = val_dataset.select(range(16))
|
||||
|
||||
batch_size = 4
|
||||
|
||||
def _map_to_encoder_decoder_inputs(batch):
|
||||
# Tokenizer will automatically set [BOS] <text> [EOS]
|
||||
inputs = tokenizer(batch["article"], padding="max_length", truncation=True, max_length=512)
|
||||
outputs = tokenizer(batch["highlights"], padding="max_length", truncation=True, max_length=128)
|
||||
batch["input_ids"] = inputs.input_ids
|
||||
batch["attention_mask"] = inputs.attention_mask
|
||||
|
||||
batch["decoder_input_ids"] = outputs.input_ids
|
||||
batch["labels"] = outputs.input_ids.copy()
|
||||
batch["labels"] = [
|
||||
[-100 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch["labels"]
|
||||
]
|
||||
batch["decoder_attention_mask"] = outputs.attention_mask
|
||||
|
||||
assert all([len(x) == 512 for x in inputs.input_ids])
|
||||
assert all([len(x) == 128 for x in outputs.input_ids])
|
||||
|
||||
return batch
|
||||
|
||||
def _compute_metrics(pred):
|
||||
labels_ids = pred.label_ids
|
||||
pred_ids = pred.predictions
|
||||
|
||||
# all unnecessary tokens are removed
|
||||
pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
|
||||
label_str = tokenizer.batch_decode(labels_ids, skip_special_tokens=True)
|
||||
|
||||
accuracy = sum([int(pred_str[i] == label_str[i]) for i in range(len(pred_str))]) / len(pred_str)
|
||||
|
||||
return {"accuracy": accuracy}
|
||||
|
||||
# map train dataset
|
||||
train_dataset = train_dataset.map(
|
||||
_map_to_encoder_decoder_inputs,
|
||||
batched=True,
|
||||
batch_size=batch_size,
|
||||
remove_columns=["article", "highlights"],
|
||||
)
|
||||
train_dataset.set_format(
|
||||
type="torch",
|
||||
columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"],
|
||||
)
|
||||
|
||||
# same for validation dataset
|
||||
val_dataset = val_dataset.map(
|
||||
_map_to_encoder_decoder_inputs,
|
||||
batched=True,
|
||||
batch_size=batch_size,
|
||||
remove_columns=["article", "highlights"],
|
||||
)
|
||||
val_dataset.set_format(
|
||||
type="torch",
|
||||
columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"],
|
||||
)
|
||||
|
||||
output_dir = self.get_auto_remove_tmp_dir()
|
||||
|
||||
training_args = Seq2SeqTrainingArguments(
|
||||
output_dir=output_dir,
|
||||
per_device_train_batch_size=batch_size,
|
||||
per_device_eval_batch_size=batch_size,
|
||||
predict_with_generate=True,
|
||||
evaluation_strategy="steps",
|
||||
do_train=True,
|
||||
do_eval=True,
|
||||
warmup_steps=0,
|
||||
eval_steps=2,
|
||||
logging_steps=2,
|
||||
)
|
||||
|
||||
# instantiate trainer
|
||||
trainer = Seq2SeqTrainer(
|
||||
model=bert2bert,
|
||||
args=training_args,
|
||||
compute_metrics=_compute_metrics,
|
||||
train_dataset=train_dataset,
|
||||
eval_dataset=val_dataset,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
|
||||
# start training
|
||||
trainer.train()
|
||||
131
tests/trainer/test_trainer_tpu.py
Normal file
131
tests/trainer/test_trainer_tpu.py
Normal file
@@ -0,0 +1,131 @@
|
||||
# Copyright 2020 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# This test is meant to be run in on an instance with TPUs like this:
|
||||
#
|
||||
# python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
|
||||
#
|
||||
# Replace 8 with the number of TPU cores you have.
|
||||
#
|
||||
|
||||
import sys
|
||||
from typing import Dict
|
||||
|
||||
from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
if is_torch_available():
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from transformers import Trainer
|
||||
|
||||
class DummyDataset(Dataset):
|
||||
def __init__(self, length: int = 101):
|
||||
self.length = length
|
||||
|
||||
def __len__(self):
|
||||
return self.length
|
||||
|
||||
def __getitem__(self, i) -> int:
|
||||
return i
|
||||
|
||||
class DummyDataCollator:
|
||||
def __call__(self, features):
|
||||
return {"input_ids": torch.tensor(features), "labels": torch.tensor(features)}
|
||||
|
||||
class DummyModel(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Add some (unused) params otherwise DDP will complain.
|
||||
self.fc = nn.Linear(120, 80)
|
||||
|
||||
def forward(self, input_ids, labels=None):
|
||||
if labels is not None:
|
||||
return torch.tensor(0.0, device=input_ids.device), input_ids
|
||||
else:
|
||||
return input_ids
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((TrainingArguments,))
|
||||
sys.argv += ["--output_dir", "./examples"]
|
||||
training_args = parser.parse_args_into_dataclasses()[0]
|
||||
|
||||
logger.warning(
|
||||
f"Process rank: {training_args.local_rank}, device: {training_args.device}, "
|
||||
f"tpu_num_cores: {training_args.tpu_num_cores}",
|
||||
)
|
||||
|
||||
# Essentially, what we want to verify in the distributed case is
|
||||
# that we get all samples back, in the right order.
|
||||
# (this is crucial for prediction for instance)
|
||||
for dataset_length in [1001, 256, 15]:
|
||||
dataset = DummyDataset(dataset_length)
|
||||
|
||||
def compute_metrics(p: EvalPrediction) -> Dict:
|
||||
sequential = list(range(len(dataset)))
|
||||
success = p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
|
||||
return {"success": success}
|
||||
|
||||
trainer = Trainer(
|
||||
model=DummyModel(),
|
||||
args=training_args,
|
||||
data_collator=DummyDataCollator(),
|
||||
eval_dataset=dataset,
|
||||
compute_metrics=compute_metrics,
|
||||
)
|
||||
metrics = trainer.evaluate()
|
||||
logger.info(metrics)
|
||||
if metrics["eval_success"] is not True:
|
||||
logger.error(metrics)
|
||||
exit(1)
|
||||
|
||||
p = trainer.predict(dataset)
|
||||
logger.info(p.metrics)
|
||||
if p.metrics["test_success"] is not True:
|
||||
logger.error(p.metrics)
|
||||
exit(1)
|
||||
|
||||
trainer.args.eval_accumulation_steps = 2
|
||||
|
||||
metrics = trainer.evaluate()
|
||||
logger.info(metrics)
|
||||
if metrics["eval_success"] is not True:
|
||||
logger.error(metrics)
|
||||
exit(1)
|
||||
|
||||
p = trainer.predict(dataset)
|
||||
logger.info(p.metrics)
|
||||
if p.metrics["test_success"] is not True:
|
||||
logger.error(p.metrics)
|
||||
exit(1)
|
||||
|
||||
trainer.args.eval_accumulation_steps = None
|
||||
|
||||
logger.info("🔥 All distributed tests successful")
|
||||
|
||||
|
||||
def _mp_fn(index):
|
||||
# For xla_spawn (TPUs)
|
||||
main()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
422
tests/trainer/test_trainer_utils.py
Normal file
422
tests/trainer/test_trainer_utils.py
Normal file
@@ -0,0 +1,422 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2018 the HuggingFace Inc. team.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import copy
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from transformers.file_utils import is_torch_available
|
||||
from transformers.testing_utils import require_torch
|
||||
|
||||
|
||||
if is_torch_available():
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import IterableDataset
|
||||
|
||||
from transformers.modeling_outputs import SequenceClassifierOutput
|
||||
from transformers.tokenization_utils_base import BatchEncoding
|
||||
from transformers.trainer_pt_utils import (
|
||||
DistributedLengthGroupedSampler,
|
||||
DistributedSamplerWithLoop,
|
||||
DistributedTensorGatherer,
|
||||
IterableDatasetShard,
|
||||
LabelSmoother,
|
||||
LengthGroupedSampler,
|
||||
SequentialDistributedSampler,
|
||||
ShardSampler,
|
||||
get_parameter_names,
|
||||
)
|
||||
|
||||
class TstLayer(nn.Module):
|
||||
def __init__(self, hidden_size):
|
||||
super().__init__()
|
||||
self.linear1 = nn.Linear(hidden_size, hidden_size)
|
||||
self.ln1 = nn.LayerNorm(hidden_size)
|
||||
self.linear2 = nn.Linear(hidden_size, hidden_size)
|
||||
self.ln2 = nn.LayerNorm(hidden_size)
|
||||
self.bias = nn.Parameter(torch.zeros(hidden_size))
|
||||
|
||||
def forward(self, x):
|
||||
h = self.ln1(nn.functional.relu(self.linear1(x)))
|
||||
h = nn.functional.relu(self.linear2(x))
|
||||
return self.ln2(x + h + self.bias)
|
||||
|
||||
class RandomIterableDataset(IterableDataset):
|
||||
# For testing, an iterable dataset of random length
|
||||
def __init__(self, p_stop=0.01, max_length=1000):
|
||||
self.p_stop = p_stop
|
||||
self.max_length = max_length
|
||||
self.generator = torch.Generator()
|
||||
|
||||
def __iter__(self):
|
||||
count = 0
|
||||
stop = False
|
||||
while not stop and count < self.max_length:
|
||||
yield count
|
||||
count += 1
|
||||
number = torch.rand(1, generator=self.generator).item()
|
||||
stop = number < self.p_stop
|
||||
|
||||
|
||||
@require_torch
|
||||
class TrainerUtilsTest(unittest.TestCase):
|
||||
def test_distributed_tensor_gatherer(self):
|
||||
# Simulate a result with a dataset of size 21, 4 processes and chunks of lengths 2, 3, 1
|
||||
world_size = 4
|
||||
num_samples = 21
|
||||
input_indices = [
|
||||
[0, 1, 6, 7, 12, 13, 18, 19],
|
||||
[2, 3, 4, 8, 9, 10, 14, 15, 16, 20, 0, 1],
|
||||
[5, 11, 17, 2],
|
||||
]
|
||||
|
||||
predictions = np.random.normal(size=(num_samples, 13))
|
||||
gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)
|
||||
for indices in input_indices:
|
||||
gatherer.add_arrays(predictions[indices])
|
||||
result = gatherer.finalize()
|
||||
self.assertTrue(np.array_equal(result, predictions))
|
||||
|
||||
# With nested tensors
|
||||
gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)
|
||||
for indices in input_indices:
|
||||
gatherer.add_arrays([predictions[indices], [predictions[indices], predictions[indices]]])
|
||||
result = gatherer.finalize()
|
||||
self.assertTrue(isinstance(result, list))
|
||||
self.assertTrue(len(result), 2)
|
||||
self.assertTrue(isinstance(result[1], list))
|
||||
self.assertTrue(len(result[1]), 2)
|
||||
self.assertTrue(np.array_equal(result[0], predictions))
|
||||
self.assertTrue(np.array_equal(result[1][0], predictions))
|
||||
self.assertTrue(np.array_equal(result[1][1], predictions))
|
||||
|
||||
def test_distributed_tensor_gatherer_different_shapes(self):
|
||||
# Simulate a result with a dataset of size 21, 4 processes and chunks of lengths 2, 3, 1
|
||||
world_size = 4
|
||||
num_samples = 21
|
||||
input_indices = [
|
||||
[0, 1, 6, 7, 12, 13, 18, 19],
|
||||
[2, 3, 4, 8, 9, 10, 14, 15, 16, 20, 0, 1],
|
||||
[5, 11, 17, 2],
|
||||
]
|
||||
sequence_lengths = [8, 10, 13]
|
||||
|
||||
predictions = np.random.normal(size=(num_samples, 13))
|
||||
gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)
|
||||
for indices, seq_length in zip(input_indices, sequence_lengths):
|
||||
gatherer.add_arrays(predictions[indices, :seq_length])
|
||||
result = gatherer.finalize()
|
||||
|
||||
# Remove the extra samples added at the end for a round multiple of num processes.
|
||||
actual_indices = [input_indices[0], input_indices[1][:-2], input_indices[2][:-1]]
|
||||
for indices, seq_length in zip(actual_indices, sequence_lengths):
|
||||
self.assertTrue(np.array_equal(result[indices, :seq_length], predictions[indices, :seq_length]))
|
||||
|
||||
# With nested tensors
|
||||
predictions = np.random.normal(size=(num_samples, 13))
|
||||
gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)
|
||||
for indices, seq_length in zip(input_indices, sequence_lengths):
|
||||
gatherer.add_arrays([predictions[indices, :seq_length], predictions[indices]])
|
||||
result = gatherer.finalize()
|
||||
|
||||
for indices, seq_length in zip(actual_indices, sequence_lengths):
|
||||
self.assertTrue(np.array_equal(result[0][indices, :seq_length], predictions[indices, :seq_length]))
|
||||
self.assertTrue(np.array_equal(result[1], predictions))
|
||||
|
||||
# Check if works if varying seq_length is second
|
||||
gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)
|
||||
for indices, seq_length in zip(input_indices, sequence_lengths):
|
||||
gatherer.add_arrays([predictions[indices], predictions[indices, :seq_length]])
|
||||
result = gatherer.finalize()
|
||||
|
||||
self.assertTrue(np.array_equal(result[0], predictions))
|
||||
for indices, seq_length in zip(actual_indices, sequence_lengths):
|
||||
self.assertTrue(np.array_equal(result[1][indices, :seq_length], predictions[indices, :seq_length]))
|
||||
|
||||
def test_label_smoothing(self):
|
||||
epsilon = 0.1
|
||||
num_labels = 12
|
||||
random_logits = torch.randn(4, 5, num_labels)
|
||||
random_labels = torch.randint(0, num_labels, (4, 5))
|
||||
loss = nn.functional.cross_entropy(random_logits.view(-1, num_labels), random_labels.view(-1))
|
||||
model_output = SequenceClassifierOutput(logits=random_logits)
|
||||
label_smoothed_loss = LabelSmoother(0.1)(model_output, random_labels)
|
||||
log_probs = -nn.functional.log_softmax(random_logits, dim=-1)
|
||||
expected_loss = (1 - epsilon) * loss + epsilon * log_probs.mean()
|
||||
self.assertTrue(torch.allclose(label_smoothed_loss, expected_loss))
|
||||
|
||||
# With a few -100 labels
|
||||
random_labels[0, 1] = -100
|
||||
random_labels[2, 1] = -100
|
||||
random_labels[2, 3] = -100
|
||||
|
||||
loss = nn.functional.cross_entropy(random_logits.view(-1, num_labels), random_labels.view(-1))
|
||||
model_output = SequenceClassifierOutput(logits=random_logits)
|
||||
label_smoothed_loss = LabelSmoother(0.1)(model_output, random_labels)
|
||||
log_probs = -nn.functional.log_softmax(random_logits, dim=-1)
|
||||
# Mask the log probs with the -100 labels
|
||||
log_probs[0, 1] = 0.0
|
||||
log_probs[2, 1] = 0.0
|
||||
log_probs[2, 3] = 0.0
|
||||
expected_loss = (1 - epsilon) * loss + epsilon * log_probs.sum() / (num_labels * 17)
|
||||
self.assertTrue(torch.allclose(label_smoothed_loss, expected_loss))
|
||||
|
||||
def test_group_by_length(self):
|
||||
# Get some inputs of random lengths
|
||||
lengths = torch.randint(0, 25, (100,)).tolist()
|
||||
# Put one bigger than the others to check it ends up in first position
|
||||
lengths[32] = 50
|
||||
|
||||
indices = list(LengthGroupedSampler(4, lengths=lengths))
|
||||
# The biggest element should be first
|
||||
self.assertEqual(lengths[indices[0]], 50)
|
||||
# The indices should be a permutation of range(100)
|
||||
self.assertEqual(list(sorted(indices)), list(range(100)))
|
||||
|
||||
def test_group_by_length_with_dict(self):
|
||||
# Get some inputs of random lengths
|
||||
data = []
|
||||
for _ in range(6):
|
||||
input_ids = torch.randint(0, 25, (100,)).tolist()
|
||||
data.append({"input_ids": input_ids})
|
||||
# Put one bigger than the others to check it ends up in first position
|
||||
data[3]["input_ids"] = torch.randint(0, 25, (105,)).tolist()
|
||||
|
||||
indices = list(LengthGroupedSampler(4, dataset=data))
|
||||
# The biggest element should be first
|
||||
self.assertEqual(len(data[indices[0]]["input_ids"]), 105)
|
||||
# The indices should be a permutation of range(6)
|
||||
self.assertEqual(list(sorted(indices)), list(range(6)))
|
||||
|
||||
def test_group_by_length_with_batch_encoding(self):
|
||||
# Get some inputs of random lengths
|
||||
data = []
|
||||
for _ in range(6):
|
||||
input_ids = torch.randint(0, 25, (100,)).tolist()
|
||||
data.append(BatchEncoding({"input_ids": input_ids}))
|
||||
# Put one bigger than the others to check it ends up in first position
|
||||
data[3]["input_ids"] = torch.randint(0, 25, (105,)).tolist()
|
||||
|
||||
indices = list(LengthGroupedSampler(4, dataset=data))
|
||||
# The biggest element should be first
|
||||
self.assertEqual(len(data[indices[0]]["input_ids"]), 105)
|
||||
# The indices should be a permutation of range(6)
|
||||
self.assertEqual(list(sorted(indices)), list(range(6)))
|
||||
|
||||
def test_distributed_length_grouped(self):
|
||||
# Get some inputs of random lengths
|
||||
lengths = torch.randint(0, 25, (100,)).tolist()
|
||||
# Put one bigger than the others to check it ends up in first position
|
||||
lengths[32] = 50
|
||||
|
||||
indices_process_0 = list(DistributedLengthGroupedSampler(4, num_replicas=2, rank=0, lengths=lengths))
|
||||
indices_process_1 = list(DistributedLengthGroupedSampler(4, num_replicas=2, rank=1, lengths=lengths))
|
||||
# The biggest element should be first
|
||||
self.assertEqual(lengths[indices_process_0[0]], 50)
|
||||
# The indices should be a permutation of range(100)
|
||||
self.assertEqual(list(sorted(indices_process_0 + indices_process_1)), list(range(100)))
|
||||
|
||||
def test_get_parameter_names(self):
|
||||
model = nn.Sequential(TstLayer(128), nn.ModuleList([TstLayer(128), TstLayer(128)]))
|
||||
# fmt: off
|
||||
self.assertEqual(
|
||||
get_parameter_names(model, [nn.LayerNorm]),
|
||||
['0.linear1.weight', '0.linear1.bias', '0.linear2.weight', '0.linear2.bias', '0.bias', '1.0.linear1.weight', '1.0.linear1.bias', '1.0.linear2.weight', '1.0.linear2.bias', '1.0.bias', '1.1.linear1.weight', '1.1.linear1.bias', '1.1.linear2.weight', '1.1.linear2.bias', '1.1.bias']
|
||||
)
|
||||
# fmt: on
|
||||
|
||||
def test_distributed_sampler_with_loop(self):
|
||||
batch_size = 16
|
||||
for length in [23, 64, 123]:
|
||||
dataset = list(range(length))
|
||||
shard1 = DistributedSamplerWithLoop(dataset, batch_size, num_replicas=2, rank=0)
|
||||
shard2 = DistributedSamplerWithLoop(dataset, batch_size, num_replicas=2, rank=1)
|
||||
|
||||
# Set seeds
|
||||
shard1.set_epoch(0)
|
||||
shard2.set_epoch(0)
|
||||
|
||||
# Sample
|
||||
samples1 = list(shard1)
|
||||
samples2 = list(shard2)
|
||||
|
||||
self.assertTrue(len(samples1) % batch_size == 0)
|
||||
self.assertTrue(len(samples2) % batch_size == 0)
|
||||
|
||||
total = []
|
||||
for sample1, sample2 in zip(samples1, samples2):
|
||||
total += [sample1, sample2]
|
||||
|
||||
self.assertEqual(set(total[:length]), set(dataset))
|
||||
self.assertEqual(set(total[length:]), set(total[: (len(total) - length)]))
|
||||
|
||||
def test_sequential_distributed_sampler(self):
|
||||
batch_size = 16
|
||||
for length in [23, 64, 123]:
|
||||
dataset = list(range(length))
|
||||
shard1 = SequentialDistributedSampler(dataset, num_replicas=2, rank=0)
|
||||
shard2 = SequentialDistributedSampler(dataset, num_replicas=2, rank=1)
|
||||
|
||||
# Sample
|
||||
samples1 = list(shard1)
|
||||
samples2 = list(shard2)
|
||||
|
||||
total = samples1 + samples2
|
||||
|
||||
self.assertListEqual(total[:length], dataset)
|
||||
self.assertListEqual(total[length:], dataset[: (len(total) - length)])
|
||||
|
||||
# With a batch_size passed
|
||||
shard1 = SequentialDistributedSampler(dataset, num_replicas=2, rank=0, batch_size=batch_size)
|
||||
shard2 = SequentialDistributedSampler(dataset, num_replicas=2, rank=1, batch_size=batch_size)
|
||||
|
||||
# Sample
|
||||
samples1 = list(shard1)
|
||||
samples2 = list(shard2)
|
||||
|
||||
self.assertTrue(len(samples1) % batch_size == 0)
|
||||
self.assertTrue(len(samples2) % batch_size == 0)
|
||||
|
||||
total = samples1 + samples2
|
||||
|
||||
self.assertListEqual(total[:length], dataset)
|
||||
self.assertListEqual(total[length:], dataset[: (len(total) - length)])
|
||||
|
||||
def check_iterable_dataset_shard(self, dataset, batch_size, drop_last, num_processes=2, epoch=0):
|
||||
# Set the seed for the base dataset to get the proper reference.
|
||||
dataset.generator.manual_seed(epoch)
|
||||
reference = list(dataset)
|
||||
|
||||
shards = [
|
||||
IterableDatasetShard(
|
||||
dataset, batch_size=batch_size, drop_last=drop_last, num_processes=num_processes, process_index=i
|
||||
)
|
||||
for i in range(num_processes)
|
||||
]
|
||||
for shard in shards:
|
||||
shard.set_epoch(epoch)
|
||||
shard_lists = [list(shard) for shard in shards]
|
||||
|
||||
for shard in shard_lists:
|
||||
# All shards have a number of samples that is a round multiple of batch size
|
||||
self.assertTrue(len(shard) % batch_size == 0)
|
||||
# All shards have the same number of samples
|
||||
self.assertEqual(len(shard), len(shard_lists[0]))
|
||||
|
||||
for shard in shards:
|
||||
# All shards know the total number of samples
|
||||
self.assertEqual(shard.num_examples, len(reference))
|
||||
|
||||
observed = []
|
||||
for idx in range(0, len(shard_lists[0]), batch_size):
|
||||
for shard in shard_lists:
|
||||
observed += shard[idx : idx + batch_size]
|
||||
|
||||
# If drop_last is False we loop through samples at the beginning to have a size that is a round multiple of
|
||||
# batch_size
|
||||
if not drop_last:
|
||||
while len(reference) < len(observed):
|
||||
reference += reference
|
||||
self.assertListEqual(observed, reference[: len(observed)])
|
||||
|
||||
# Check equivalence between IterableDataset and ShardSampler
|
||||
dataset.generator.manual_seed(epoch)
|
||||
reference = list(dataset)
|
||||
|
||||
sampler_shards = [
|
||||
ShardSampler(
|
||||
reference, batch_size=batch_size, drop_last=drop_last, num_processes=num_processes, process_index=i
|
||||
)
|
||||
for i in range(num_processes)
|
||||
]
|
||||
for shard, sampler_shard in zip(shard_lists, sampler_shards):
|
||||
self.assertListEqual(shard, list(sampler_shard))
|
||||
|
||||
def test_iterable_dataset_shard(self):
|
||||
dataset = RandomIterableDataset()
|
||||
|
||||
self.check_iterable_dataset_shard(dataset, 4, drop_last=True, num_processes=2, epoch=0)
|
||||
self.check_iterable_dataset_shard(dataset, 4, drop_last=False, num_processes=2, epoch=0)
|
||||
|
||||
self.check_iterable_dataset_shard(dataset, 4, drop_last=True, num_processes=3, epoch=42)
|
||||
self.check_iterable_dataset_shard(dataset, 4, drop_last=False, num_processes=3, epoch=42)
|
||||
|
||||
def test_iterable_dataset_shard_with_length(self):
|
||||
sampler_shards = [
|
||||
IterableDatasetShard(list(range(100)), batch_size=4, drop_last=True, num_processes=2, process_index=i)
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
# Build expected shards: each process will have batches of size 4 until there is not enough elements to
|
||||
# form two full batches (so we stop at 96 = (100 // (4 * 2)) * 4)
|
||||
expected_shards = [[], []]
|
||||
current_shard = 0
|
||||
for i in range(0, 96, 4):
|
||||
expected_shards[current_shard].extend(list(range(i, i + 4)))
|
||||
current_shard = 1 - current_shard
|
||||
|
||||
self.assertListEqual([list(shard) for shard in sampler_shards], expected_shards)
|
||||
self.assertListEqual([len(shard) for shard in sampler_shards], [len(shard) for shard in expected_shards])
|
||||
|
||||
sampler_shards = [
|
||||
IterableDatasetShard(list(range(100)), batch_size=4, drop_last=False, num_processes=2, process_index=i)
|
||||
for i in range(2)
|
||||
]
|
||||
# When drop_last=False, we get two last full batches by looping back to the beginning.
|
||||
expected_shards[0].extend(list(range(96, 100)))
|
||||
expected_shards[1].extend(list(range(0, 4)))
|
||||
|
||||
self.assertListEqual([list(shard) for shard in sampler_shards], expected_shards)
|
||||
self.assertListEqual([len(shard) for shard in sampler_shards], [len(shard) for shard in expected_shards])
|
||||
|
||||
def check_shard_sampler(self, dataset, batch_size, drop_last, num_processes=2):
|
||||
shards = [
|
||||
ShardSampler(
|
||||
dataset, batch_size=batch_size, drop_last=drop_last, num_processes=num_processes, process_index=i
|
||||
)
|
||||
for i in range(num_processes)
|
||||
]
|
||||
shard_lists = [list(shard) for shard in shards]
|
||||
|
||||
for shard in shard_lists:
|
||||
# All shards have a number of samples that is a round multiple of batch size
|
||||
self.assertTrue(len(shard) % batch_size == 0)
|
||||
# All shards have the same number of samples
|
||||
self.assertEqual(len(shard), len(shard_lists[0]))
|
||||
|
||||
observed = []
|
||||
for idx in range(0, len(shard_lists[0]), batch_size):
|
||||
for shard in shard_lists:
|
||||
observed += shard[idx : idx + batch_size]
|
||||
|
||||
# If drop_last is False we loop through samples at the beginning to have a size that is a round multiple of
|
||||
# batch_size
|
||||
reference = copy.copy(dataset)
|
||||
if not drop_last:
|
||||
while len(reference) < len(observed):
|
||||
reference += reference
|
||||
self.assertListEqual(observed, reference[: len(observed)])
|
||||
|
||||
def test_shard_sampler(self):
|
||||
for n_elements in [64, 123]:
|
||||
dataset = list(range(n_elements))
|
||||
|
||||
self.check_shard_sampler(dataset, 4, drop_last=True, num_processes=2)
|
||||
self.check_shard_sampler(dataset, 4, drop_last=False, num_processes=2)
|
||||
|
||||
self.check_shard_sampler(dataset, 4, drop_last=True, num_processes=3)
|
||||
self.check_shard_sampler(dataset, 4, drop_last=False, num_processes=3)
|
||||
Reference in New Issue
Block a user