Spaces:
Sleeping
Sleeping
File size: 4,567 Bytes
add142e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import numpy as np
from torch.utils.data import DataLoader
from torch.utils.data.sampler import Sampler
from torch.utils.data.distributed import DistributedSampler
from util.torch import worker_init_fn
class OrderedDistributedSampler(Sampler):
"""
Sampler that orders samples in a specified order and assigns them to different processes.
Attributes:
dataset (Dataset): The dataset to sample from.
num_replicas (int): The total number of processes participating in the distributed training.
rank (int): The rank of the current process.
num_samples (int): The number of samples per process.
total_size (int): The total number of samples across all processes.
Methods:
__iter__(): Returns an iterator over the indices of the samples for the current process.
__len__(): Returns the number of samples per process.
"""
def __init__(self, dataset, num_replicas=None, rank=None):
"""
Constructor.
Args:
dataset (Dataset): The dataset to sample from.
num_replicas (int): Number of processes in the distributed training. Defaults to None.
rank (int): The rank of the current process. Defaults to None.
Raises:
AssertionError: If num_replicas or rank is None.
"""
assert num_replicas is not None
assert rank is not None
self.dataset = dataset
self.num_replicas = num_replicas
self.rank = rank
self.num_samples = int(np.ceil(len(self.dataset) * 1.0 / self.num_replicas))
self.total_size = self.num_samples * self.num_replicas
def __iter__(self):
"""
Returns an iterator over the indices of the samples to be used in the current process.
"""
indices = list(range(len(self.dataset)))
# add extra samples to make it evenly divisible
indices += indices[: (self.total_size - len(indices))]
assert len(indices) == self.total_size
# subsample
indices = indices[
self.rank * self.num_samples: self.rank * self.num_samples
+ self.num_samples
]
assert len(indices) == self.num_samples
return iter(indices)
def __len__(self):
"""
Returns the number of samples per process.
"""
return self.num_samples
def define_loaders(
train_dataset,
val_dataset,
batch_size=32,
val_bs=32,
distributed=False,
world_size=0,
local_rank=0,
num_workers=0,
):
"""
Define data loaders for training and validation datasets.
If `distributed` is True, the data loaders will use DistributedSampler for shuffling the
training dataset and OrderedDistributedSampler for sampling the validation dataset.
Args:
train_dataset (Dataset): The training dataset.
val_dataset (Dataset): The validation dataset.
batch_size (int): The batch size for training data loader. Default to 32.
val_bs (int): The batch size for validation data loader. Default to 32.
distributed (bool): Whether to use distributed training. Default to False.
world_size (int): The total number of processes for distributed training. Default to 0.
local_rank (int): The rank of the current process. Default to 0.
num_workers (int): Number of workers to use for the dataloaders. Default to 0.
Returns:
tuple: A tuple containing the training data loader and the validation data loader.
"""
sampler, val_sampler = None, None
if distributed:
sampler = DistributedSampler(
train_dataset,
num_replicas=world_size,
rank=local_rank,
shuffle=True,
seed=world_size + local_rank,
)
val_sampler = OrderedDistributedSampler(
val_dataset, num_replicas=world_size, rank=local_rank
)
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
sampler=sampler,
shuffle=sampler is None,
drop_last=True,
num_workers=num_workers,
pin_memory=True,
worker_init_fn=worker_init_fn,
collate_fn=None,
persistent_workers=num_workers > 0,
)
val_loader = DataLoader(
val_dataset,
batch_size=val_bs,
sampler=val_sampler,
shuffle=False,
num_workers=num_workers,
pin_memory=True,
collate_fn=None,
persistent_workers=num_workers > 0,
)
return train_loader, val_loader
|