04_cifar10_distributeddataparallel.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from tqdm import tqdm
import torchvision
import torchvision.transforms as transforms

# define Net
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

def get_dataset(mode="cpu", batch_size=64, num_workers=2, local_rank=0):
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
    trainsampler = None
    trainloader = None

    # 1. cpu or cuda or DataParallel
    if mode == "cpu" or mode == "gpu" or mode == "DataParallel":
        trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers)

    # 2. DistributedDataParallel
    elif mode == "DistributedDataParallel":
        ####### 1. DistributedSampler设置分布式采样；有了sampler，shuffle=False即可；这里的batch_size不是总和，应该是每个GPU的batch_size；#######
        trainsampler = torch.utils.data.distributed.DistributedSampler(trainset, rank=local_rank)
        trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=False, sampler=trainsampler)

    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

    return trainset, testset, trainloader, testloader, trainsampler

# show some images
def sample_imshow(trainloader, batch_size=64, save_dir='./sample'):
    dataiter = iter(trainloader)
    images, labels = dataiter.next()
    os.makedirs(save_dir, exist_ok=True)
    filename = os.path.join(save_dir, 'sample_grid.png')
    sample_grid = torchvision.utils.make_grid(images, nrow=int(batch_size ** 0.5), normalize=True, pad_value=1)
    torchvision.utils.save_image(sample_grid, filename)
    print(">>> save image to {}".format(filename))
    # print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

def train(log_step_interval=100, save_step_interval=1000, eval_step_interval=200, save_path="ckpt", mode="cpu", epochs=10, batch_size=64, resume=""):
    save_path = f"{save_path}_{mode}"

    if mode == "cpu":
        device = torch.device("cpu")
    elif mode == "gpu" or mode == "DataParallel":
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    elif mode == "DistributedDataParallel":
        n_gpus = 2
        ####### 2. init_process_group #######
        torch.distributed.init_process_group("nccl", world_size=n_gpus)
        ####### 3. set local device #######
        local_rank = torch.distributed.get_rank()
        device = torch.device("cuda", local_rank)
        torch.cuda.set_device(local_rank)
    print(">>> use device : {}".format(device))

    # 1. load dataset
    _, _, trainloader, _, trainsampler = get_dataset(mode=mode, batch_size=batch_size, local_rank=local_rank)

    # 2. define model
    model = Net()

    if local_rank == 0:
        sample_imshow(trainloader)
        print(">>> model param sum : {}".format(sum(p.numel() for p in model.parameters())))
    
    # resume
    if resume != "":
        # load resume pth -> use [cpu/cuda/cuda:index]
        checkpoint = torch.load(resume, map_location=torch.device("cuda:0"))
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        start_epoch = checkpoint['epoch']
        start_step = checkpoint['step']

    ####### 4. model move to cuda && use DataParallel #######
    model.to(device)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])        

    # 3. define loss + optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 4. train
    epoch_index = 0
    start_epoch = 0
    num_batches = len(trainloader)
    for epoch_index in range(start_epoch, start_epoch+epochs+1):
        ema_loss = 0.0
        ####### 5. random dataset in every rank #######
        trainsampler.set_epoch(epoch_index)
        for batch_index, data in enumerate(trainloader):
            step = num_batches*(epoch_index) + batch_index + 1
            ####### 6. data move to cuda #######
            inputs, labels = data[0].to(device), data[1].to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            ema_loss = 0.9*ema_loss + 0.1*loss

            # log loss
            if step % log_step_interval == 0 and local_rank == 0:
                print(">>> epoch: {:5}, step: {:7}, loss: {:2.4f}".format(epoch_index, step, ema_loss.item()))
            
            # 5. save model pth
            if step % save_step_interval == 0 and local_rank == 0:
                os.makedirs(save_path, exist_ok=True)
                save_file = os.path.join(save_path, f"cifar_net_step_{step}.pth")
                torch.save({'epoch': epoch_index,
                            ####### 7. use model.module.state_dict() instead of model.state_dict() #######
                            'model_state_dict': model.module.state_dict(),
                            'optimizer_state_dict': optimizer.state_dict(),
                            'loss': loss,
                            }, save_file)
                print(">>> save pth to {}".format(save_file))
    if local_rank == 0:
        print('>>> Finished Training')
    torch.distributed.destroy_process_group()

if __name__ == '__main__':
    ####### batch_size = batch_size_origin #######
    batch_size = 4
    train(log_step_interval=1000, save_step_interval=10000, eval_step_interval=200, mode="DistributedDataParallel", epochs=100, batch_size=batch_size)

# CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 04_cifar10_gpu_distributeddataparallel.py