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 71962d9..9ef6331 100644 Binary files a/norch/__pycache__/__init__.cpython-38.pyc and b/norch/__pycache__/__init__.cpython-38.pyc differ 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 5b7fd1c..9af1d40 100644 Binary files a/norch/nn/__pycache__/__init__.cpython-38.pyc and b/norch/nn/__pycache__/__init__.cpython-38.pyc differ diff --git a/norch/nn/__pycache__/parameter.cpython-38.pyc b/norch/nn/__pycache__/parameter.cpython-38.pyc index 10a5bf3..fc69338 100644 Binary files a/norch/nn/__pycache__/parameter.cpython-38.pyc and b/norch/nn/__pycache__/parameter.cpython-38.pyc differ 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