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

EMA for vqgan vae comes from ema_pytorch now

Browse Source
This commit is contained in:
Phil Wang
2022-06-20 15:29:08 -07:00
parent 0021535c26
commit c098f57e09
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
dalle2_pytorch/vqgan_vae_trainer.py
View File

@@ -16,10 +16,11 @@ from torchvision.utils import make_grid, save_image
from einops import rearrange from einops import rearrange
from dalle2_pytorch.train import EMA
from dalle2_pytorch.vqgan_vae import VQGanVAE from dalle2_pytorch.vqgan_vae import VQGanVAE
from dalle2_pytorch.optimizer import get_optimizer from dalle2_pytorch.optimizer import get_optimizer
from ema_pytorch import EMA
# helpers # helpers
def exists(val): def exists(val):
@@ -97,7 +98,7 @@ class VQGanVAETrainer(nn.Module):
valid_frac = 0.05, valid_frac = 0.05,
random_split_seed = 42, random_split_seed = 42,
ema_beta = 0.995, ema_beta = 0.995,
ema_update_after_step = 2000, ema_update_after_step = 500,
ema_update_every = 10, ema_update_every = 10,
apply_grad_penalty_every = 4, apply_grad_penalty_every = 4,
amp = False amp = False