If you only look at the time spent on a single epoch, distributed training is indeed faster than singleton training. But our goal is to have a smaller loss, so I did an experiment. In the same time, the loss of single-instance training is smaller than the loss of distributed training.
The data set is imagenet, the network is resnet50, and the following is my code.
distributed_train.py
# -*- coding: utf-8 -*-
import datetime
import os
import shutil
import sys
import time
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.nn as nn
import torch.nn.parallel
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
best_acc1 = 0
def get_dataset(world_size, rank, batch_size, num_workers):
traindir = os.path.join('/root/data/ILSVRC2012', 'train')
valdir = os.path.join('/root/data/ILSVRC2012', 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset,
num_replicas=world_size,
rank=rank,
)
train_batch_size = int(batch_size / world_size)
train_set = torch.utils.data.DataLoader(
train_dataset, batch_size=train_batch_size, shuffle=(train_sampler is None),
num_workers=num_workers, pin_memory=True, sampler=train_sampler)
val_set = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=batch_size, shuffle=False,
num_workers=num_workers, pin_memory=True)
return train_set, val_set
def init_print(rank, size, debug_print=True):
if not debug_print:
""" In case run on hundreds of nodes, you may want to mute all the nodes except master """
if rank > 0:
sys.stdout = open(os.devnull, 'w')
sys.stderr = open(os.devnull, 'w')
else:
# labelled print with info of [rank/size]
old_out = sys.stdout
class LabeledStdout:
def __init__(self, rank, size):
self._r = rank
self._s = size
self.flush = sys.stdout.flush
self._start_time = datetime.datetime.now()
def write(self, x):
if x == '\n':
old_out.write(x)
else:
now = datetime.datetime.now()
seconds = (now - self._start_time).total_seconds()
old_out.write('%s [%d/%d] %s' % (seconds, self._r, self._s, x))
sys.stdout = LabeledStdout(rank, size)
def main():
multiprocessing_distributed = True
lr = 0.01
weight_decay = 1e-4
momentum = 0.9
start_epoch = 0
epochs = 2
num_workers = 16
batch_size = 64
ngpus_per_node = torch.cuda.device_count()
print('ngpus per node = %s' % ngpus_per_node)
print('epochs = %s' % epochs)
dist_backend = dist.Backend.MPI
rank = int(os.environ.get('RANK', '0'))
world_size = int(os.environ.get('WORLD_SIZE', '0'))
print('world size = %s, rank = %s' % (world_size, rank))
dist.init_process_group(backend=dist_backend, rank=rank, world_size=world_size)
ngpus_per_node = torch.cuda.device_count()
global best_acc1
init_print(rank, world_size)
device = torch.device('cuda')
# create model
model = models.resnet50(pretrained=False)
model = model.to(device)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
model = torch.nn.parallel.DistributedDataParallel(
model,
)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss()
# criterion = criterion.cuda()
criterion = criterion.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr,
momentum=momentum,
weight_decay=weight_decay)
cudnn.benchmark = True
# partition dataset
print('batch size = %s' % batch_size)
print('num_workers = %s' % num_workers)
train_set, val_set = get_dataset(world_size, rank, batch_size, num_workers)
print('start epoch = %s, epochs = %s' % (start_epoch, epochs))
start_time = datetime.datetime.now()
for epoch in range(start_epoch, epochs):
print('epoch = %s' % epoch)
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_set, model, criterion, optimizer, epoch, device)
# evaluate on validation set
acc1 = validate(val_set, model, criterion, start_time)
# remember best acc@1 and save checkpoint
print('acc1 = %s, best_acc1 = %s' % (acc1, best_acc1))
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
print('acc1 = %s, best_acc1 = %s, is_best = %s' % (acc1, best_acc1, is_best))
arch = 'resnet50'
if not multiprocessing_distributed or (multiprocessing_distributed
and rank % ngpus_per_node == 0):
save_checkpoint({
'epoch': epoch + 1,
'arch': arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}, is_best)
def train(train_loader, model, criterion, optimizer, epoch, device):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to train mode
model.train()
end = time.time()
print('end = %s' % end)
print_freq = 20
for i, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input = input.to(device)
target = target.to(device)
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1[0], input.size(0))
top5.update(acc5[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
epoch, i, len(train_loader), batch_time=batch_time,
data_time=data_time, loss=losses, top1=top1, top5=top5))
def validate(val_loader, model, criterion, start_time):
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
for i, (input, target) in enumerate(val_loader):
input = Variable(input).cuda(non_blocking=True)
target = Variable(target).cuda(non_blocking=True)
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1[0], input.size(0))
top5.update(acc5[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
print_freq = 10
if i % print_freq == 0:
print('Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
i, len(val_loader), batch_time=batch_time, loss=losses,
top1=top1, top5=top5))
now = datetime.datetime.now()
delta_seconds = (now - start_time).total_seconds()
print('delta seconds {delta_seconds} * Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(delta_seconds=delta_seconds, top1=top1, top5=top5, loss=losses))
return top1.avg
def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
torch.save(state, filename)
if is_best:
shutil.copyfile(filename, 'model_best.pth.tar')
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def adjust_learning_rate(optimizer, epoch):
lr = 0.01
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
lr = lr * (0.1 ** (epoch // 30))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified values of k"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res
if __name__ == "__main__":
main()
distributed_train.yaml
apiVersion: "kubeflow.org/v1beta1"
kind: "PyTorchJob"
metadata:
name: "pytorch-mpi-imagenet-gpu"
spec:
pytorchReplicaSpecs:
Master:
replicas: 1
restartPolicy: OnFailure
template:
spec:
volumes:
- name: data
nfs:
server: 10.255.1.115
path: "/mnt/data"
- name: aaaa
nfs:
server: 172.12.5.221
path: "/mnt/aaaa"
- name: shm
hostPath:
path: "/dev/shm"
- name: nvidia-driver
hostPath:
path: "/var/lib/nvidia-docker/volumes/nvidia_driver/375.66"
containers:
- name: pytorch
image: "172.16.3.205:5000/nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04-mpi"
command:
- "mpirun"
- "-n"
- "1"
- "--cpu-set"
- "0,1,2,3,4,5,6,7"
- "--allow-run-as-root"
- "python"
- "-u"
- "/root/aaaa/distributed_train.py"
workingDir: "/root/aaaa"
imagePullPolicy: Always
resources:
limits:
cpu: 8
memory: 30Gi
alpha.kubernetes.io/nvidia-gpu: 1
volumeMounts:
- name: data
mountPath: "/root/data"
- name: aaaa
mountPath: "/root/aaaa"
- name: shm
mountPath: "/dev/shm"
- name: nvidia-driver
mountPath: "/usr/local/nvidia"
tolerations:
- key: CriticalAddonsOnly
operator: Exists
Worker:
replicas: 1
restartPolicy: OnFailure
template:
spec:
volumes:
- name: data
nfs:
server: 10.255.1.115
path: "/mnt/data"
- name: aaaa
nfs:
server: 172.12.5.221
path: "/mnt/aaaa"
- name: shm
hostPath:
path: "/dev/shm"
- name: nvidia-driver
hostPath:
path: "/var/lib/nvidia-docker/volumes/nvidia_driver/375.66"
containers:
- name: pytorch
image: "172.16.3.205:5000/nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04-mpi"
command:
- "mpirun"
- "-n"
- "1"
- "--cpu-set"
- "0,1,2,3,4,5,6,7"
- "--allow-run-as-root"
- "python"
- "-u"
- "/root/aaaa/distributed_train.py"
workingDir: "/root/aaaa"
imagePullPolicy: Always
resources:
limits:
cpu: 8
memory: 30Gi
alpha.kubernetes.io/nvidia-gpu: 1
volumeMounts:
- name: data
mountPath: "/root/data"
- name: aaaa
mountPath: "/root/aaaa"
- name: shm
mountPath: "/dev/shm"
- name: nvidia-driver
mountPath: "/usr/local/nvidia"
tolerations:
- key: CriticalAddonsOnly
operator: Exists
single_train.py
# -*- coding: utf-8 -*-
import datetime
import os
import shutil
import sys
import time
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
best_acc1 = 0
def main():
global best_acc1
lr = 0.01
momentum = 0.9
weight_decay = 1e-4
start_epoch = 0
epochs = 2
num_workers = 16
batch_size = 64
print('num_workers = %s' % num_workers)
print('batch_size = %s' % batch_size)
print('epochs = %s' % epochs)
init_print(0, 1)
device = torch.device('cuda')
# create model
model = models.resnet50(pretrained=False)
model = model.to(device)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss()
criterion = criterion.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr,
momentum=momentum,
weight_decay=weight_decay)
cudnn.benchmark = True
# Data loading code
traindir = os.path.join('/root/data/ILSVRC2012', 'train')
valdir = os.path.join('/root/data/ILSVRC2012', 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
train_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=batch_size, shuffle=(train_sampler is None),
num_workers=num_workers, pin_memory=True, sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=batch_size, shuffle=False,
num_workers=num_workers, pin_memory=True)
sys.stdout.flush()
for epoch in range(start_epoch, epochs):
adjust_learning_rate(optimizer, epoch, lr)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, device)
# evaluate on validation set
acc1 = validate(val_loader, model, criterion)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
arch = 'resnet50'
save_checkpoint({
'epoch': epoch + 1,
'arch': arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}, is_best)
def train(train_loader, model, criterion, optimizer, epoch, device):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to train mode
model.train()
end = time.time()
print_freq = 20
for i, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input = input.to(device)
target = target.to(device)
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1[0], input.size(0))
top5.update(acc5[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
epoch, i, len(train_loader), batch_time=batch_time,
data_time=data_time, loss=losses, top1=top1, top5=top5))
sys.stdout.flush()
def validate(val_loader, model, criterion):
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
print_freq = 20
for i, (input, target) in enumerate(val_loader):
input = Variable(input).cuda(non_blocking=True)
target = Variable(target).cuda(non_blocking=True)
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1[0], input.size(0))
top5.update(acc5[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % print_freq == 0:
print('Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
i, len(val_loader), batch_time=batch_time, loss=losses,
top1=top1, top5=top5))
print(' * Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(top1=top1, top5=top5, loss=losses))
return top1.avg
def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
torch.save(state, filename)
if is_best:
shutil.copyfile(filename, 'model_best.pth.tar')
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def adjust_learning_rate(optimizer, epoch, lr):
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
lr = lr * (0.1 ** (epoch // 30))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified values of k"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res
def init_print(rank, size, debug_print=True):
if not debug_print:
""" In case run on hundreds of nodes, you may want to mute all the nodes except master """
if rank > 0:
sys.stdout = open(os.devnull, 'w')
sys.stderr = open(os.devnull, 'w')
else:
# labelled print with info of [rank/size]
old_out = sys.stdout
class LabeledStdout:
def __init__(self, rank, size):
self._r = rank
self._s = size
self.flush = sys.stdout.flush
self._start_time = datetime.datetime.now()
def write(self, x):
if x == '\n':
old_out.write(x)
else:
now = datetime.datetime.now()
seconds = (now - self._start_time).total_seconds()
old_out.write('%s [%d/%d] %s' % (seconds, self._r, self._s, x))
sys.stdout = LabeledStdout(rank, size)
if __name__ == '__main__':
main()
single_train.json
{
"kind": "Job",
"apiVersion": "batch/v1",
"metadata": {
"name": "hqx-pytorch-1wxud3d",
"labels": {
"name": "hqx-pytorch-1wxud3d"
}
},
"spec": {
"parallelism": 1,
"completions": 1,
"backoffLimit": 4,
"template": {
"spec": {
"volumes": [
{
"name": "aaaa",
"nfs": {
"server": "172.12.5.221",
"path": "/mnt/aaaa"
}
},
{
"name": "data",
"nfs": {
"server": "11.255.3.115",
"path": "/mnt/data"
}
},
{
"name": "shm",
"hostPath": {
"path": "/dev/shm"
}
},
{
"name": "nvidia-driver",
"hostPath": {
"path": "/var/lib/nvidia-docker/volumes/nvidia_driver/375.66"
}
}
],
"containers": [
{
"name": "hqx-pytorch-1wxud3d",
"image": "172.16.3.205:5000/pytorch/pytorch:latest",
"command": [
"python3",
"/root/aaaa/single_train.py"
],
"workingDir": "/root/aaaa/",
"env": [
{
"name": "POD_HOST_IP",
"valueFrom": {
"fieldRef": {
"fieldPath": "status.hostIP"
}
}
}
],
"resources": {
"limits": {
"cpu": "8",
"memory": "30Gi",
"alpha.kubernetes.io/nvidia-gpu": 1
}
},
"volumeMounts": [
{
"name": "aaaa",
"mountPath": "/root/aaaa"
},
{
"name": "data",
"mountPath": "/root/data"
},
{
"name": "shm",
"mountPath": "/dev/shm"
},
{
"name": "nvidia-driver",
"mountPath": "/usr/local/nvidia"
}
],
"imagePullPolicy": "IfNotPresent"
}
],
"tolerations": [
{
"key": "CriticalAddonsOnly",
"operator": "Exists"
}
],
"restartPolicy": "Never"
}
}
}
}
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent