aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

multigpu training support v0.0.5

Browse source
This commit is contained in:
lucasdelimanogueira 2024-06-03 15:38:16 -03:00
parent 64f08a8963
commit 12eb38a940
2 changed files with 0 additions and 207 deletions
Download patch file Download diff file Expand all files Collapse all files

87
train.py
View file

@ -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()

120
train_multigpu.py
View file

@ -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()