| 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 |
| ) |
|
|
