I'm new to GANs and I have trouble in training DCGAN on mnist, when I was working with GAN with linear layers everything was fine and the generator generates pretty good images. But when I started working with Convolutional GAN, the generator generators only noise. Does anybody know how to fix this issue?
Discriminator loss and generator loss
Here are my Neural network and my training loop,
Neural network :
Discriminator
class Discriminator(nn.Module):
def __init__(self, ch, F):
super(Discriminator, self).__init__()
def block(in_ch,out_ch,k,s,p,final=False,bn=True):
block = []
block.append(nn.Conv2d(in_ch,out_ch,kernel_size=k,stride=s,padding=p))
if not final and bn:
block.append(nn.BatchNorm2d(out_ch))
block.append(nn.LeakyReLU(0.2))
elif not bn and not final:
block.append(nn.LeakyReLU(0.2))
elif final and not bn:
block.append(nn.Sigmoid())
return block
self.D = nn.Sequential(
*block(ch,F,k=4,s=2,p=1,bn=False),
*block(F,F*2,k=4,s=2,p=1),
*block(F*2,F*4,k=4,s=2,p=1),
*block(F*4,F*8,k=4,s=2,p=1),
*block(F*8,1,k=4,s=2,p=0,final=True,bn=False)
)
def forward(self,x): return self.D(x)
Generator
class Generator(nn.Module):
def __init__(self, ch_noise, ch_img, features_g):
super(Generator, self).__init__()
def block(in_ch,out_ch,k,s,p,final=False):
block = []
block.append(nn.ConvTranspose2d(in_ch,out_ch,kernel_size=k,stride=s,padding=p))
if not final:
block.append(nn.BatchNorm2d(out_ch))
block.append(nn.ReLU())
if final:
block.append(nn.Tanh())
return block
self.G = nn.Sequential(
*block(ch_noise, features_g*16, k=4,s=1,p=0),
*block(features_g*16, features_g*8, k=4, s=2, p=1),
*block(features_g*8, features_g*4, k=4, s=2, p=1),
*block(features_g*4, features_g*2, k=4, s=2, p=1),
*block(features_g*2, 1, k=4, s=2, p=1,final=True)
)
def forward(self,z): return self.G(z)
Training Loop:
loss_G, loss_D = [],[]
for i in range(epochs):
D.train()
G.train()
st = time.time()
for idx, (img, _ ) in enumerate(mnist):
img = img.to(device)
## Discriminator ##
D.zero_grad(set_to_none=True)
lable = torch.ones(bs,device=device)*0.9
pred = D(img).reshape(-1)
loss_d_real = criterion(pred,lable)
z = torch.randn(img.shape[0],ch_z, 1,1,device=device)
fake_img = G(z)
lable = torch.ones(bs,device=device)*0.1
pred = D(fake_img.detach()).reshape(-1)
loss_d_fake = criterion(pred,lable)
D_loss = loss_d_real + loss_d_fake
D_loss.backward()
optim_d.step()
## Generator ##
G.zero_grad(True)
lable = torch.randn(bs,device=device)
pred = D(fake_img).reshape(-1)
G_loss = criterion(pred,lable)
G_loss.backward()
optim_g.step()
## printing on terminal
if idx % 100 == 0:
print(f'\nBatches done : {idx}/{len(mnist)}')
print(f'Loss_D : {D_loss.item():.4f}\tLoss_G : {G_loss.item():.4f}')
et = time.time()
print(f'\nEpoch : {i+1}\n{time_cal(st,et)}')
G.eval()
with torch.no_grad():
fake_image = G(fixed_noise)
save_image(fake_image[:25],fp=f'{path_to_img}/{i+1}_fake.png',nrow=5,normalize=True)
loss_G.append(G_loss.item())
loss_D.append(D_loss.item())
Here's the link to my Colab notebook in which I have been working.