| from __future__ import division |
| import torch |
| from torch.utils.data import DataLoader |
| from torch.utils.data.sampler import Sampler |
|
|
| class RandomSampler(Sampler): |
|
|
| def __init__(self, data_source, checkpoint): |
| self.data_source = data_source |
| if checkpoint is not None and checkpoint['dataset_perm'] is not None: |
| self.dataset_perm = checkpoint['dataset_perm'] |
| self.perm = self.dataset_perm[checkpoint['batch_size']*checkpoint['batch_idx']:] |
| else: |
| self.dataset_perm = torch.randperm(len(self.data_source)).tolist() |
| self.perm = torch.randperm(len(self.data_source)).tolist() |
|
|
| def __iter__(self): |
| return iter(self.perm) |
| |
| def __len__(self): |
| return len(self.perm) |
|
|
| class SequentialSampler(Sampler): |
|
|
| def __init__(self, data_source, checkpoint): |
| self.data_source = data_source |
| if checkpoint is not None and checkpoint['dataset_perm'] is not None: |
| self.dataset_perm = checkpoint['dataset_perm'] |
| self.perm = self.dataset_perm[checkpoint['batch_size']*checkpoint['batch_idx']:] |
| else: |
| self.dataset_perm = list(range(len(self.data_source))) |
| self.perm = self.dataset_perm |
|
|
| def __iter__(self): |
| return iter(self.perm) |
| |
| def __len__(self): |
| return len(self.perm) |
|
|
| class CheckpointDataLoader(DataLoader): |
| """ |
| Extends torch.utils.data.DataLoader to handle resuming training from an arbitrary point within an epoch. |
| """ |
| def __init__(self, dataset, checkpoint=None, batch_size=1, |
| shuffle=False, num_workers=0, pin_memory=False, drop_last=True, |
| timeout=0, worker_init_fn=None): |
|
|
| if shuffle: |
| sampler = RandomSampler(dataset, checkpoint) |
| else: |
| sampler = SequentialSampler(dataset, checkpoint) |
| if checkpoint is not None: |
| self.checkpoint_batch_idx = checkpoint['batch_idx'] |
| else: |
| self.checkpoint_batch_idx = 0 |
|
|
| super(CheckpointDataLoader, self).__init__(dataset, sampler=sampler, shuffle=False, batch_size=batch_size, num_workers=num_workers, |
| drop_last=drop_last, pin_memory=pin_memory, timeout=timeout, worker_init_fn=None) |
|
|
