Image-to-Image Translation Using Conditional DCGANs

train.py

#G(x,z), D(x,y)
factor=1
G.train()
D.train()
for epoch in range(50):
  for i,(x,y) in enumerate(dataloader):
    opt_D.zero_grad()
    opt_G.zero_grad()

    x=x.to(device)
    y=y.to(device)
    if i%factor==0:
      #train Discriminator
      #for real data
      loss_real=criterion(D(x,y),torch.ones((batch_size,)).to(device)*0.9)
      
      #for fake data
      z=torch.randn((batch_size,1,128,128)).to(device)
      loss_fake=criterion(D(x,G(x,z).detach()),torch.ones((batch_size,)).to(device)*0.1)

      loss_D=loss_real+loss_fake
      loss_D.backward()
      opt_D.step()

    #train Generator
    z=torch.randn((batch_size,1,128,128)).to(device)
    loss_G=criterion(D(x,G(x,z)),torch.ones((batch_size,)).to(device))
    loss_G.backward()
    opt_G.step()