Here’s My training loop code…
for epoch in range(epochs):
for count, data in enumerate(data_list):
train_image = []
sample_image = []
for image in data:
img_train = cv2.imread(path + image).T/255
img_train = img_train.reshape(1, img_train.shape[0], img_train.shape[1], img_train.shape[2])
img_sample = cv2.imread(path2 + image,0).T/255
img_sample = img_sample.reshape(1, 1, img_sample.shape[0], img_sample.shape[1])
train_image.append(img_train)
sample_image.append(img_sample)
assert(img_sample.shape == (1, 1, 144, 144))
train_image = torch.from_numpy(np.concatenate(train_image, axis = 0)).type('torch.cuda.FloatTensor')
sample_image = torch.from_numpy(np.concatenate(sample_image, axis = 0)).type('torch.cuda.FloatTensor')
label = torch.full((train_image.shape[0],), real_label, device=device)
#Training the discriminator... minimizing -(log(D(x)) + log(1 - D(G(Z))))
dis.zero_grad()
Gen.zero_grad()
G_z = Gen(sample_image.detach())
disc_real_out = dis(train_image.detach()).view(-1)
error_real = GAN_loss(disc_real_out, label)
error_real.backward()
disc_fake_out = dis(G_z.detach()).view(-1)
label.fill_(fake_label)
error_fake = GAN_loss(disc_fake_out, label)
error_fake.backward()
total_disc_error = error_real + error_fake
D_optimizer.step()
#Training the Generator... maximizing log(D(G(Z))))
D_G_z = dis(G_z.detach()).view(-1)
label.fill_(real_label)
error_gen = GAN_loss(D_G_z, label)
error_gen.backward()
G_optimizer.step()
G_losses.append(error_gen.item())
D_losses.append(total_disc_error.item())
print("Discriminator Loss : ", total_disc_error.item(), "\t", "Generator Loss : ", total_gen_error.item())
```