https-deeplearning-ai / gans-public Goto Github PK

from torchvision.models import vgg19

class Loss(nn.Module):

     ...

    def adv_loss(self, x, is_real):
        target = torch.zeros_like(x) if is_real else torch.ones_like(x)
        return F.binary_cross_entropy_with_logits(x, target)

    def forward(self, generator, discriminator, hr_real, lr_real):
        ''' Performs forward pass and returns total losses for G and D '''
        hr_fake = generator(lr_real)
        fake_preds_for_g = discriminator(hr_fake)
        fake_preds_for_d = discriminator(hr_fake.detach())
        real_preds_for_d = discriminator(hr_real.detach())

        g_loss = (
            0.001 * self.adv_loss(fake_preds_for_g, False) + \
            0.006 * self.vgg_loss(hr_real, hr_fake) + \
            self.img_loss(hr_real, hr_fake)
        )
        d_loss = 0.5 * (
            self.adv_loss(real_preds_for_d, True) + \
            self.adv_loss(fake_preds_for_d, False)
        )

        return g_loss, d_loss, hr_fake

        d_loss = 0.5 * (
            self.adv_loss(real_preds_for_d, False) + \
            self.adv_loss(fake_preds_for_d, True)
        )

    def adv_loss(self, x, is_real):
        target = torch.ones_like(x) if is_real else torch.zeros_like(x)
        return F.binary_cross_entropy_with_logits(x, target)

        g_loss = (
            0.001 * self.adv_loss(fake_preds_for_g, True) + \
            0.006 * self.vgg_loss(hr_real, hr_fake) + \
            self.img_loss(hr_real, hr_fake)
        )