transformers/trainer.py stops after some iterations for iterative dataloaders.
Author: rabeehk
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 3.5.1
- Platform: linux
- Python version: 3.7
- PyTorch version (GPU?): yes 
- Tensorflow version (GPU?): - 
- Using GPU in script?: - 
- Using distributed or parallel set-up in script?: - 

### Who can help
 Trainer: @sgugger
 Text Generation: @patrickvonplaten @TevenLeScao
 T5: @patrickvonplaten
 examples/seq2seq: @patil-suraj
 examples/bert-loses-patience: @JetRunner

## Information
Hi
I am using dataloader of below, after first epoch it finishes and trainer does not continue with max_steps, could you point me to the issue? I set is_sized_dataset to False.  Thank you. 

```
class TaskDataLoader:
    """Wrapper around dataloader to keep the task names."""
    def __init__(self, task_name, dataset, batch_size=8,
                 collate_fn=None, drop_last=False, num_workers=0, sampler=None):
        self.dataset = dataset
        self.task_name = task_name
        self.data_loader = DataLoader(self.dataset,
                                      batch_size=batch_size,
                                      sampler=sampler,
                                      collate_fn=collate_fn,
                                      drop_last=drop_last,
                                      num_workers=num_workers)
    def __len__(self):
        return len(self.data_loader) #self.dataset.num_rows

    def __iter__(self):
        for batch in self.data_loader:
            yield batch



class MultiTaskDataLoader:
    """Given a dictionary of task: dataset, returns a multi-task dataloader
    which uses temperature sampling to sample different datasets."""

    def __init__(self,  tasks_to_datasets, batch_size=8, collate_fn=None,
                 drop_last=False, num_workers=0, temperature=100.0):

        # Computes a mapping from task to dataloaders.
        self.task_to_dataloaders = {}
        for task, dataset in tasks_to_datasets.items():
            dataloader = TaskDataLoader(task, dataset, batch_size,
                collate_fn=collate_fn, drop_last=drop_last, num_workers=num_workers)
            self.task_to_dataloaders.update({task: dataloader})
        self.tasknames = list(self.task_to_dataloaders.keys())

        # Computes the temperature sampling weights.
        self.sampling_weights = self.temperature_sampling(self.dataloader_sizes.values(), temperature)
        self.dataiters = {k: cycle(v) for k, v in self.task_to_dataloaders.items()}

    def temperature_sampling(self, dataset_sizes, temp):
        total_size = sum(dataset_sizes)
        weights = np.array([(size / total_size) ** (1.0 / temp) for size in dataset_sizes])
        return weights/np.sum(weights)

    @property
    def dataloader_sizes(self):
        if not hasattr(self, '_dataloader_sizes'):
            self._dataloader_sizes = {k: len(v) for k, v in self.task_to_dataloaders.items()}
        return self._dataloader_sizes

    def __len__(self):
        return sum(v for k, v in self.dataloader_sizes.items())

    def __iter__(self):
        outputs = {}
        for i in range(len(self)):
            taskname = np.random.choice(self.tasknames, p=self.sampling_weights)
            dataiter = self.dataiters[taskname]
            outputs["batch"] = next(dataiter)
            outputs["task"] = taskname
            yield outputs



class Trainer():
    """This is the trainer class which is responsible for distributing the data
    in case of multiple TPUs/GPUs."""
    def __init__(self, dataset_names_to_datasets):
        self.dataset_names_to_datasets = dataset_names_to_datasets
        self.batch_size = 8
        self.local_rank = -1 # this is not -1 in case of multi-gpu
        self.collate_fn = None
        self.drop_last = False
        self.num_workers = 0

    def get_sharded_data(self, num_replicas, rank):
        """Returns the sharded data belonging to the given rank."""
        sharded_dataset_names_to_datasets = {}
        for dataset_name, dataset in self.dataset_names_to_datasets:
            sharded_data = dataset.shard(num_replicas, rank)
            sharded_dataset_names_to_datasets.update({dataset_name: sharded_data})
        return sharded_dataset_names_to_datasets


    def get_train_dataset_shards(self):
        """In case of multiprocessing, returns the sharded data for the given rank."""
        if is_torch_tpu_available():
            if xm.xrt_world_size() > 1:
                return self.get_sharded_data(num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())
        elif self.local_rank != -1:
                return self.get_sharded_data(num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())
        else:
            return self.dataset_names_to_datasets


    def get_train_dataloader(self):
        """Returns the multi-task dataloader, each batch belongs
        to one task dataset."""
        dataset_names_to_datasets = self.get_train_dataset_shards()
        dataloader = MultiTaskDataLoader(dataset_names_to_datasets,
                                         batch_size=self.batch_size,
                                         collate_fn=self.collate_fn,
                                         drop_last=self.drop_last,
                                         num_workers=self.num_workers)
        return dataloader

```


 

Comments:
could you tell me please how should be the format of iterative dataloaders for the trainer funtion? I mean in the current implementation, this just goes till the end of length of dataloader and then it terminates, it does not loop again, could you explain how I can use trainer.py with iterative dataloaders please? thanks 

It's hard to know what's going on without seeing the command/script you are executing. In particular, the `Trainer` logs a lot of info regarding the number of steps/epochs at the beginning of training that could be useful to debug this.

Hi @sgugger  I spent really the whole day long hours continously, cannot see what is going on, this really needs someone of more expertise. this is hard for me to see  the reason. To me callbaxks are changing the dataloader but not sure where it is happening. 


@sgugger I added the codes in https://github.com/rabeehk/debug , here is how to run:
```
pip install -r requirements.txt
python setup.py develop
cd seq2seq 
python finetune_t5_trainer.py configs/mrpc_adapter_local.json
```
The results of running the codes is that after epoch 1 the dataloader is no more called resulting in labels not being inside the batch. I made the test case small to run fast, could you have a look, this is my only hope to fix this issue. thank you 

```
### epochs_trained, num_train_epochs  0 4000
#### epoch  0
step  0   dict_keys(['input_ids', 'attention_mask', 'decoder_input_ids', 'labels'])
### in the loss  dict_keys(['input_ids', 'attention_mask', 'decoder_input_ids', 'labels'])
@@@ after
#### epoch  1
step  0   dict_keys(['input_ids', 'attention_mask', 'decoder_input_ids'])
### in the loss  dict_keys(['input_ids', 'attention_mask', 'decoder_input_ids'])
Traceback (most recent call last):
  File "finetune_t5_trainer.py", line 250, in <module>
    main()
  File "finetune_t5_trainer.py", line 183, in main
    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
  File "/remote/idiap.svm/user.active/rkarimi/dev/internship/seq2seq/trainers/trainer.py", line 789, in train
    tr_loss += self.training_step(model, inputs)
  File "/remote/idiap.svm/user.active/rkarimi/dev/internship/seq2seq/trainers/trainer.py", line 1141, in training_step
    loss = self.compute_loss(model, inputs)
  File "/remote/idiap.svm/user.active/rkarimi/dev/internship/seq2seq/trainers/t5_trainer.py", line 339, in compute_loss
    labels = inputs.pop("labels")
KeyError: 'labels'
  0%| 
```




Hi
here is the solution, the way trainer works is having 1 iterative datasets with max_steps, then the issue was that cycle has caching in memory under the hood, then after first epoch, when modifying inputs, this was using it in the next iter and was craching, fixed with defining cycle as:

```
def cycle(iterable):
    while True:
        for x in iterable:
            yield x
```

and iterating over max_steps in the MultiTaskDataLoader.

I'm really sorry to bother you. Could you please tell me how to modify it specifically? I have also encountered this problem.