aljaz/DALLE2-pytorch
1
1
Fork 0
mirror of https://github.com/lucidrains/DALLE2-pytorch.git synced 2025-12-20 10:14:19 +01:00
Code Issues Packages Projects Releases Wiki Activity

prepare for turning off gradient penalty, as shown in GAN literature, GP needs to be only applied 1 out of 4 iterations

Browse Source
This commit is contained in:
Phil Wang
2022-04-23 07:52:10 -07:00
parent 05b74be69a
commit f82917e1fd
2 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
dalle2_pytorch/vqgan_vae.py
View File

@@ -477,7 +477,8 @@ class VQGanVAE(nn.Module):
img,
return_loss = False,
return_discr_loss = False,
return_recons = False
return_recons = False,
add_gradient_penalty = True
):
batch, channels, height, width, device = *img.shape, img.device
assert height == self.image_size and width == self.image_size, 'height and width of input image must be equal to {self.image_size}'
@@ -502,11 +503,11 @@ class VQGanVAE(nn.Module):
fmap_discr_logits, img_discr_logits = map(self.discr, (fmap, img))
gp = gradient_penalty(img, img_discr_logits)
discr_loss = self.discr_loss(fmap_discr_logits, img_discr_logits)
loss = discr_loss + gp
if add_gradient_penalty:
gp = gradient_penalty(img, img_discr_logits)
loss = discr_loss + gp
if return_recons:
return loss, fmap