Spaces:
Running
Running
File size: 1,318 Bytes
9d61c9b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
from typing import List
import torch
from torch import Tensor
from torch.nn.modules.loss import _Loss
class DiscriminatorLoss(_Loss):
"""Discriminator Loss module"""
def forward(
self,
disc_real_outputs: List[Tensor],
disc_generated_outputs: List[Tensor],
):
loss = 0
r_losses = []
g_losses = []
for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
r_loss = torch.mean((1 - dr) ** 2)
g_loss = torch.mean(dg**2)
loss += r_loss + g_loss
r_losses.append(r_loss.item())
g_losses.append(g_loss.item())
return loss, r_losses, g_losses
class FeatureMatchingLoss(_Loss):
"""Feature Matching Loss module"""
def forward(self, fmap_r: List[Tensor], fmap_g: List[Tensor]):
loss = 0
for dr, dg in zip(fmap_r, fmap_g):
for rl, gl in zip(dr, dg):
loss += torch.mean(torch.abs(rl - gl))
return loss * 2
class GeneratorLoss(_Loss):
"""Generator Loss module"""
def forward(self, disc_outputs: List[Tensor]):
loss = 0
gen_losses = []
for dg in disc_outputs:
l = torch.mean((1 - dg) ** 2)
gen_losses.append(l)
loss += l
return loss, gen_losses
|