GAN loss remains same during training

SinghJivjot: <p>import torch<br>
from torch import nn, optim<br>
from torch.utils.data import DataLoader<br>
import torchvision<br>
from torchvision import datasets, transforms<br>
from torch.utils.tensorboard import SummaryWriter</p>
<p>class Discriminator(nn.Module):<br>
def <strong>init</strong>(self, img_channels=3, features=64): # 3x64x64<br>
super().<strong>init</strong>()<br>
self.disc = nn.Sequential(<br>
self._make_conv2d_block(img_channels, features, use_bn=False), # 64x32x32</p>
<pre><code>        self._make_conv2d_block(features, features*2), # 128x16x16
        self._make_conv2d_block(features*2, features*4), # 256x8x8
        self._make_conv2d_block(features*4, features*8), # 512x4x4

        self._make_conv2d_block(features*8, 1, 4, 2, 0, use_act=False), # 1x1x1
        nn.Sigmoid(),
    )

def _make_conv2d_block(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, use_bn=True, use_act=True, leak=0.2):
    layers = [
        nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            bias=False,
        ),
    ]
    if use_bn:
        layers.insert(1, nn.BatchNorm2d(out_channels))
    if use_act:
        layers.insert(2, nn.LeakyReLU(leak))
    return nn.Sequential(*layers)

def forward(self, x):
    return self.disc(x)
</code></pre>
<p>class Generator(nn.Module):<br>
def <strong>init</strong>(self, z_dim=100, features=64, img_channels=3): # 100x1x1<br>
super().<strong>init</strong>()<br>
self.gen = nn.Sequential(<br>
self._make_convT2d_block(z_dim, features*16, 4, 1, 0), # 1024x4x4</p>
<pre><code>        self._make_convT2d_block(features*16, features*8), # 512x8x8
        self._make_convT2d_block(features*8, features*4), # 256x16x16
        self._make_convT2d_block(features*4, features*2), # 512x32x32

        self._make_convT2d_block(features*2, img_channels, use_bn=False, use_act=False), # 3x64x64
        nn.Tanh(),
    )

def _make_convT2d_block(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, use_bn=True, use_act=True):
    layers = [
        nn.ConvTranspose2d(
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            bias=False,
        )
    ]
    if use_bn:
        layers.insert(1, nn.BatchNorm2d(out_channels))
    if use_act:
        layers.insert(2, nn.ReLU())
    return nn.Sequential(*layers)        

def forward(self, x):
    return self.gen(x)
</code></pre>
<p>def initialize_weights(model):<br>
for m in model.modules():<br>
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):<br>
nn.init.normal_(<br>
tensor=m.weight.data,<br>
mean=0.0,<br>
std=0.02,<br>
)</p>
<p>BATCH_SIZE, NUM_CHANNELS, H, W = 8, 3, 64, 64<br>
Z_DIM = 100<br>
FEATURES = 64<br>
img_test = torch.randn((BATCH_SIZE, NUM_CHANNELS, H, W))<br>
disc = Discriminator(NUM_CHANNELS, FEATURES)<br>
assert disc(img_test).shape == (BATCH_SIZE, 1, 1, 1), ‘Disc Test Failed’<br>
noise_test = torch.randn((BATCH_SIZE, Z_DIM, 1, 1))<br>
gen = Generator(Z_DIM, FEATURES, NUM_CHANNELS)<br>
assert gen(noise_test).shape == (BATCH_SIZE, NUM_CHANNELS, H, W), ‘Gen Test Failed’<br>
print(‘Success!’)  # Success</p>
<p>device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)<br>
NUM_CHANNELS = 1<br>
BATCH_SIZE = 64<br>
IMG_SIZE = 64<br>
Z_DIM = 100<br>
FEATURES_DISC = 64<br>
FEATURES_GEN = 64<br>
LR = 2e-4<br>
NUM_EPOCHS = 5<br>
DISC_ITER = 1<br>
GRID_SHOW = 8</p>
<p>disc = Discriminator(NUM_CHANNELS, FEATURES_DISC).to(device)<br>
gen = Generator(Z_DIM, FEATURES_GEN, NUM_CHANNELS).to(device)<br>
initialize_weights(disc), initialize_weights(gen)<br>
fixed_noise = torch.randn((GRID_SHOW, Z_DIM, 1, 1)).to(device)</p>
<p>transformations = transforms.Compose(<br>
[<br>
transforms.Resize((IMG_SIZE, IMG_SIZE)),<br>
transforms.ToTensor(),<br>
transforms.Normalize(<br>
[0.5 for _ in range(NUM_CHANNELS)], [0.5 for _ in range(NUM_CHANNELS)]<br>
),<br>
]<br>
)<br>
images = datasets.MNIST(root=‘mnist’, train=True, transform=transformations, download=True)<br>
images_loader = DataLoader(images, batch_size=BATCH_SIZE, shuffle=True)</p>
<p>optim_disc = optim.Adam(disc.parameters(), lr=LR, betas=(0.5, 0.999))<br>
optim_gen = optim.Adam(gen.parameters(), lr=LR, betas=(0.5, 0.999))<br>
criterion = nn.BCELoss()</p>
<p>writer_fake = SummaryWriter(‘logs/fake’)<br>
writer_real = SummaryWriter(‘logs/real’)<br>
steps = 1</p>
<p>disc.train(), gen.train()</p>
<p>for epoch in range(NUM_EPOCHS):<br>
for batch_idx, (real, _) in enumerate(images_loader):<br>
real = real.to(device)<br>
for _ in range(DISC_ITER):<br>
noise = torch.randn((real.shape[0], Z_DIM, 1, 1)).to(device)<br>
fake = gen(noise)<br>
disc_real = disc(real).view(-1)<br>
loss_disc_real = criterion(<br>
input=disc_real,<br>
target=torch.ones_like(disc_real),<br>
)<br>
disc_fake = disc(fake.detach()).view(-1)<br>
loss_disc_fake = criterion(<br>
input=disc_fake,<br>
target=torch.zeros_like(disc_fake),<br>
)<br>
loss_disc = loss_disc_real + loss_disc_fake<br>
disc.zero_grad()<br>
loss_disc.backward()<br>
optim_disc.step()</p>
<pre><code>    disc_fake = disc(fake).view(-1)
    loss_gen = criterion(
        input=disc_fake,
        target=torch.ones_like(disc_fake),
    )
    gen.zero_grad()
    loss_gen.backward()
    optim_gen.step()

    if batch_idx % 50 == 0:
        print(
            f'Epoch: {epoch+1}/{NUM_EPOCHS} -- Step: {steps} -- Batch: {batch_idx+1}/{len(images_loader)} -- Disc Loss: {loss_disc:.4f} -- Gen Loss: {loss_gen:.4f}'
        )
        with torch.no_grad():
            fake = gen(fixed_noise)
            fake_images = torchvision.utils.make_grid(fake, nrow=GRID_SHOW, normalize=True)
            real_images = torchvision.utils.make_grid(real[:GRID_SHOW], nrow=GRID_SHOW, normalize=True)
            writer_fake.add_image(
                'Fake', fake_images, global_step=steps
            )
            writer_real.add_image(
                'Real', real_images, global_step=steps
            )
            steps += 1
</code></pre>
<p>Epoch: 1/5 – Step: 1 – Batch: 1/938 – Disc Loss: 1.3863 – Gen Loss: 0.6932<br>
Epoch: 1/5 – Step: 2 – Batch: 51/938 – Disc Loss: 1.3863 – Gen Loss: 0.6932<br>
Epoch: 1/5 – Step: 3 – Batch: 101/938 – Disc Loss: 1.3863 – Gen Loss: 0.6932<br>
Epoch: 1/5 – Step: 4 – Batch: 151/938 – Disc Loss: 1.3863 – Gen Loss: 0.6931<br>
Epoch: 1/5 – Step: 5 – Batch: 201/938 – Disc Loss: 1.3863 – Gen Loss: 0.6931<br>
Epoch: 1/5 – Step: 6 – Batch: 251/938 – Disc Loss: 1.3863 – Gen Loss: 0.6932<br>
Epoch: 1/5 – Step: 7 – Batch: 301/938 – Disc Loss: 1.3863 – Gen Loss: 0.6931</p>

SinghJivjot: <p>what the actual efff !!!<br>
I had been at the verge of a breakdown coz of this problem, suffering for more than an entire week<br>
So the moral of story is that there was nothing wrong with pytorch routine, it was actually a GAN issue<br>
I just had to remove the BatchNorm2d from the last layer of discriminator and there goes the loss wandering like a ‘—’<br>
GANs are extremely sensitive to architectures and hyper-parameters</p>