| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class GANLoss(nn.Module): | |
| """ | |
| Define different GAN objectives. | |
| The GANLoss class abstracts away the need to create the target label tensor | |
| that has the same size as the input. | |
| """ | |
| def __init__(self, loss_mode="vanilla", real_label=1.0, fake_label=0.0): | |
| """ | |
| --------- | |
| Arguments | |
| --------- | |
| loss_mode : str | |
| GAN loss mode (default="vanilla") | |
| real_label : bool | |
| label for real image | |
| fake_label : bool | |
| label for fake image | |
| """ | |
| super().__init__() | |
| self.loss_mode = loss_mode | |
| self.register_buffer("real_label", torch.tensor(real_label)) | |
| self.register_buffer("fake_label", torch.tensor(fake_label)) | |
| self.loss = None | |
| if self.loss_mode == "vanilla": | |
| self.loss = nn.BCEWithLogitsLoss() | |
| else: | |
| raise NotImplementedError( | |
| f"GANLoss with {self.loss_mode} mode - not implemented yet" | |
| ) | |
| def get_target_tensor(self, prediction, target_is_real): | |
| """ | |
| --------- | |
| Arguments | |
| --------- | |
| prediction : tensor | |
| prediction from a discriminator | |
| target_is_real : bool | |
| whether the groundtruth label is for a real image or a fake image | |
| ------- | |
| Returns | |
| ------- | |
| tensor : A label tensor filled with groundtruth label with the same size as that of input | |
| """ | |
| if target_is_real: | |
| target_tensor = self.real_label | |
| else: | |
| target_tensor = self.fake_label | |
| return target_tensor.expand_as(prediction) | |
| def __call__(self, prediction, target_is_real): | |
| """ | |
| --------- | |
| Arguments | |
| --------- | |
| prediction : tensor | |
| prediction from a discriminator | |
| target_is_real : bool | |
| whether the groundtruth label is for a real image or a fake image | |
| ------- | |
| Returns | |
| ------- | |
| loss : the computed loss | |
| """ | |
| if self.loss_mode == "vanilla": | |
| target_tensor = self.get_target_tensor(prediction, target_is_real) | |
| loss = self.loss(prediction, target_tensor) | |
| else: | |
| loss = 0 | |
| return loss | |