-
Notifications
You must be signed in to change notification settings - Fork 0
/
difusion.py
30 lines (26 loc) · 949 Bytes
/
difusion.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
model = Unet(
dim = 64,
dim_mults = (1, 2),
channels = 1
)
diffusion = GaussianDiffusion(
model,
image_size = 84,
# channels = 1,
timesteps = 1000, # number of steps
sampling_timesteps = 250, # number of sampling timesteps (using ddim for faster inference [see citation for ddim paper])
loss_type = 'l1' # L1 or L2
)
trainer = Trainer(
diffusion,
'./pong_pic',
train_batch_size = 16,
train_lr = 1e-4,
train_num_steps = 200000, # total training steps
gradient_accumulate_every = 1, # gradient accumulation steps
ema_decay = 0.995, # exponential moving average decay
amp = False, # turn on mixed precision
calculate_fid = True # whether to calculate fid during training
)
trainer.train()