Using DistributedDataParallel onn GANs

username: Naren_Dhyani
cooked: <p>I am trying to train a simple GAN using distributed data parallel. This is my complete code that creates a model, data loader, initializes the process and run it.</p>
<p>The only output I get is of the first epoch</p>
<pre><code class="lang-auto">Epoch: 1        Discriminator Loss: 0.013536    Generator Loss: 0.071964        D(x): 0.724387  D(G(z)): 0.316473 / 0.024269
</code></pre>
<p>My code file below for your reference:</p>
<pre><code class="lang-auto">import os
import numpy as np
import torch
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
from matplotlib import pyplot as plt
import seaborn as sns
from torchvision.utils import make_grid
from torchvision import datasets
import torchvision.utils as vutils
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import torch.multiprocessing as mp
import time
from torchvision.utils import save_image
from torch.utils.data.distributed import DistributedSampler
import argparse

parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')
parser.add_argument('--ex', default="test", type=str, help='name')
parser.add_argument('--resume', '-r', action='store_true',
                    help='resume from checkpoint')
args = parser.parse_args()
path = './archive'
os.path.exists(path)
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(100, 64 * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(64 * 8),
            nn.ReLU(True),
            nn.ConvTranspose2d(64 * 8, 64 * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64 * 4),
            nn.ReLU(True),
            nn.ConvTranspose2d(64 * 4, 64 * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64 * 2),
            nn.ReLU(True),
            nn.ConvTranspose2d(64 * 2, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 3, 4, 2, 1, bias=False),
            nn.Tanh()
        )

    def forward(self, input):
        return self.main(input)

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(3, 64, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 64 * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64 * 2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64 * 2, 64 * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64 * 4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64 * 4, 64 * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64 * 8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64 * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)



def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)


def unnorm(images, means, stds):
    means = torch.tensor(means).reshape(1,3,1,1)
    stds = torch.tensor(stds).reshape(1,3,1,1)
    return images*stds+means

def show_batch(data_loader):
    for images, labels in data_loader:
        fig, ax = plt.subplots(figsize=(15, 15))
        ax.set_xticks([]); ax.set_yticks([])
        unnorm_images = unnorm(images, *norm)
        ax.imshow(make_grid(unnorm_images[:batch_size], nrow=8).permute(1, 2, 0).clamp(0,1))
        break

def example(rank, world_size):
    # create default process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    # Create Data loader
    torch.cuda.set_device(rank)
    norm=((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    batch_size = 64
    image_size = 64
    epochs = 5

    transf = transforms.Compose([
        transforms.Resize(image_size),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize(*norm,inplace=True),
    ])

    dataset = datasets.ImageFolder(root=path,transform=transf)
    train_sampler = DistributedSampler(dataset=dataset,num_replicas=2,rank=rank)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=train_sampler, num_workers=0,pin_memory = True)
    # print(len(dataloader))

    # create local Discriminator model
    model = nn.Linear(10, 10).to(rank)
    modelD = Discriminator()
    modelD.to(rank)
    modelD.apply(weights_init)

    # create local Generator model
    modelG = Generator()
    modelG.to(rank)
    modelG.apply(weights_init)
    
    # construct DDP model
    # ddp_model = DDP(model, device_ids=[rank])
    ddp_modelD = DDP(modelD, device_ids=[rank])
    ddp_modelG = DDP(modelG, device_ids=[rank])
    

    real_label = 1.
    fake_label = 0.
    fixed_noise = torch.randn(64, 100, 1, 1, device=rank)

    # define loss function and optimizer
    criterion = nn.BCELoss()
    optimizerD = optim.Adam(ddp_modelD.parameters(), lr=0.0001, betas=(0.5, 0.999))
    optimizerG = optim.Adam(ddp_modelG.parameters(), lr=0.0001, betas=(0.5, 0.999))
    # optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    img_list = []
    G_losses, D_losses = [], []
    total = 0
    
    if rank == 0:
        sample_dir = './generated'+str(args.ex)
        os.makedirs(sample_dir, exist_ok=True)

    # iterate over epochs
    for epoch in range(epochs):
        g_loss = 0.0
        d_loss = 0.0
        D_x = 0.0
        D_G_z1 = 0.0
        D_G_z2 = 0.0
        calc_time = 0.0
        dataloadingtime_start = time.monotonic ()
        for i, data in enumerate(dataloader, 0):
            # t0 = time.monotonic()
            ddp_modelD.zero_grad()
            real_cpu = data[0].to(rank)
            b_size = real_cpu.size(0)
            label = torch.full((b_size,), real_label, dtype=torch.float, device=rank)
            # Discriminator start
            # dis_start = time.monotonic()
            output = ddp_modelD(real_cpu).view(-1)
            errD_real = criterion(output, label)
            errD_real.backward()
            D_x += output.mean().item()
            noise = torch.randn(b_size, 100, 1, 1, device=rank)
            fake = ddp_modelG(noise)
            label.fill_(fake_label)
            output = ddp_modelD(fake.detach()).view(-1)
            errD_fake = criterion(output, label)
            errD_fake.backward()
            D_G_z1 += output.mean().item()
            errD = errD_real + errD_fake
            optimizerD.step()
            #Discriminator end
            # dis_end = time.monotonic()
            # avg_d_time += (dis_end - dis_start)

            #Generator start
            # gen_start = time.monotonic()
            ddp_modelG.zero_grad()
            label.fill_(real_label)
            output = ddp_modelD(fake).view(-1)
            errG = criterion(output, label)
            errG.backward()
            D_G_z2 += output.mean().item()
            optimizerG.step()
            # gen_end = time.monotonic()
            #Generator end
            # avg_g_time += (gen_end - gen_start)

            g_loss += errG.item()
            d_loss += errD.item()
            total += b_size
            # t1 = time.monotonic()
            # calc_time += (t1-t0)
        
        if rank == 0:
            
            avg_g_loss = g_loss / total
            G_losses.append(avg_g_loss)
            avg_d_loss = d_loss / total
            D_losses.append(avg_d_loss)
            avg_D_x = D_x / len(dataloader)
            avg_D_G_z1 = D_G_z1 / len(dataloader)
            avg_D_G_z2 = D_G_z2 / len(dataloader)
            print('Epoch: {} \tDiscriminator Loss: {:.6f} \tGenerator Loss: {:.6f} \tD(x): {:.6f} \tD(G(z)): {:.6f} / {:.6f}'.format(
                epoch + 1,
                avg_d_loss,
                avg_g_loss,
                avg_D_x,
                avg_D_G_z1,
                avg_D_G_z2
            ))
            with torch.no_grad():
                fake = ddp_modelG(fixed_noise).detach().cpu()
            img_list.append(vutils.make_grid(unnorm(fake, *norm), padding=2, normalize=True))
            fake_fname = 'generated-images-{0:0=4d}.png'.format(epoch+1)
            save_image(unnorm(fake, *norm), os.path.join(sample_dir, fake_fname), nrow=8)
    
    if rank == 0:
        torch.save(ddp_modelG.state_dict(), './'+str(args.ex)+'G.pth')
        torch.save(ddp_modelD.state_dict(), './'+str(args.ex)+'D.pth')
        print('Finished Training')
        end = time.monotonic()

        
        plt.figure(figsize=(20,12))
        plt.plot(G_losses,label="Generator")
        plt.plot(D_losses,label="Discriminator")
        plt.xlabel("Epoch")
        plt.ylabel("Loss")
        plt.legend()
        plt.savefig(str(args.ex)+'_GDLoss.png')
        plt.show()


        import matplotlib.animation as animation
        fig = plt.figure(figsize=(8, 8))
        plt.axis("off")
        ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated=True)] for i in img_list[::6]]
        ani = animation.ArtistAnimation(fig, ims, interval=250, repeat_delay=250, blit=True)
        f = './'+str(args.ex)+'animation.gif'
        writergif = animation.PillowWriter(fps=30) 
        ani.save(f, writer=writergif)
        plt.figure(figsize=(20,12))
        plt.plot(G_losses,label="Generator")
        plt.plot(D_losses,label="Discriminator")
        plt.xlabel("Epoch")
        plt.ylabel("Loss")
        plt.legend()
        plt.savefig(str(args.ex)+'_GDLoss.png')
        plt.show()


        import matplotlib.animation as animation
        fig = plt.figure(figsize=(8, 8))
        plt.axis("off")
        ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated=True)] for i in img_list[::6]]
        ani = animation.ArtistAnimation(fig, ims, interval=250, repeat_delay=250, blit=True)
        f = './'+str(args.ex)+'animation.gif'
        writergif = animation.PillowWriter(fps=30) 
        ani.save(f, writer=writergif)

        

# def init_process(rank, size, fn, backend='nccl'):
#     """ Initialize the distributed environment. """
#     os.environ['MASTER_ADDR'] = '127.0.0.1'
#     os.environ['MASTER_PORT'] = '29500'
#     dist.init_process_group(backend, rank=rank, world_size=size)
#     fn(rank, size)
def main():
    world_size = 2
    # size = 2
    # processes = []
    # mp.set_start_method("spawn")
    # for rank in range(size):
    #     p = mp.Process(target=init_process, args=(rank, size, example))
    #     p.start()
    #     print("Process started for: "+str(rank))
    #     processes.append(p)

    # for p in processes:
    #     p.join()
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    # Environment variables which need to be
    # set when using c10d's default "env"
    # initialization mode.
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29501"
    start = time.monotonic
    main()
    end = time.monotonic
    # train_time  = end-start
    print("Time taken for distributed training over 200 epochs: ", end-start)
    dist.destroy_process_group()
</code></pre>
<p>nvidia-smi gives me this output.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/4/347c9bc401fdeae15d8cb54cff63be57d00fb91f.png" data-download-href="https://discuss.pytorch.org/uploads/default/347c9bc401fdeae15d8cb54cff63be57d00fb91f" title="Screen Shot 2022-05-16 at 1.58.47 AM"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/4/347c9bc401fdeae15d8cb54cff63be57d00fb91f.png" alt="Screen Shot 2022-05-16 at 1.58.47 AM" data-base62-sha1="7ujNEuABefkl5iOhlpAMmkAnSYT" width="690" height="393" data-dominant-color="2E2E2E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-05-16 at 1.58.47 AM</span><span class="informations">1246×710 39.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>You can simply run the code using</p>
<pre><code class="lang-auto">python distdataparalleldcgan.py --ex test
</code></pre>
<p>I am kind of clueless as to what might be going wrong as there are no error outputs ?</p>