From e590f57555f731de56d79f8f7621d507905a39de Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Sat, 25 May 2024 10:53:09 -0300 Subject: [PATCH] implemented distributed sampler --- norch/__init__.py | 2 +- norch/__pycache__/__init__.cpython-38.pyc | Bin 365 -> 356 bytes norch/distributed/distributed.py | 6 ++++ norch/nn/__pycache__/__init__.cpython-38.pyc | Bin 241 -> 261 bytes norch/nn/__pycache__/parameter.cpython-38.pyc | Bin 728 -> 736 bytes norch/nn/parameter.py | 4 +-- norch/utils/__init__.py | 4 +-- norch/utils/data/dataloader.py | 18 ++++++----- norch/utils/data/dataset.py | 2 +- norch/utils/data/distributed.py | 13 ++++---- norch/utils/{utils.py => functions.py} | 0 train.py | 29 +++++++++++++++++- 12 files changed, 56 insertions(+), 22 deletions(-) rename norch/utils/{utils.py => functions.py} (100%) diff --git a/norch/__init__.py b/norch/__init__.py index 80da58b..2aae8f0 100644 --- a/norch/__init__.py +++ b/norch/__init__.py @@ -3,7 +3,7 @@ from .nn import * from .optim import * from .utils import * from .norchvision import * -from . import data +from .utils import * __version__ = "0.0.4" __author__ = 'Lucas de Lima Nogueira' diff --git a/norch/__pycache__/__init__.cpython-38.pyc b/norch/__pycache__/__init__.cpython-38.pyc index 71962d9e9e29bd7e443f82302b507406eb8b5080..9ef633144738bc5135fd57fbb835823ece228f9b 100644 GIT binary patch delta 180 zcmaFM^n{5wl$V!_0SMY32By85$SWJ_0OX`Fq%h_%gfOHq0cl1k%><>Hp)?B;pE-&( zg(Zp&h}lzEgBdj0UV_ZjWW2?w^^y_Dn%JtS#1X}lmtT~ep;wZcSDarI#gvy9#hPDG pl9_vpwX`HNr});yv!-n9K-ER;lQ|fJMFfCM4n{a+Vqs)q1OU6HCH^ns~^RjSVPO#5VD7un-@R#ly$}flMrnEQ|ou C-zsVV diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py index bd87bc3..98bc2a0 100644 --- a/norch/distributed/distributed.py +++ b/norch/distributed/distributed.py @@ -9,6 +9,12 @@ def init_process_group(rank, world_size, backend='nccl'): Tensor._C.init_process_group(rank, world_size) +def get_rank(): + return os.getenv('OMPI_COMM_WORLD_RANK', 0) + +def get_world_size(): + return os.getenv('OMPI_COMM_WORLD_SIZE', 1) + def broadcast_tensor(tensor, src=0): Tensor._C.broadcast_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int] diff --git a/norch/nn/__pycache__/__init__.cpython-38.pyc b/norch/nn/__pycache__/__init__.cpython-38.pyc index 5b7fd1c7b09471efcfce61414f7fe60ec9c05c2f..9af1d40ac65a2aab524d9d068a81346d9cc372be 100644 GIT binary patch delta 82 zcmey!*viBk%FD~e00cdk0@KPS^2##mOw?A9Vn|`kVF_VKVG3r@WPZsAl-Fdu#i-?{ d$vQDnoimD~Ah9ShCnt4cdjby!P?CcY0s;M45hVZs delta 62 zcmZo=`pC!|%FD~e00g`~-f3|Yd1V<@CTgn)F$FVdGQVU5ifc07V$|}}WSN+$K5pF diff --git a/norch/nn/parameter.py b/norch/nn/parameter.py index 2324a1a..5a29f3f 100644 --- a/norch/nn/parameter.py +++ b/norch/nn/parameter.py @@ -1,5 +1,5 @@ from norch.tensor import Tensor -from norch.utils import utils +from norch.utils import functions import random class Parameter(Tensor): @@ -7,5 +7,5 @@ class Parameter(Tensor): A parameter is a trainable tensor. """ def __init__(self, shape): - data = utils.generate_random_list(shape=shape) + data = functions.generate_random_list(shape=shape) super().__init__(data, requires_grad=True) \ No newline at end of file diff --git a/norch/utils/__init__.py b/norch/utils/__init__.py index c6adfd1..b6b5a68 100644 --- a/norch/utils/__init__.py +++ b/norch/utils/__init__.py @@ -1,2 +1,2 @@ -from .utils import * -from .data import * \ No newline at end of file +from .data import * +from .functions import * \ No newline at end of file diff --git a/norch/utils/data/dataloader.py b/norch/utils/data/dataloader.py index 74a2800..13c02c5 100644 --- a/norch/utils/data/dataloader.py +++ b/norch/utils/data/dataloader.py @@ -11,17 +11,19 @@ class Dataloader: def __iter__(self): if self.sampler is not None: - indices = list(self.sampler) + indices = iter(self.sampler) + else: - indices = np.arange(len(self.dataset)) - - for start in range(0, len(indices), self.batch_size): - end = start + self.batch_size - batch_indices = indices[start:end] - batch_size = len(batch_indices) - yield Batch([self.dataset[i] for i in batch_indices], batch_size) + indices = range(len(self.dataset)) + + for idx in indices: + start = idx * self.batch_size + end = min(start + self.batch_size, len(self.dataset)) + yield Batch(self.dataset[start:end], end - start) + def __len__(self): if self.sampler is not None: return len(self.sampler) // self.batch_size + return len(self.dataset) // self.batch_size diff --git a/norch/utils/data/dataset.py b/norch/utils/data/dataset.py index d38c1ea..2130d52 100644 --- a/norch/utils/data/dataset.py +++ b/norch/utils/data/dataset.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod import os -from norch.utils import extract_to_dir, download_from_url +from norch.utils.functions import extract_to_dir, download_from_url from .example import Example import norch diff --git a/norch/utils/data/distributed.py b/norch/utils/data/distributed.py index 6af54bc..c83dfc4 100644 --- a/norch/utils/data/distributed.py +++ b/norch/utils/data/distributed.py @@ -1,5 +1,5 @@ -import math import numpy as np +import math class DistributedSampler: def __init__(self, dataset, num_replicas, rank): @@ -10,18 +10,17 @@ class DistributedSampler: self.total_size = self.num_samples * self.num_replicas def __iter__(self): - # Create indices for the dataset - indices = np.arange(len(self.dataset)) + indices = list(range(len(self.dataset))) # Add extra samples to make it evenly divisible - indices = np.concatenate([indices, indices[:(self.total_size - len(indices)) % len(indices)]]) + indices = indices[:self.total_size] assert len(indices) == self.total_size - # Subsample + # subsample indices = indices[self.rank:self.total_size:self.num_replicas] assert len(indices) == self.num_samples - + return iter(indices) def __len__(self): - return self.num_samples \ No newline at end of file + return self.num_samples diff --git a/norch/utils/utils.py b/norch/utils/functions.py similarity index 100% rename from norch/utils/utils.py rename to norch/utils/functions.py diff --git a/train.py b/train.py index 8f07a41..7f9a613 100644 --- a/train.py +++ b/train.py @@ -1,4 +1,4 @@ -import os +"""import os import norch import norch.distributed as dist @@ -29,3 +29,30 @@ def main(): if __name__ == "__main__": main() + +""" +import norch +import matplotlib.pyplot as plt +import numpy as np +import random + +train_data, test_data = norch.norchvision.datasets.MNIST.splits() +train_sampler = norch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=10, rank=2) +train_loader = norch.utils.data.Dataloader(train_data, batch_size = 1, sampler=train_sampler) +input_sample, target_sample = train_data[0] + +fig = plt.figure(figsize = (20, 10)) +columns = 4 +rows = 2 + +# Choose a random image +for image_index, batch in enumerate(train_loader): + + image, label = batch + fig.add_subplot(rows, columns, image_index+1) + plt.imshow(np.array(image).reshape(28, 28)) + plt.title(label) + plt.axis('off') + if image_index > 6: + break +plt.show() \ No newline at end of file