diff --git a/README.md b/README.md index dd310b5..5ed5ebd 100644 --- a/README.md +++ b/README.md @@ -65,7 +65,7 @@ class MyModel(nn.Module): ### 3.3 - Example training ```python import norch -from norch.utils.data.dataloader import Dataloader +from norch.utils.data.dataloader import DataLoader from norch.norchvision import transforms import norch import norch.nn as nn @@ -77,21 +77,21 @@ BATCH_SIZE = 32 device = "cuda" #cpu epochs = 10 -transform = transforms.Sequential( +transform = transforms.Compose( [ transforms.ToTensor(), transforms.Reshape([-1, 784, 1]) ] ) -target_transform = transforms.Sequential( +target_transform = transforms.Compose( [ transforms.ToTensor() ] ) train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform) -train_loader = Dataloader(train_data, batch_size = BATCH_SIZE) +train_loader = DataLoader(train_data, batch_size = BATCH_SIZE) class MyModel(nn.Module): def __init__(self): @@ -146,4 +146,4 @@ for epoch in range(epochs): | Loss | in progress | | | Data | in progress | | | Convolutional Neural Network | in progress | | -| Distributed | in progress | +| Distributed | in progress | diff --git a/build/distributed.o b/build/distributed.o index f9c8535..11a92ef 100644 Binary files a/build/distributed.o and b/build/distributed.o differ diff --git a/examples/train.ipynb b/examples/train.ipynb index c53ea86..427933a 100644 --- a/examples/train.ipynb +++ b/examples/train.ipynb @@ -33,7 +33,7 @@ "import norch\n", "import norch.nn as nn\n", "import norch.optim as optim\n", - "from norch.utils.data.dataloader import Dataloader\n", + "from norch.utils.data.dataloader import DataLoader\n", "from norch.norchvision import transforms\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", @@ -102,7 +102,7 @@ "device = \"cuda\" #cpu\n", "epochs = 10\n", "\n", - "transform = transforms.Sequential(\n", + "transform = transforms.Compose(\n", " [\n", " transforms.ToTensor(),\n", " transforms.Reshape([-1, 784, 1])\n", @@ -116,7 +116,7 @@ ")\n", "\n", "train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)\n", - "train_loader = Dataloader(train_data, batch_size = BATCH_SIZE)\n", + "train_loader = DataLoader(train_data, batch_size = BATCH_SIZE)\n", "\n", "class MyModel(nn.Module):\n", " def __init__(self):\n", diff --git a/norch/__init__.py b/norch/__init__.py index c2a3122..2aae8f0 100644 --- a/norch/__init__.py +++ b/norch/__init__.py @@ -3,6 +3,7 @@ from .nn import * from .optim import * from .utils import * from .norchvision import * +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/csrc/distributed.cpp b/norch/csrc/distributed.cpp index 6b25365..bb753b2 100644 --- a/norch/csrc/distributed.cpp +++ b/norch/csrc/distributed.cpp @@ -35,12 +35,12 @@ void init_process_group(int env_rank, int env_world_size) { } -void broadcast_tensor(Tensor* tensor) { +void broadcast_tensor(Tensor* tensor, int src) { cudaStream_t stream; cudaStreamCreate(&stream); - NCCL_CHECK(ncclBroadcast(tensor->data, tensor->data, tensor->size * sizeof(float), ncclFloat, 0, nccl_comm, stream)); + NCCL_CHECK(ncclBroadcast(tensor->data, tensor->data, tensor->size * sizeof(float), ncclFloat, src, nccl_comm, stream)); cudaStreamSynchronize(stream); cudaStreamDestroy(stream); } @@ -56,6 +56,19 @@ void allreduce_sum_tensor(Tensor* tensor) { cudaStreamDestroy(stream); } +void allreduce_mean_tensor(Tensor* tensor) { + cudaStream_t stream; + cudaStreamCreate(&stream); + + // Perform NCCL AllReduce operation to calculate the mean of all tensors across all processes + NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream)); + + tensor_div_scalar_cuda(tensor, world_size, tensor->data); + + cudaStreamSynchronize(stream); + cudaStreamDestroy(stream); +} + void end_process_group() { MPI_CHECK(MPI_Finalize()); NCCL_CHECK(ncclCommDestroy(nccl_comm)); diff --git a/norch/csrc/distributed.h b/norch/csrc/distributed.h index 1992b8f..53e318e 100644 --- a/norch/csrc/distributed.h +++ b/norch/csrc/distributed.h @@ -22,8 +22,9 @@ extern "C" { void init_process_group(int rank, int world_size); - void broadcast_tensor(Tensor* tensor); + void broadcast_tensor(Tensor* tensor, int src); void allreduce_sum_tensor(Tensor* tensor); + void allreduce_mean_tensor(Tensor* tensor); } #endif /* DISTRIBUTED_H */ diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py index 9c797f8..98bc2a0 100644 --- a/norch/distributed/distributed.py +++ b/norch/distributed/distributed.py @@ -9,12 +9,18 @@ def init_process_group(rank, world_size, backend='nccl'): Tensor._C.init_process_group(rank, world_size) -def broadcast_tensor(tensor): +def get_rank(): + return os.getenv('OMPI_COMM_WORLD_RANK', 0) - Tensor._C.broadcast_tensor.argtypes = [ctypes.POINTER(CTensor)] +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] Tensor._C.broadcast_tensor.restype = None - Tensor._C.broadcast_tensor(tensor.tensor) + Tensor._C.broadcast_tensor(tensor.tensor, src) def allreduce_sum_tensor(tensor): @@ -22,3 +28,10 @@ def allreduce_sum_tensor(tensor): Tensor._C.allreduce_sum_tensor.restype = None Tensor._C.allreduce_sum_tensor(tensor.tensor) + +def allreduce_mean_tensor(tensor): + + Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)] + Tensor._C.allreduce_mean_tensor.restype = None + + Tensor._C.allreduce_mean_tensor(tensor.tensor) diff --git a/norch/libtensor.so b/norch/libtensor.so index 9e4f3e6..dc9871e 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/nn/__init__.py b/norch/nn/__init__.py index 8f7997b..c131ea4 100644 --- a/norch/nn/__init__.py +++ b/norch/nn/__init__.py @@ -1,4 +1,5 @@ from .modules import * from .activation import * from .loss import * -from .functional import * \ No newline at end of file +from .functional import * +from .parallel import * \ No newline at end of file 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/parallel.py b/norch/nn/parallel.py new file mode 100644 index 0000000..7d5060b --- /dev/null +++ b/norch/nn/parallel.py @@ -0,0 +1,21 @@ +from .module import * +import norch.distributed as dist +import os + +class DistributedDataParallel(Module): + def __init__(self, module): + super().__init__() + + self.module = module + + def forward(self, *inputs, **kwargs): + return self.module(*inputs, **kwargs) + + def backward(self): + self.module.backward() + for module, name, parameter in self.parameters(): + dist.allreduce_mean_tensor(parameter) + + + + 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/norchvision/transforms.py b/norch/norchvision/transforms.py index 03dafad..944e426 100644 --- a/norch/norchvision/transforms.py +++ b/norch/norchvision/transforms.py @@ -11,7 +11,7 @@ class Reshape: def __call__(self, x): return x.reshape(self.shape) -class Sequential: +class Compose: def __init__(self, transforms): self.transforms = transforms 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/__init__.py b/norch/utils/data/__init__.py index 2397554..e5ea955 100644 --- a/norch/utils/data/__init__.py +++ b/norch/utils/data/__init__.py @@ -1,4 +1,5 @@ from .dataset import * from .example import * from .dataloader import * -from .batch import * \ No newline at end of file +from .batch import * +from .distributed import * \ No newline at end of file diff --git a/norch/utils/data/dataloader.py b/norch/utils/data/dataloader.py index 01bd41e..fa823ad 100644 --- a/norch/utils/data/dataloader.py +++ b/norch/utils/data/dataloader.py @@ -2,19 +2,27 @@ import numpy as np from .batch import Batch -class Dataloader(object): - - def __init__(self, dataset, batch_size=32): +class DataLoader: + + def __init__(self, dataset, batch_size, sampler=None): self.dataset = dataset self.batch_size = batch_size + self.sampler = sampler def __iter__(self): - starts = np.arange(0, len(self.dataset), self.batch_size) - - for start in starts: - end = start + self.batch_size - batch_size = min(end, len(self.dataset)) - start - yield Batch(self.dataset[start:end], batch_size) - + if self.sampler is not None: + indices = list(iter(self.sampler)) + + else: + indices = list(range(len(self.dataset))) + + for i in range(0, len(indices), self.batch_size): + batch_indices = indices[i:i + self.batch_size] + batch_data = self.dataset[batch_indices] + yield Batch(batch_data, len(batch_data)) + 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 new file mode 100644 index 0000000..5d9f294 --- /dev/null +++ b/norch/utils/data/distributed.py @@ -0,0 +1,28 @@ +import numpy as np +import math + +class DistributedSampler: + def __init__(self, dataset, num_replicas, rank): + self.dataset = dataset + self.num_replicas = num_replicas + self.rank = rank + self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas)) + self.total_size = self.num_samples * self.num_replicas + + def __iter__(self): + indices = list(range(len(self.dataset))) + + # Add extra samples to make it evenly divisible + indices = indices[:self.total_size] + if len(indices) < self.total_size: + indices += indices[:(self.total_size - len(indices))] + assert len(indices) == self.total_size + + # 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 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/tests/test_dataset.py b/tests/test_dataset.py new file mode 100644 index 0000000..aceb53c --- /dev/null +++ b/tests/test_dataset.py @@ -0,0 +1,93 @@ +import unittest +import norch +from norch.utils import utils_unittests as utils +import torch +import torchvision +import numpy as np +from norch.norchvision import transforms as norch_transforms +from torchvision import transforms as torch_transforms + +class TestDataset(unittest.TestCase): + def test_dataloader_batchsize_1(self): + transforms = norch_transforms.Compose( + [ + norch_transforms.ToTensor(), + ] + ) + train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transforms, target_transform=transforms) + train_loader = norch.utils.data.DataLoader(train_data, batch_size = 1) + + labels_norch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_norch.append(utils.to_torch(label)) + + if i > 10: + break + + transforms = torch_transforms.Compose( + [ + torch_transforms.ToTensor(), + ] + ) + + train_data = torchvision.datasets.MNIST(root='./.data/', download=True, transform=transforms) + train_loader = torch.utils.data.DataLoader(train_data, batch_size = 1) + + labels_torch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_torch.append(label) + + if i > 10: + break + + for label_norch, label_torch in zip(labels_norch, labels_torch): + self.assertTrue(utils.compare_torch(label_norch, label_torch)) + + + def test_dataloader_batchsize_32(self): + transforms = norch_transforms.Compose( + [ + norch_transforms.ToTensor(), + ] + ) + train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transforms, target_transform=transforms) + train_loader = norch.utils.data.DataLoader(train_data, batch_size = 32) + + labels_norch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_norch.append(utils.to_torch(label)) + + if i > 10: + break + + transforms = torch_transforms.Compose( + [ + torch_transforms.ToTensor(), + ] + ) + + train_data = torchvision.datasets.MNIST(root='./.data/', download=True, transform=transforms) + train_loader = torch.utils.data.DataLoader(train_data, batch_size = 32) + + labels_torch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_torch.append(label) + + if i > 10: + break + + for label_norch, label_torch in zip(labels_norch, labels_torch): + self.assertTrue(utils.compare_torch(label_norch, label_torch)) + + + + + diff --git a/tests/test_distributed.py b/tests/test_distributed.py new file mode 100644 index 0000000..1b16a20 --- /dev/null +++ b/tests/test_distributed.py @@ -0,0 +1,139 @@ +import unittest +import norch +from norch.utils import utils_unittests as utils +import torch +import torchvision +import numpy as np +import os +from norch.norchvision import transforms as norch_transforms +from torchvision import transforms as torch_transforms + +class TestDistributed(unittest.TestCase): + def test_distributed_sampler_batch_1(self): + transforms = norch_transforms.Compose( + [ + norch_transforms.ToTensor(), + ] + ) + train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transforms, target_transform=transforms) + distributed_sampler = norch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=8, rank=2) + train_loader = norch.utils.data.DataLoader(train_data, batch_size=1, sampler=distributed_sampler) + + labels_norch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_norch.append(utils.to_torch(label)) + + if i > 100: + break + + transforms = torch_transforms.Compose( + [ + torch_transforms.ToTensor(), + ] + ) + + train_data = torchvision.datasets.MNIST(root='./.data/', download=True, transform=transforms) + distributed_sampler = torch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=8, rank=2, shuffle=False) + train_loader = torch.utils.data.DataLoader(train_data, batch_size=1, sampler=distributed_sampler) + + labels_torch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_torch.append(label) + + if i > 100: + break + + for label_norch, label_torch in zip(labels_norch, labels_torch): + self.assertTrue(utils.compare_torch(label_norch, label_torch)) + + + def test_distributed_sampler_batch_32_rank_2(self): + transforms = norch_transforms.Compose( + [ + norch_transforms.ToTensor(), + ] + ) + train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transforms, target_transform=transforms) + distributed_sampler = norch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=8, rank=2) + train_loader = norch.utils.data.DataLoader(train_data, batch_size=32, sampler=distributed_sampler) + + labels_norch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_norch.append(utils.to_torch(label)) + + if i > 100: + break + + transforms = torch_transforms.Compose( + [ + torch_transforms.ToTensor(), + ] + ) + + train_data = torchvision.datasets.MNIST(root='./.data/', download=True, transform=transforms) + distributed_sampler = torch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=8, rank=2, shuffle=False) + train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, sampler=distributed_sampler) + + labels_torch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_torch.append(label) + + if i > 100: + break + + for label_norch, label_torch in zip(labels_norch, labels_torch): + self.assertTrue(utils.compare_torch(label_norch, label_torch)) + + def test_distributed_sampler_batch_32_rank_7(self): + transforms = norch_transforms.Compose( + [ + norch_transforms.ToTensor(), + ] + ) + train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transforms, target_transform=transforms) + distributed_sampler = norch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=8, rank=7) + train_loader = norch.utils.data.DataLoader(train_data, batch_size=32, sampler=distributed_sampler) + + labels_norch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_norch.append(utils.to_torch(label)) + + if i > 100: + break + + transforms = torch_transforms.Compose( + [ + torch_transforms.ToTensor(), + ] + ) + + train_data = torchvision.datasets.MNIST(root='./.data/', download=True, transform=transforms) + distributed_sampler = torch.utils.data.distributed.DistributedSampler(dataset=train_data, num_replicas=8, rank=7, shuffle=False) + train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, sampler=distributed_sampler) + + labels_torch = [] + for i, batch in enumerate(train_loader): + + image, label = batch + labels_torch.append(label) + + if i > 100: + break + + for label_norch, label_torch in zip(labels_norch, labels_torch): + self.assertTrue(utils.compare_torch(label_norch, label_torch)) + + + + + diff --git a/tests/test_nn.py b/tests/test_nn.py index 532f5b6..50c1a79 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -200,5 +200,4 @@ class TestNNModuleActivationFn(unittest.TestCase): softmax_torch_expected = softmax_fn_torch.forward(torch_input) # Compare the results - self.assertTrue(utils.compare_torch(softmax_torch_result, softmax_torch_expected)) - + self.assertTrue(utils.compare_torch(softmax_torch_result, softmax_torch_expected)) \ No newline at end of file diff --git a/train.py b/train.py index 8f07a41..9ad1187 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,31 @@ def main(): if __name__ == "__main__": main() + +""" +import norch +import matplotlib.pyplot as plt +import numpy as np +import random +from norch.norchvision import transforms + +train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transforms.ToTensor()) +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 = 50, 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