Initialize the GANLoss class. Parameters: gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp. target_real_label (bool) - - label for a real image target_fake_label (bool) - - label of a fake image Note
(self, gan_mode, target_real_label=1.0, target_fake_label=0.0)
| 222 | """ |
| 223 | |
| 224 | def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0): |
| 225 | """ Initialize the GANLoss class. |
| 226 | |
| 227 | Parameters: |
| 228 | gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp. |
| 229 | target_real_label (bool) - - label for a real image |
| 230 | target_fake_label (bool) - - label of a fake image |
| 231 | |
| 232 | Note: Do not use sigmoid as the last layer of Discriminator. |
| 233 | LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss. |
| 234 | """ |
| 235 | super(GANLoss, self).__init__() |
| 236 | self.register_buffer('real_label', torch.tensor(target_real_label)) |
| 237 | self.register_buffer('fake_label', torch.tensor(target_fake_label)) |
| 238 | self.gan_mode = gan_mode |
| 239 | if gan_mode == 'lsgan': |
| 240 | self.loss = nn.MSELoss() |
| 241 | elif gan_mode == 'vanilla': |
| 242 | self.loss = nn.BCEWithLogitsLoss() |
| 243 | elif gan_mode in ['wgangp']: |
| 244 | self.loss = None |
| 245 | else: |
| 246 | raise NotImplementedError('gan mode %s not implemented' % gan_mode) |
| 247 | |
| 248 | def get_target_tensor(self, prediction, target_is_real): |
| 249 | """Create label tensors with the same size as the input. |