Spaces:
Running on Zero
Running on Zero
| # Copyright 2021 Tomoki Hayashi | |
| # MIT License (https://opensource.org/licenses/MIT) | |
| # Adapted by Florian Lux 2021 | |
| import torch | |
| import torch.nn.functional as F | |
| class FeatureMatchLoss(torch.nn.Module): | |
| def __init__(self, | |
| average_by_layers=False, | |
| 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): | |
| """ | |
| Calcualate feature matching loss. | |
| Args: | |
| feats_hat (list): List of list of discriminator outputs | |
| calcuated from generator outputs. | |
| feats (list): List of list of discriminator outputs | |
| calcuated from groundtruth. | |
| 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_ = 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 + feat_match_loss_ | |
| if self.average_by_discriminators: | |
| feat_match_loss /= i + 1 | |
| return feat_match_loss | |