Repository: huggingface/transformers (Python) — Bug Fix

Issue:
Trainer: exception raised when calling len() on IterableDataset
# 🐛 Bug

## Information
While pre-training a Longformer model from scratch, the text is delivered through an `IterableDataset` object. The code which is called by `Trainer.train()` still calls `len()` on this object, which raises an exception.
#5829 addressed the proper creation of the Dataloader.

The problem arises when using:
* [x] my own modified scripts: see code

The tasks I am working on is:
* [x] my own task or dataset: pre-train a LM from scratch

## To reproduce
Here is my entire code, but it can be reproduced with any `PreTrainedModel` by using an `IterableDataset`.

```python
import logging
import random

from dataclasses import dataclass, field

from transformers import LongformerConfig, LongformerForMaskedLM, LongformerTokenizerFast
from transformers import Trainer, TrainingArguments
from transformers import TextDataset, DataCollatorForLanguageModeling
from transformers import HfArgumentParser

from sklearn.model_selection import train_test_split

from pathlib import Path

from utils_pretrain import MultiTextDataset

logger = logging.getLogger(__name__)





@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    max_seq_len: int = field(
        metadata={"help": "Input Sequence Length"}
    )
    num_hidden_layers: int = field(
        metadata={'help': 'Number of transformer layers in Longformer'}
    )
    tok_dir: str = field(
        metadata={
            'help': 'Folder with tokenizer files'
        }
    )
    txt_dir: str = field(
        metadata={"help": "Folder with txt files for tokenizer training"}
    )
    filter_files: str = field(
        default='[a-c]*.txt',
        metadata={"help": "regex to select specific files"}
    )
    test_size: float = field(
        default=0.05,
        metadata={'help': 'proportion of the data that will be used for evaluation'} 
    )


def main():
    parser = HfArgumentParser((ModelArguments, TrainingArguments))
    model_args, train_args = parser.parse_args_into_dataclasses()
    model_args: ModelArguments

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
        train_args.local_rank,
        train_args.device,
        train_args.n_gpu,
        bool(train_args.local_rank != -1),
        train_args.fp16,
    )
    logger.info("Training/evaluation parameters %s", train_args)

    MODEL_NAME = 'allenai/longformer-base-4096'
    tokenizer: LongformerTokenizerFast = LongformerTokenizerFast.from_pretrained(model_args.tok_dir)

    # Customize an existing config rather than create from scratch
    config: LongformerConfig = LongformerConfig.from_pretrained(MODEL_NAME)
    config.max_position_embeddings = model_args.max_seq_len + 2
    config.num_hidden_layers = model_args.num_hidden_layers
    config.attention_window = [512] * model_args.num_hidden_layers
    config.vocab_size = tokenizer.vocab_size
    model = LongformerForMaskedLM(config)

    data_files = list(Path(model_args.txt_dir).glob(model_args.filter_files))
    shuffled_files = random.sample(data_files, len(data_files))
    train_files, val_files = train_test_split(shuffled_files, test_size=model_args.test_size)

    train_ds, val_ds = list(
        map(
            lambda x: MultiTextDataset(
                files=x,
                tokenizer=tokenizer,
                block_size=model_args.max_seq_len
            ),
            [train_files, val_files]
        )
    )

    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=True,
        mlm_probability=0.15
    )

    train_args: TrainingArguments
    train_args.do_train = True
    train_args.evaluate_during_training = True

    trainer = Trainer(
        model=model,
        args=train_args,
        data_collator=data_collator,
        train_dataset=train_ds,
        eval_dataset=val_ds,
    )

    trainer.train(train_args.output_dir)
```

The class `MultiTextDataset` inherits `IterableDataset`. It has no `__len__` method, and the length would require the whole dataset to be parsed at once to be known.

Here is the exception and stack trace:
```
Traceback (most recent call last):
  File "longformer_pretrain.py", line 131, in <module>
    main()
  File "longformer_pretrain.py", line 122, in main
    trainer.train(train_args.output_dir)
  File "/home/jrossi/anaconda3/envs/COLIEE/lib/python3.7/site-packages/transformers/trainer.py", line 392, in train
    self.args.max_steps // (len(train_dataloader) // self.args.gradient_accumulation_steps) + 1
  File "/home/jrossi/anaconda3/envs/COLIEE/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 313, in __len__
    length = self._IterableDataset_len_called = len(self.dataset)
TypeError: object of type 'MultiTextDataset' has no len()
```

## Expected behavior

The call to `Trainer.train()` starts the training. A case has to be made in the code to accomodate the usage of `IterableDataset`, which means not assuming that `len()` can be called on the dataset at any point.

- If a number of epochs is given, one epoch corresponds to consuming the iterable dataset until StopIteration
- If a number of steps is given, training stops after performing MAX_STEPS or catching a StopIteration, whichever comes first
- During training, the progress bar should be either a % of epochs performed, or a % of steps performed
- (optional) If a number of epochs is given, register how many steps it took to consume the iterator so a better progress bar can be shown for the next epochs (each epoch will consume the same iterator once)

With regards to [Pytorch documentation](https://pytorch.org/docs/stable/data.html#), there is no certainty that `__len__` method will be implemented, even on `Dataset` objects. 
The distinction should be made between objects that implement `__len__` and those that do not implement it.
The current code __assumes__ that the `Dataset` objects given when creating a `Trainer` implement `len()`, but there is no guarantee of this.

```python
import collections
if isinstance(bar, collections.Sized): (...)
```

## Environment info

- `transformers` version: 3.0.2
- Platform: Linux-5.7.8-1.el7.elrepo.x86_64-x86_64-with-centos-7.8.2003-Core
- Python version: 3.7.7
- PyTorch version (GPU?): 1.5.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: YES
- Using distributed or parallel set-up in script?: NO (for the moment)

## Fix
I can contribute. I will suggest a PR to fix this.


Discussion:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Resolving patch:
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -16,7 +16,9 @@
 The Trainer class, to easily train a 🤗 Transformers from scratch or finetune it on a new task.
 """
 
+import collections
 import inspect
+import math
 import os
 import re
 import shutil
@@ -283,6 +285,15 @@ def __init__(
                 FutureWarning,
             )
 
+        if args.max_steps > 0:
+            logger.info("max_steps is given, it will override any value given in num_train_epochs")
+
+        # Enforce rules on using datasets with no __len__
+        if train_dataset is not None and not isinstance(train_dataset, collections.abc.Sized) and args.max_steps <= 0:
+            raise ValueError("train_dataset does not implement __len__, max_steps has to be specified")
+        if eval_dataset is not None and not isinstance(eval_dataset, collections.abc.Sized):
+            raise ValueError("eval_dataset must implement __len__")
+
         if is_datasets_available():
             if isinstance(train_dataset, datasets.Dataset):
                 self._remove_unused_columns(self.train_dataset, description="training")
@@ -361,7 +372,7 @@ def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optio
         dataset.set_format(type=dataset.format["type"], columns=columns)
 
     def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
-        if isinstance(self.train_dataset, torch.utils.data.IterableDataset):
+        if not isinstance(self.train_dataset, collections.abc.Sized):
             return None
         elif is_torch_tpu_available():
             return get_tpu_sampler(self.train_dataset)
@@ -376,7 +387,7 @@ def get_train_dataloader(self) -> DataLoader:
         """
         Returns the training :class:`~torch.utils.data.DataLoader`.
 
-        Will use no sampler if :obj:`self.train_dataset` is a :obj:`torch.utils.data.IterableDataset`, a random sampler
+        Will use no sampler if :obj:`self.train_dataset` does not implement :obj:`__len__`, a random sampler
         (adapted to distributed training if necessary) otherwise.
 
         Subclass and override this method if you want to inject some custom behavior.
@@ -395,9 +406,7 @@ def get_train_dataloader(self) -> DataLoader:
         )
 
     def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.sampler.Sampler]:
-        if isinstance(eval_dataset, torch.utils.data.IterableDataset):
-            return None
-        elif is_torch_tpu_available():
+        if is_torch_tpu_available():
             return SequentialDistributedSampler(eval_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())
         elif self.args.local_rank != -1:
             return SequentialDistributedSampler(eval_dataset)
@@ -408,19 +417,18 @@ def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoa
         """
         Returns the evaluation :class:`~torch.utils.data.DataLoader`.
 
-        Will use no sampler if :obj:`self.eval_dataset` is a :obj:`torch.utils.data.IterableDataset`, a sequential
-        sampler (adapted to distributed training if necessary) otherwise.
-
         Subclass and override this method if you want to inject some custom behavior.
 
         Args:
             eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
                 If provided, will override :obj:`self.eval_dataset`. If it is an :obj:`datasets.Dataset`, columns not
-                accepted by the ``model.forward()`` method are automatically removed.
+                accepted by the ``model.forward()`` method are automatically removed. It must implement :obj:`__len__`.
         """
         if eval_dataset is None and self.eval_dataset is None:
             raise ValueError("Trainer: evaluation requires an eval_dataset.")
-        elif eval_dataset is not None and is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
+        elif eval_dataset is not None and not isinstance(eval_dataset, collections.abc.Sized):
+            raise ValueError("eval_dataset must implement __len__")
+        elif is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
             self._remove_unused_columns(eval_dataset, description="evaluation")
         eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
         eval_sampler = self._get_eval_sampler(eval_dataset)
@@ -438,17 +446,16 @@ def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
         """
         Returns the test :class:`~torch.utils.data.DataLoader`.
 
-        Will use no sampler if :obj:`test_dataset` is a :obj:`torch.utils.data.IterableDataset`, a sequential
-        sampler (adapted to distributed training if necessary) otherwise.
-
         Subclass and override this method if you want to inject some custom behavior.
 
         Args:
-            eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
+            test_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
                 The test dataset to use. If it is an :obj:`datasets.Dataset`, columns not accepted by the
-                ``model.forward()`` method are automatically removed.
+                ``model.forward()`` method are automatically removed. It must implement :obj:`__len__`.
         """
-        if is_datasets_available() and isinstance(test_dataset, datasets.Dataset):
+        if not isinstance(test_dataset, collections.abc.Sized):
+            raise ValueError("test_dataset must implement __len__")
+        elif is_datasets_available() and isinstance(test_dataset, datasets.Dataset):
             self._remove_unused_columns(test_dataset, description="test")
         test_sampler = self._get_eval_sampler(test_dataset)
 
@@ -494,6 +501,8 @@ def create_optimizer_and_scheduler(self, num_training_steps: int):
     def num_examples(self, dataloader: DataLoader) -> int:
         """
         Helper to get number of samples in a :class:`~torch.utils.data.DataLoader` by accessing its dataset.
+
+        Will raise an exception if the underlying dataset dese not implement method :obj:`__len__`
         """
         return len(dataloader.dataset)
 
@@ -579,19 +588,32 @@ def train(self, model_path: Optional[str] = None, trial: Union["optuna.Trial", D
             # Reinitializes optimizer and scheduler
             self.optimizer, self.lr_scheduler = None, None
 
+        # Keeping track whether we can can len() on the dataset or not
+        train_dataset_is_sized = isinstance(self.train_dataset, collections.abc.Sized)
+
         # Data loader and number of training steps
         train_dataloader = self.get_train_dataloader()
-        num_update_steps_per_epoch = len(train_dataloader) // self.args.gradient_accumulation_steps
-        num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
-        if self.args.max_steps > 0:
-            max_steps = self.args.max_steps
-            num_train_epochs = self.args.max_steps // num_update_steps_per_epoch + int(
-                self.args.max_steps % num_update_steps_per_epoch > 0
-            )
+
+        # Setting up training control variables:
+        # number of training epochs: num_train_epochs
+        # number of training steps per epoch: num_update_steps_per_epoch
+        # total number of training steps to execute: max_steps
+        if train_dataset_is_sized:
+            num_update_steps_per_epoch = len(train_dataloader) // self.args.gradient_accumulation_steps
+            num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
+            if self.args.max_steps > 0:
+                max_steps = self.args.max_steps
+                num_train_epochs = self.args.max_steps // num_update_steps_per_epoch + int(
+                    self.args.max_steps % num_update_steps_per_epoch > 0
+                )
+            else:
+                max_steps = math.ceil(self.args.num_train_epochs * num_update_steps_per_epoch)
+                num_train_epochs = math.ceil(self.args.num_train_epochs)
         else:
-            max_steps = int(num_update_steps_per_epoch * self.args.num_train_epochs)
-            num_train_epochs = self.args.num_train_epochs
-        num_train_epochs = int(np.ceil(num_train_epochs))
+            # see __init__. max_steps is set when the dataset has no __len__
+            max_steps = self.args.max_steps
+            num_train_epochs = 1
+            num_update_steps_per_epoch = max_steps
 
         self.create_optimizer_and_scheduler(num_training_steps=max_steps)
         self.state = TrainerState()
@@ -645,8 +667,15 @@ def train(self, model_path: Optional[str] = None, trial: Union["optuna.Trial", D
                 * self.args.gradient_accumulation_steps
                 * (torch.distributed.get_world_size() if self.args.local_rank != -1 else 1)
             )
+
+        num_examples = (
+            self.num_examples(train_dataloader)
+            if train_dataset_is_sized
+            else total_train_batch_size * self.args.max_steps
+        )
+
         logger.info("***** Running training *****")
-        logger.info("  Num examples = %d", self.num_examples(train_dataloader))
+        logger.info("  Num examples = %d", num_examples)
         logger.info("  Num Epochs = %d", num_train_epochs)
         logger.info("  Instantaneous batch size per device = %d", self.args.per_device_train_batch_size)
         logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d", total_train_batch_size)
@@ -703,6 +732,7 @@ def train(self, model_path: Optional[str] = None, trial: Union["optuna.Trial", D
             if self.args.past_index >= 0:
                 self._past = None
 
+            steps_in_epoch = len(epoch_iterator) if train_dataset_is_sized else self.args.max_steps
             self.control = self.callback_handler.on_epoch_begin(self.args, self.state, self.control)
 
             for step, inputs in enumerate(epoch_iterator):
@@ -728,8 +758,8 @@ def train(self, model_path: Optional[str] = None, trial: Union["optuna.Trial", D
 
                 if (step + 1) % self.args.gradient_accumulation_steps == 0 or (
                     # last step in epoch but step is always smaller than gradient_accumulation_steps
-                    len(epoch_iterator) <= self.args.gradient_accumulation_steps
-                    and (step + 1) == len(epoch_iterator)
+                    steps_in_epoch <= self.args.gradient_accumulation_steps
+                    and (step + 1) == steps_in_epoch
                 ):
                     if self.args.fp16 and _use_native_amp:
                         self.scaler.unscale_(self.optimizer)
@@ -750,7 +780,7 @@ def train(self, model_path: Optional[str] = None, trial: Union["optuna.Trial", D
                     self.lr_scheduler.step()
                     model.zero_grad()
                     self.state.global_step += 1
-                    self.state.epoch = epoch + (step + 1) / len(epoch_iterator)
+                    self.state.epoch = epoch + (step + 1) / steps_in_epoch
                     self.control = self.callback_handler.on_step_end(self.args, self.state, self.control)
 
                     self._maybe_log_save_evalute(tr_loss, model, trial, epoch)
@@ -1207,11 +1237,15 @@ def evaluate(self, eval_dataset: Optional[Dataset] = None) -> Dict[str, float]:
         Args:
             eval_dataset (:obj:`Dataset`, `optional`):
                 Pass a dataset if you wish to override :obj:`self.eval_dataset`. If it is an :obj:`datasets.Dataset`,
-                columns not accepted by the ``model.forward()`` method are automatically removed.
+                columns not accepted by the ``model.forward()`` method are automatically removed. It must implement
+                the :obj:`__len__` method.
 
         Returns:
             A dictionary containing the evaluation loss and the potential metrics computed from the predictions.
         """
+        if eval_dataset is not None and not isinstance(eval_dataset, collections.abc.Sized):
+            raise ValueError("eval_dataset must implement __len__")
+
         eval_dataloader = self.get_eval_dataloader(eval_dataset)
 
         output = self.prediction_loop(eval_dataloader, description="Evaluation")
@@ -1234,7 +1268,7 @@ def predict(self, test_dataset: Dataset) -> PredictionOutput:
         Args:
             test_dataset (:obj:`Dataset`):
                 Dataset to run the predictions on. If it is an :obj:`datasets.Dataset`, columns not accepted by the
-                ``model.forward()`` method are automatically removed.
+                ``model.forward()`` method are automatically removed. Has to implement the method :obj:`__len__`
 
         Returns:
             `NamedTuple`:
@@ -1245,6 +1279,9 @@ def predict(self, test_dataset: Dataset) -> PredictionOutput:
             metrics (:obj:`Dict[str, float]`, `optional`):
                 The potential dictionary of metrics (if the dataset contained labels).
         """
+        if test_dataset is not None and not isinstance(test_dataset, collections.abc.Sized):
+            raise ValueError("test_dataset must implement __len__")
+
         test_dataloader = self.get_test_dataloader(test_dataset)
 
         return self.prediction_loop(test_dataloader, description="Prediction")
@@ -1264,6 +1301,8 @@ def prediction_loop(
             )
             return self._prediction_loop(dataloader, description, prediction_loss_only=prediction_loss_only)
 
+        if not isinstance(dataloader.dataset, collections.abc.Sized):
+            raise ValueError("dataset must implement __len__")
         prediction_loss_only = (
             prediction_loss_only if prediction_loss_only is not None else self.args.prediction_loss_only
         )


Test patch:
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
old mode 100755
new mode 100644
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -31,11 +31,14 @@
     from torch.utils.data import IterableDataset
 
     from transformers import (
+        AutoModelForMaskedLM,
         AutoModelForSequenceClassification,
+        DataCollatorForLanguageModeling,
         GlueDataset,
         GlueDataTrainingArguments,
         LineByLineTextDataset,
         PreTrainedModel,
+        TextDataset,
         Trainer,
         TrainerState,
     )
@@ -83,15 +86,16 @@ def __init__(self, a=0, b=0, double_output=False, **kwargs):
 if is_torch_available():
 
     class SampleIterableDataset(IterableDataset):
-        def __init__(self, file_path):
-            self.file_path = file_path
+        """
+        Criteria is not whether it is IterableDataset or not, criteria is whether __len__ is implemented
+        """
 
-        def parse_file(self):
-            f = open(self.file_path, "r")
-            return f.readlines()
+        def __init__(self, file_path, tokenizer):
+            self.ds = TextDataset(file_path=file_path, tokenizer=tokenizer, block_size=64)
 
         def __iter__(self):
-            return iter(self.parse_file())
+            for i in range(len(self.ds)):
+                yield self.ds[i]
 
     class RegressionModel(torch.nn.Module):
         def __init__(self, a=0, b=0, double_output=False):
@@ -538,13 +542,51 @@ def test_trainer_eval_lm(self):
         self.assertEqual(len(dataset), 31)
 
     def test_trainer_iterable_dataset(self):
+        # Simulate Language Modeling with an IterableDataset, with no __len__ method
+        # Pick-up a tiny model, so it works on CPU
+        # See Issue #5990: https://github.com/huggingface/transformers/issues/5990
         MODEL_ID = "sshleifer/tiny-distilbert-base-cased"
-        model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
-        train_dataset = SampleIterableDataset(PATH_SAMPLE_TEXT)
-        training_args = TrainingArguments(output_dir="./examples", no_cuda=True)
-        trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset)
+        model = AutoModelForMaskedLM.from_pretrained(MODEL_ID)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        train_dataset = SampleIterableDataset(file_path=PATH_SAMPLE_TEXT, tokenizer=tokenizer)
+        training_args = TrainingArguments(output_dir="./examples", no_cuda=True, max_steps=2)
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.15)
+
+        training_args = TrainingArguments(output_dir="./examples", no_cuda=True, max_steps=2)
+        trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, data_collator=data_collator)
+        trainer.train()
+
         loader = trainer.get_train_dataloader()
         self.assertIsInstance(loader, torch.utils.data.DataLoader)
+        self.assertIsInstance(loader.sampler, torch.utils.data.dataloader._InfiniteConstantSampler)
+
+        # Exception if giving iterable dataset and no max_steps
+        with self.assertRaises(ValueError):
+            training_args = TrainingArguments(output_dir="./examples", no_cuda=True)
+            _ = Trainer(model=model, args=training_args, train_dataset=train_dataset, data_collator=data_collator)
+
+        # Exception if eval_dataset is iterable in __init__
+        with self.assertRaises(ValueError):
+            training_args = TrainingArguments(output_dir="./examples", no_cuda=True, max_steps=2)
+            _ = Trainer(
+                model=model,
+                args=training_args,
+                train_dataset=train_dataset,
+                eval_dataset=train_dataset,
+                data_collator=data_collator,
+            )
+
+        # Exception if predicting with iterable dataset
+        with self.assertRaises(ValueError):
+            training_args = TrainingArguments(output_dir="./examples", no_cuda=True)
+            trainer = Trainer(model=model, args=training_args, data_collator=data_collator)
+            trainer.predict(train_dataset)
+
+        # Exception if evaluating with iterable dataset
+        with self.assertRaises(ValueError):
+            training_args = TrainingArguments(output_dir="./examples", no_cuda=True)
+            trainer = Trainer(model=model, args=training_args, data_collator=data_collator)
+            trainer.evaluate(train_dataset)
 
     def test_num_train_epochs_in_training(self):
         # len(train_dl) < gradient_accumulation_steps shouldn't give ``ZeroDivisionError`` when ``max_steps`` is given.