Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| # Copyright 2021 Tomoki Hayashi | |
| # MIT License (https://opensource.org/licenses/MIT) | |
| """Feature matching loss modules.""" | |
| import torch | |
| import torch.nn.functional as F | |
| class FeatureMatchLoss(torch.nn.Module): | |
| """Feature matching loss module.""" | |
| def __init__( | |
| self, | |
| average_by_layers=True, | |
| average_by_discriminators=True, | |
| include_final_outputs=False, | |
| ): | |
| """Initialize FeatureMatchLoss module.""" | |
| 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 generater 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_ += 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 | |