MultiGPU를 이용하기위해서 처음에는 Pytorch DataParallel를 써봤지만... GPU가 하나밖에 돌아가지않거나 유틸이 영 좋지못했다
이러한 문제를 해결하기위해서 이것저것 찾아보다가 Distributed DataParallel 라는것을 알게됬고 고생을 꽤많이했다.
일단 여기서는 Distributed DataParallel을 하기위해서 Apex를 이용해서 한번 해보자!
일단 Apex를 설치하기위해서 필요한것은 다음과 같다.
전제조건
Cuda 설치
CUDNN 설치
NCCL 설치
Pytorch 설치
이게 다 설치되었다면 먼저 Apex를 받아보자
git clone [https://github.com/NVIDIA/apex.git](https://github.com/NVIDIA/apex.git)
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
이까지하면 apex는 설치될것이다
그다음 pytorch 에서 apex 를이용해 dataparallel를 진행해보자
import os
from datetime import datetime
import argparse
import torch.multiprocessing as mp
import torchvision
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
import torch
import torch.nn as nn
import torch.distributed as dist
from apex.parallel import DistributedDataParallel as DDP
from apex import amp
import torch.nn as nn
import torch.nn.functional as F
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N',
help='number of data loading workers (default: 4)')
parser.add_argument('-g', '--gpus', default=1, type=int,
help='number of gpus per node')
parser.add_argument('-nr', '--nr', default=0, type=int,
help='ranking within the nodes')
parser.add_argument('--epochs', default=2, type=int, metavar='N',
help='number of total epochs to run')
args = parser.parse_args()
args.world_size = args.gpus * args.nodes
os.environ['MASTER_ADDR'] = '123.456.0.0.1'
os.environ['MASTER_PORT'] = '13555'
mp.spawn(train, nprocs=args.gpus, args=(args,))
class my_network(nn.Module):
def __init__(self):
super(my_network, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=2)
self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=2)
self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=2)
self.conv4 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=2)
self.conv5 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=2)
self.bn1 = nn.BatchNorm2d(32)
self.bn2 = nn.BatchNorm2d(64)
self.bn3 = nn.BatchNorm2d(128)
self.bn4 = nn.BatchNorm2d(256)
self.bn5 = nn.BatchNorm2d(512)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.avg = nn.AvgPool2d(8)
self.fc = nn.Linear(512 * 1 * 1, 2)
def forward(self, x):
x = self.pool(F.leaky_relu(self.bn1(self.conv1(x))))
x = self.pool(F.leaky_relu(self.bn2(self.conv2(x))))
x = self.pool(F.leaky_relu(self.bn3(self.conv3(x))))
x = self.pool(F.leaky_relu(self.bn4(self.conv4(x))))
x = self.pool(F.leaky_relu(self.bn5(self.conv5(x))))
x = self.avg(x)
x = x.view(-1, 512 * 1 * 1)
x = self.fc(x)
return x
def train(gpu, args):
rank = args.nr * args.gpus + gpu
dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
torch.manual_seed(0)
model = my_network()
torch.cuda.set_device(gpu)
model.cuda(gpu)
batch_size = 100
criterion = nn.CrossEntropyLoss().cuda(gpu)
optimizer = torch.optim.SGD(model.parameters(), 1e-4)
model, optimizer = amp.initialize(model, optimizer,
opt_level='O1')
model = DDP(model)
trans = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
train_dataset = torchvision.datasets.ImageFolder(root="/data/train",transform = trans)
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,
num_replicas=args.world_size,
rank=rank)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8,
pin_memory=True,
sampler=train_sampler)
start = datetime.now()
total_step = len(train_loader)
for epoch in range(args.epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
optimizer.step()
if (i + 1) % 100 == 0 and gpu == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
loss.item()))
if gpu == 0:
print("Training complete in: " + str(datetime.now() - start))
if __name__ == '__main__':
main()
대략적으로 코드는 이렇다. init group process 설정 후 Port 는 자신의 ip를 써야한다
이런식으로 py파일을 하나 만들어주자
그다음은 py을 실행해보자
py파일을 실행할때는 기존의 파이썬 실행코드와 다르게 각종 옵션을 선택하여 실행해줘야한다
python3 pytorchexample.py -n 1 -g 8 -nr 0 --epochs 50
n는 노드수
g는 gpu수
epoch 수를 지정하고 실행하면 실행이될것이다.
유틸을 확인해보고싶으면
나는 gpustat 을 이용해서 확인해봤다.
나는init group process 지정하는것이랑 nccl설치에 애를 좀먹었다...생각보다 별거아니었던거같은데 다른사람들은 이글을보고 고생하지않았으면한다