Merge pull request #86 from lucasdelimanogueira/tmp

Tmp

README.md

@@ -1,13 +1,13 @@
 # PyNorch
-Recreating PyTorch from scratch (C/C++, CUDA and Python, with GPU support and automatic differentiation!)
+Recreating PyTorch from scratch (C/C++, CUDA and Python, with multi-GPU support and automatic differentiation!)
 
 Project details explanations can also be found on [medium](https://towardsdatascience.com/recreating-pytorch-from-scratch-with-gpu-support-and-automatic-differentiation-8f565122a3cc).
 
 # 1 - About
-**PyNorch** is a deep learning framework constructed using C/C++, CUDA and Python. This is a personal project with educational purpose only! `Norch` means **NOT** PyTorch, and we have **NO** claims to rivaling the already established PyTorch. The main objective of **PyNorch** was to give a brief understanding of how a deep learning framework works internally. It implements the Tensor object, GPU support and an automatic differentiation system. 
+**PyNorch** is a deep learning framework constructed using C/C++, CUDA and Python. This is a personal project with educational purpose only! `Norch` means **NOT** PyTorch, and we have **NO** claims to rivaling the already established PyTorch. The main objective of **PyNorch** was to give a brief understanding of how a deep learning framework works internally. It implements the Tensor object, multi-GPU support and an automatic differentiation system. 
 
 # 2 - Installation
-Install this package from PyPi (you can test on Colab!)
+Install this package from PyPi (you can test on Colab! Also tested on AWS g4dn.12xlarge instance with image ami-061debf863768593d)
 
 ```css
 $ pip install norch
@@ -60,8 +60,10 @@ class MyModel(nn.Module):
         return out
 ```
 
-### 3.3 - Example training
+### 3.3 - Example single GPU training
 ```python
+# examples/train_singlegpu.py
+
 import norch
 from norch.utils.data.dataloader import DataLoader
 from norch.norchvision import transforms as T
@@ -135,6 +137,107 @@ for epoch in range(epochs):
 
 ```
 
+### 3.4 - Example multi-GPU training
+First create a file .py as the example below
+
+```python
+# examples/train_multigpu.py
+
+import os
+import norch
+import norch.distributed as dist
+import norch.distributed
+import norch.nn as nn
+import norch.optim as optim
+from norch.nn.parallel import DistributedDataParallel
+from norch.utils.data.distributed import DistributedSampler
+from norch.norchvision import transforms as T
+import random
+random.seed(1)
+
+local_rank = int(os.getenv('OMPI_COMM_WORLD_LOCAL_RANK', -1))
+rank = int(os.getenv('OMPI_COMM_WORLD_RANK', -1))
+world_size = int(os.getenv('OMPI_COMM_WORLD_SIZE', -1))
+
+dist.init_process_group(
+    rank, 
+    world_size
+)
+
+BATCH_SIZE = 32
+device = local_rank
+epochs = 10
+
+transform = T.Compose(
+    [
+        T.ToTensor(),
+        T.Reshape([-1, 784, 1])
+    ]
+)
+
+target_transform = T.Compose(
+    [
+        T.ToTensor()
+    ]
+)
+
+train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
+distributed_sampler = DistributedSampler(dataset=train_data, num_replicas=world_size, rank=rank)
+train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, sampler=distributed_sampler)
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+        self.fc1 = nn.Linear(784, 30)
+        self.sigmoid1 = nn.Sigmoid()
+        self.fc2 = nn.Linear(30, 10)
+        self.sigmoid2 = nn.Sigmoid()
+
+    def forward(self, x):
+        out = self.fc1(x)
+        out = self.sigmoid1(out)
+        out = self.fc2(out)
+        out = self.sigmoid2(out)
+        
+        return out
+
+model = MyModel().to(device)
+model = DistributedDataParallel(model)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.SGD(model.parameters(), lr=0.01)
+loss_list = []
+
+print(f"Starting training on Rank {rank}/{world_size}\n\n")
+
+for epoch in range(epochs):    
+    for idx, batch in enumerate(train_loader):
+
+        inputs, target = batch
+
+        inputs = inputs.to(device)
+        target = target.to(device)
+
+        outputs = model(inputs)
+        
+        loss = criterion(outputs, target)
+        
+        optimizer.zero_grad()
+        
+        loss.backward()
+
+        optimizer.step()
+    
+    if rank == 0:
+        print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')
+        loss_list.append(loss[0])
+```
+Then you can run using
+
+```css
+$ python3 -m norch.distributed.run --nproc_per_node 4 examples/train_multigpu.py
+```
+
+
 
 # 4 - Progress
 
@@ -144,4 +247,4 @@ for epoch in range(epochs):
 | Loss                         | in progress | <ul><li>[x] MSE</li><li>[X] Cross Entropy</li></ul>    |
 | Data                         | in progress    | <ul><li>[X] Dataset</li><li>[X] Batch</li><li>[X] Iterator</li></ul>   |
 | Convolutional Neural Network | in progress    | <ul><li>[ ] Conv2d</li><li>[ ] MaxPool2d</li><li>[ ] Dropout</li></ul> |
-| Distributed                  | in progress | <ul><li>[ ] All-reduce</li><li>[X] DistributedSampler</li><li>[ ] DistributedDataParallel</li></ul>                 |
+| Distributed                  | in progress | <ul><li>[X] All-reduce</li><li>[X] Broadcast</li><li>[X] DistributedSampler</li><li>[X] DistributedDataParallel</li></ul>                 |

examples/train_multigpu.py

@@ -2,50 +2,24 @@ import os
 import norch
 import norch.distributed as dist
 import norch.distributed
+import norch.nn as nn
+import norch.optim as optim
+from norch.nn.parallel import DistributedDataParallel
+from norch.utils.data.distributed import DistributedSampler
+from norch.norchvision import transforms as T
+import random
+random.seed(1)
 
 def main():
-
+    
     local_rank = int(os.getenv('OMPI_COMM_WORLD_LOCAL_RANK', -1))
     rank = int(os.getenv('OMPI_COMM_WORLD_RANK', -1))
     world_size = int(os.getenv('OMPI_COMM_WORLD_SIZE', -1))
 
-    dist.init_process_group(rank, world_size)
-
-    tensor = norch.Tensor([1,1,1]).to(rank)
-    tensor = (rank + 1) * tensor
-    print(f"BEFORE on rank {rank}: {tensor} \n\n")
-
-    dist.allreduce_sum_tensor(tensor)
-
-    print(f"AFTER ALLREDUCE on rank {rank}: {tensor} \n\n")
-
-    print("###############\n\n\n")
-
-    tensor = tensor * 10
-    print(f"BEFORE BROADCAST on rank {rank}: {tensor} \n\n")
-
-    dist.broadcast_tensor(tensor)
-
-    print(f"AFTER BROADCAST on rank {rank}: {tensor} \n\n")
-
-def main2():
-    import norch
-    import norch.nn as nn
-    import norch.optim as optim
-    from norch.utils.data.dataloader import DataLoader
-    from norch.nn.parallel import DistributedDataParallel
-    from norch.utils.data.distributed import DistributedSampler
-    from norch.norchvision import transforms as T
-    import numpy as np
-    import matplotlib.pyplot as plt
-    import random
-    random.seed(1)
-
-    local_rank = int(os.getenv('OMPI_COMM_WORLD_LOCAL_RANK', -1))
-    rank = int(os.getenv('OMPI_COMM_WORLD_RANK', -1))
-    world_size = int(os.getenv('OMPI_COMM_WORLD_SIZE', -1))
-
-    dist.init_process_group(rank, world_size)
+    dist.init_process_group(
+        rank, 
+        world_size
+    )
 
     BATCH_SIZE = 32
     device = local_rank
@@ -65,7 +39,7 @@ def main2():
     )
 
     train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
-    distributed_sampler = DistributedSampler(dataset=train_data, num_replicas=world_size, rank=local_rank)
+    distributed_sampler = DistributedSampler(dataset=train_data, num_replicas=world_size, rank=rank)
     train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, sampler=distributed_sampler)
 
     class MyModel(nn.Module):
@@ -86,11 +60,12 @@ def main2():
 
     model = MyModel().to(device)
     model = DistributedDataParallel(model)
-    print(f"parameter bias on Rank {rank}: {model.module.fc1.bias}\n\n")
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(model.parameters(), lr=0.01)
     loss_list = []
 
+    print(f"Starting training on Rank {rank}/{world_size}\n\n")
+
     for epoch in range(epochs):    
         for idx, batch in enumerate(train_loader):
 
@@ -106,14 +81,12 @@ def main2():
             optimizer.zero_grad()
             
             loss.backward()
-            print(f"GRADIENT AFTER rank {local_rank}: {model.module.fc2.bias.grad}")
-            print("\n\n")
-
 
             optimizer.step()
-            break
-
-        break
+        
+        if rank == 0:
+            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')
+            loss_list.append(loss[0])
 
 if __name__ == "__main__":
     main()

examples/train_singlegpu.py

@@ -1,23 +1,14 @@
-import os
 import norch
-import norch.distributed as dist
-import norch.distributed
-
 import norch.nn as nn
 import norch.optim as optim
-from norch.utils.data.dataloader import DataLoader
-from norch.nn.parallel import DistributedDataParallel
-from norch.utils.data.distributed import DistributedSampler
 from norch.norchvision import transforms as T
-import numpy as np
-import matplotlib.pyplot as plt
 import random
 random.seed(1)
 
 def main():
 
     BATCH_SIZE = 32
-    device = "cpu"
+    device = "cuda"
     epochs = 10
 
     transform = T.Compose(
@@ -53,7 +44,6 @@ def main():
             return out
 
     model = MyModel().to(device)
-    model = DistributedDataParallel(model)
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(model.parameters(), lr=0.01)
     loss_list = []
@@ -70,17 +60,11 @@ def main():
             loss = criterion(outputs, target)
             
             optimizer.zero_grad()
-            print(f"GRADIENT BEFORE: {model.module.fc2.bias.grad}")
             
             loss.backward()
-            print(f"GRADIENT AFTER: {model.module.fc2.bias.grad}")
-            print("\n\n")
-
 
             optimizer.step()
-            break
 
-        break
 
 if __name__ == "__main__":
     main()

norch/__init__.py

@@ -5,6 +5,6 @@ from .utils import *
 from .norchvision import *
 from .utils import *
 
-__version__ = "0.0.4"
+__version__ = "0.0.5"
 __author__ = 'Lucas de Lima Nogueira'
 __credits__ = 'Lucas de Lima Nogueira'

norch/csrc/distributed.cpp

@@ -58,19 +58,6 @@ 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));

norch/distributed/distributed.py

@@ -28,10 +28,3 @@ 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)

setup.py

@@ -76,7 +76,7 @@ class CustomInstall(install):
 
 setuptools.setup(
     name="norch",
-    version="0.0.4",
+    version="0.0.5",
     author="Lucas de Lima",
     author_email="nogueiralucasdelima@gmail.com",
     description="A deep learning framework",