Repository: pytorch
Issue: state_dict and load_state_dict methods for DataLoader and Sampler to continue training at specific epoch and batch
Labels: feature,module: dataloader,triaged

## 🚀 Feature
Based on training checkpoint data, allow training to restart at the epoch and **precise** batch, after the last batch trained on when the training checkpoint data was saved.

## Motivation

I'm training (relatively) large sequence to sequence models (> 100M parameters) with a (relatively large) dataset (1M conversations), using limited training hardware (one machine with 4x 1080Ti); one epoch can easily take 12 hours or (much) longer.

There are situations where I want to restart training based on a training checkpoint I saved during the training process that was stopped before, for some reason. In order to not lose many hours of training time, and to continue in exactly the same way, as if I would if the training process had not stopped, I want to be able to continue at the epoch and **precise** batch, after the last batch trained on when the training checkpoint data was saved.

## Pitch

All relevant components in of a PyTorch training process, such as `nn.Module`s, optimizers and schedulers have `state_dict` and `load_state_dict` methods in order to retrieve and load its states, such that training can be restarted with the exactly the same state as it was stopped.

The only exception, as far as I can tell, is for `DataLoader` and `Sampler`. These (base) classes do not have `state_dict` and `load_state_dict` methods.

It would be very consistent and concise, if a PyTorch user could do the following in order to save and load training state:

1. Get the state of your `DataLoader`s, models, optimizers and schedulers and save it together with any other training process variables (e.g. current epoch and batch index) as a training checkpoint.

2. Load your training checkpoint and subsequently load the state of `DataLoader`s, models, optimizers and schedulers

3. Restart your training loop at the loaded epoch index and batch index and continue training

Please note that it is assumed that the `Sampler` state is nested as a part of the `DataLoader` state.

## Alternatives

`DistributedSampler` has the possibility to set the current epoch, such that the random shuffling of samples is deterministic. However, I can't tell the DataLoader from which specific batch to continue.

Edit:
## Additional thoughts

Assume that `k` of `N` batches of an epoch have been trained on, before training was stopped and a training checkpoint was saved. When the state of a `DataLoader` is reloaded in to an `DataLoader` instance `my_data_loader`, to stay consistent, `len(my_data_loader)` should be the total number of batches in an epoch, `N`.

```
my_sampler = SomeSampler(...)
my_data_loader = torch.utils.data.DataLoader(my_dataset, ... sampler=my_sampler ...) 

my_data_loader.load_state_dict(checkpoint['data_loader'])

# Should print the number of batches in a normal epoch, N 
print(len(my_data_loader))
```

However, when iterating over batches using the first `iter` object after restarting training:
```
for batch in iter(my_data_loader):
    ...
```
The number of iterations will be `N-k`, for the last `N-k` batches in the epoch from where we stopped training and saved our training checkpoint. In the next epoch `iter(my_data_loader)` will again yield in the normal `N` batches.

cc @SsnL