-
Notifications
You must be signed in to change notification settings - Fork 39
/
train.py
137 lines (127 loc) · 4.27 KB
/
train.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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import os
import sys
from src.denoising_diffusion_pytorch import GaussianDiffusion
from src.residual_denoising_diffusion_pytorch import (ResidualDiffusion,
Trainer, Unet, UnetRes,
set_seed)
# init
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in [0])
sys.stdout.flush()
set_seed(10)
debug = False
if debug:
save_and_sample_every = 2
sampling_timesteps = 10
sampling_timesteps_original_ddim_ddpm = 10
train_num_steps = 200
else:
save_and_sample_every = 1000
if len(sys.argv) > 1:
sampling_timesteps = int(sys.argv[1])
else:
sampling_timesteps = 10
sampling_timesteps_original_ddim_ddpm = 250
train_num_steps = 100000
original_ddim_ddpm = False
if original_ddim_ddpm:
condition = False
input_condition = False
input_condition_mask = False
else:
condition = False
input_condition = False
input_condition_mask = False
if condition:
# Image restoration
if input_condition:
folder = ["xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_shadow_free_train.flist",
"xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_shadow_train.flist",
"xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_mask_train.flist",
"xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_shadow_free_test.flist",
"xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_shadow_test.flist",
"xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_mask_test.flist"]
else:
folder = ["xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_shadow_free_train.flist",
"xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_shadow_train.flist",
"xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_shadow_free_test.flist",
"xxx/dataset/ISTD_Dataset_arg/data_val/ISTD_shadow_test.flist"]
train_batch_size = 1
num_samples = 1
sum_scale = 0.01
image_size = 256
else:
# Image Generation
folder = 'xxx/CelebA/img_align_celeba'
train_batch_size = 128
num_samples = 64
sum_scale = 1
image_size = 64
num_unet = 2
objective = 'pred_res_noise'
test_res_or_noise = "noise"
if original_ddim_ddpm:
model = Unet(
dim=64,
dim_mults=(1, 2, 4, 8)
)
diffusion = GaussianDiffusion(
model,
image_size=image_size,
timesteps=1000, # number of steps
sampling_timesteps=sampling_timesteps_original_ddim_ddpm,
loss_type='l1', # L1 or L2
)
else:
model = UnetRes(
dim=64,
dim_mults=(1, 2, 4, 8),
num_unet=num_unet,
condition=condition,
input_condition=input_condition,
objective=objective,
test_res_or_noise = test_res_or_noise
)
diffusion = ResidualDiffusion(
model,
image_size=image_size,
timesteps=1000, # number of steps
# number of sampling timesteps (using ddim for faster inference [see citation for ddim paper])
sampling_timesteps=sampling_timesteps,
objective=objective,
loss_type='l2', # L1 or L2
condition=condition,
sum_scale=sum_scale,
input_condition=input_condition,
input_condition_mask=input_condition_mask,
test_res_or_noise = test_res_or_noise
)
trainer = Trainer(
diffusion,
folder,
train_batch_size=train_batch_size,
num_samples=num_samples,
train_lr=2e-4,
train_num_steps=train_num_steps, # total training steps
gradient_accumulate_every=2, # gradient accumulation steps
ema_decay=0.995, # exponential moving average decay
amp=False, # turn on mixed precision
convert_image_to="RGB",
condition=condition,
save_and_sample_every=save_and_sample_every,
equalizeHist=False,
crop_patch=False,
generation=True,
num_unet=num_unet,
)
# train
trainer.train()
# test
if not trainer.accelerator.is_local_main_process:
pass
else:
trainer.load(trainer.train_num_steps//save_and_sample_every)
trainer.set_results_folder(
'./results/test_timestep_'+str(sampling_timesteps))
trainer.test(last=True)
# trainer.set_results_folder('./results/test_sample')
# trainer.test(sample=True)