From 12eb38a9400daeed0919f41c87d3ddabce698a29 Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Mon, 3 Jun 2024 15:38:16 -0300 Subject: [PATCH] multigpu training support v0.0.5 --- train.py | 87 --------------------------------- train_multigpu.py | 120 ---------------------------------------------- 2 files changed, 207 deletions(-) delete mode 100644 train.py delete mode 100644 train_multigpu.py diff --git a/train.py b/train.py deleted file mode 100644 index 666355c..0000000 --- a/train.py +++ /dev/null @@ -1,87 +0,0 @@ -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" - 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) - model = DistributedDataParallel(model) - 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() - 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() - diff --git a/train_multigpu.py b/train_multigpu.py deleted file mode 100644 index 20e3b6e..0000000 --- a/train_multigpu.py +++ /dev/null @@ -1,120 +0,0 @@ -import os -import norch -import norch.distributed as dist -import norch.distributed - -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) - - 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=local_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) - 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 = [] - - 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() - print(f"GRADIENT AFTER rank {local_rank}: {model.module.fc2.bias.grad}") - print("\n\n") - - - optimizer.step() - break - - break - -if __name__ == "__main__": - main() -