def Loss_on_layer_GAN(self, image, finding, impression, decoder, D):
'''
Pretrain genertaor with batch
:image image batch
:text text batch
'''
image1 = F.interpolate(image, size=(64, 64))
image2 = F.interpolate(image, size=(128, 128))
image3 = F.interpolate(image, size=(256, 256))
real_img = [image1,image2,image3]
bsz = self.batch_size
#hidden = self.encoder.init_hidden(bsz)
words_emb, sent_emb = self.encoder(finding, impression)
words_emb, sent_emb = words_emb.detach(), sent_emb.detach()
mask1 = (finding == 0)
mask2 = (impression == 0)
mask = torch.cat((mask1, mask2), 1)
num_words = words_emb.size(2)
if mask.size(1) > num_words:
mask = mask[:, :num_words]
pre_image, _, mu, logvar = decoder(sent_emb, words_emb, mask)
# Train Discriminator
self.D_optimizer.zero_grad()
total_D_loss = 0
for i in range(3):
pre_fake = D[i](pre_image[i], mu.detach())
pre_real = D[i](real_img[i], mu.detach())
gradient_penalty, gradients = cal_gradient_penalty(D[i], real_img[i], pre_image[i], mu.detach(), "cuda")
D_loss = pre_fake.mean() - pre_real.mean() + gradient_penalty
total_D_loss += D_loss
D_loss.backward(retain_graph=True)
self.D_optimizer.step()
# Train Generator
self.G_optimizer.zero_grad()
G_loss = 0
for i in range(3):
pre_fake = D[i](pre_image[i], mu.detach())
adv_loss = - self.adv_loss_ratio * pre_fake.mean()
adv_loss.backward(retain_graph=True)
content_loss = self.pix_loss_ratio * self.G_criterion(pre_image[i].float(),
real_img[i].float())
content_loss.backward(retain_graph=True)
G_loss += content_loss
G_loss += adv_loss
self.G_optimizer.step()
return total_D_loss, G_loss, pre_image[-1], real_img[-1]
def Loss_on_layer_GAN(self, image, finding, impression, decoder, D):
'''
Pretrain genertaor with batch
:image image batch
:text text batch
'''
image1 = F.interpolate(image, size=(64, 64))
image2 = F.interpolate(image, size=(128, 128))
image3 = F.interpolate(image, size=(256, 256))
real_img = [image1, image2, image3]
txt_emded, hidden = self.encoder(finding, impression)
pre_image, mu, logvar = decoder(txt_emded)
# Train Discriminator
self.D_optimizer.zero_grad()
total_D_loss = 0
for i in range(3):
pre_fake = D[i](pre_image[i], mu.detach())
pre_real = D[i](real_img[i], mu.detach())
gradient_penalty, gradients = cal_gradient_penalty(
D[i], real_img[i], pre_image[i], mu.detach(), "cuda")
D_loss = pre_fake.mean() - pre_real.mean() + gradient_penalty
total_D_loss += D_loss
D_loss.backward(retain_graph=True)
self.D_optimizer.step()
# Train Generator
self.G_optimizer.zero_grad()
G_loss = 0
for i in range(3):
pre_fake = D[i](pre_image[i], mu.detach())
adv_loss = -self.adv_loss_ratio * pre_fake.mean()
adv_loss.backward(retain_graph=True)
content_loss = self.pix_loss_ratio * self.G_criterion(
pre_image[i].float(), real_img[i].float())
content_loss.backward(retain_graph=True)
G_loss += content_loss
G_loss += adv_loss
self.G_optimizer.step()
return total_D_loss, G_loss, pre_image[-1], real_img[-1]
def Loss_on_layer_GAN(self, image, finding, impression, decoder, D):
'''
Pretrain genertaor with batch
:image image batch
:text text batch
'''
txt_emded, hidden = self.encoder(finding, impression)
pre_image = decoder(txt_emded)
# Train Discriminator
self.D_optimizer.zero_grad()
pre_fake = D(pre_image, txt_emded)
pre_real = D(image, txt_emded)
gradient_penalty, gradients = cal_gradient_penalty(
D, image, pre_image, txt_emded, "cuda")
D_loss = pre_fake.mean() - pre_real.mean() + gradient_penalty
D_loss.backward(retain_graph=True)
self.D_optimizer.step()
# Train Generator
self.G_optimizer.zero_grad()
pre_fake = D(pre_image, txt_emded)
adv_loss = -self.adv_loss_ratio * pre_fake.mean()
adv_loss.backward(retain_graph=True)
content_loss = self.pix_loss_ratio * self.G_criterion(
pre_image.float(), image.float())
content_loss.backward(retain_graph=True)
G_loss = content_loss + adv_loss
self.G_optimizer.step()
return D_loss, G_loss, pre_image, image