| |
| |
| |
|
|
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def feature_loss(fmap_r, fmap_g): |
| loss = 0 |
| for dr, dg in zip(fmap_r, fmap_g): |
| loss += torch.mean(torch.abs(dr - dg)) |
|
|
| return loss / len(fmap_g) |
|
|
|
|
| class FeatureMatchLoss(torch.nn.Module): |
|
|
| def __init__(self, |
| average_by_layers=True, |
| average_by_discriminators=False, |
| include_final_outputs=False, ): |
| super().__init__() |
| self.average_by_layers = average_by_layers |
| self.average_by_discriminators = average_by_discriminators |
| self.include_final_outputs = include_final_outputs |
|
|
| def forward(self, feats_hat, feats): |
| """ |
| Calculate feature matching loss. |
| |
| Args: |
| feats_hat (list): List of lists of discriminator outputs |
| calculated from generator outputs. |
| feats (list): List of lists of discriminator outputs |
| calculated from ground-truth. |
| |
| Returns: |
| Tensor: Feature matching loss value. |
| """ |
| feat_match_loss = 0.0 |
| for i, (feats_hat_, feats_) in enumerate(zip(feats_hat, feats)): |
| feat_match_loss_ = 0.0 |
| if not self.include_final_outputs: |
| feats_hat_ = feats_hat_[:-1] |
| feats_ = feats_[:-1] |
| for j, (feat_hat_, feat_) in enumerate(zip(feats_hat_, feats_)): |
| feat_match_loss_ += F.l1_loss(feat_hat_, feat_.detach()) |
| if self.average_by_layers: |
| feat_match_loss_ /= j + 1 |
| feat_match_loss += feat_match_loss_ |
| if self.average_by_discriminators: |
| feat_match_loss /= i + 1 |
|
|
| return feat_match_loss |
|
|
