diff --git a/README.md b/README.md
index a6491be..d6eb1bb 100644
--- a/README.md
+++ b/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 |
|
| Data | in progress | - [X] Dataset
- [X] Batch
- [X] Iterator
|
| Convolutional Neural Network | in progress | - [ ] Conv2d
- [ ] MaxPool2d
- [ ] Dropout
|
-| Distributed | in progress | - [ ] All-reduce
- [X] DistributedSampler
- [ ] DistributedDataParallel
|
+| Distributed | in progress | - [X] All-reduce
- [X] Broadcast
- [X] DistributedSampler
- [X] DistributedDataParallel
|
diff --git a/build/cuda.cu.o b/build/cuda.cu.o
index 1d8c326..427e573 100644
Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ
diff --git a/build/distributed.o b/build/distributed.o
index 58e0ef2..caf2c5f 100644
Binary files a/build/distributed.o and b/build/distributed.o differ
diff --git a/examples/train_multigpu.py b/examples/train_multigpu.py
new file mode 100644
index 0000000..8cbaaff
--- /dev/null
+++ b/examples/train_multigpu.py
@@ -0,0 +1,93 @@
+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
+ )
+
+ 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])
+
+if __name__ == "__main__":
+ main()
+
diff --git a/examples/train.ipynb b/examples/train_nn_mnist.ipynb
similarity index 100%
rename from examples/train.ipynb
rename to examples/train_nn_mnist.ipynb
diff --git a/examples/train_singlegpu.py b/examples/train_singlegpu.py
new file mode 100644
index 0000000..8c58ed6
--- /dev/null
+++ b/examples/train_singlegpu.py
@@ -0,0 +1,71 @@
+import norch
+import norch.nn as nn
+import norch.optim as optim
+from norch.norchvision import transforms as T
+import random
+random.seed(1)
+
+def main():
+
+ BATCH_SIZE = 32
+ device = "cuda"
+ 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)
+ train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE)
+
+ 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)
+ criterion = nn.CrossEntropyLoss()
+ optimizer = optim.SGD(model.parameters(), lr=0.01)
+ loss_list = []
+
+ 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 __name__ == "__main__":
+ main()
+
diff --git a/norch/__init__.py b/norch/__init__.py
index 2aae8f0..355e312 100644
--- a/norch/__init__.py
+++ b/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'
\ No newline at end of file
diff --git a/norch/csrc/distributed.cpp b/norch/csrc/distributed.cpp
index 0a1cc69..bc993fa 100644
--- a/norch/csrc/distributed.cpp
+++ b/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));
diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py
index d534ec4..7d41dec 100644
--- a/norch/distributed/distributed.py
+++ b/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)
diff --git a/norch/libtensor.so b/norch/libtensor.so
index bc35ee2..68eb2f2 100755
Binary files a/norch/libtensor.so and b/norch/libtensor.so differ
diff --git a/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc b/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc
index 0ee226c..0daa1ff 100644
Binary files a/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc and b/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc differ
diff --git a/setup.py b/setup.py
index 7425e0e..2dfaca7 100644
--- a/setup.py
+++ b/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",