dnddnjs / pytorch-multigpu Goto Github PK

ngpus_per_node = torch.cuda.device_count()

args.world_size = ngpus_per_node * args.world_size
mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))

args.rank = args.rank * ngpus_per_node + gpu    
dist.init_process_group(backend=args.dist_backend
                        ,init_method=args.dist_url
                        ,world_size=args.world_size
                        ,rank=args.rank)

print('==> Making model..')
net = pyramidnet()
torch.cuda.set_device(args.gpu)
net.cuda(args.gpu)


args.batch_size = int(args.batch_size / ngpus_per_node)
args.num_workers = int(args.num_workers / ngpus_per_node)

net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.gpu])


num_params = sum(p.numel() for p in net.parameters() if p.requires_grad)
print('The number of parameters of model is', num_params)

print('==> Preparing data..')
transforms_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])

dataset_train = CIFAR10(root='../data',
                        train=True,
                        download=True,
                        transform=transforms_train)
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset_train)
train_loader = DataLoader(dataset_train,
                          batch_size=args.batch_size,
                          shuffle=(train_sampler is None),
                          num_workers=args.num_workers,
                          sampler=train_sampler)

# there are 10 classes so the dataset name is cifar-10
classes = ('plane', 'car', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(),
                      lr=args.lr,
                      momentum=0.9,
                      weight_decay=1e-4)

train(net, criterion, optimizer, train_loader, args.gpu)

train_loss = 0
correct = 0
total = 0

epoch_start = time.time()
for batch_idx, (inputs, targets) in enumerate(train_loader):
    start = time.time()
    
    inputs = inputs.cuda(device)
    targets = targets.cuda(device)
    outputs = net(inputs)
    loss = criterion(outputs, targets)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    train_loss += loss.item()
    _, predicted = outputs.max(1)
    total += targets.size(0)
    correct += predicted.eq(targets).sum().item()

    acc = 100 * correct / total
    
    batch_time = time.time() - start
    
    if batch_idx % 20 == 0:
        print('Epoch: [{}/{}]| loss: {:.3f} | acc: {:.3f} | batch time: {:.3f}s '.format(
            batch_idx, len(train_loader), train_loss/(batch_idx+1), acc, batch_time))

elapse_time = time.time() - epoch_start
elapse_time = datetime.timedelta(seconds=elapse_time)
print("Training time {}".format(elapse_time))