-
Notifications
You must be signed in to change notification settings - Fork 0
/
ecdm_second_stage.yaml
115 lines (112 loc) · 2.84 KB
/
ecdm_second_stage.yaml
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
model:
class_path: ecdm.models.diffusion.ecdm_second_stage.ECDMSecondStage
init_args:
learning_rate: 4.5e-6
linear_start: 0.0015
linear_end: 0.0155
log_every_t: 200
timesteps: 1000
loss_type: l2
first_stage_key: "tir_img"
cond_stage_key: "tir_edge"
image_size: [512,640]
channels: 3
conditioning_key: concat
monitor: 'val/loss_ema'
first_stage_ckpt: "checkpoints/last.ckpt"
scheduler_config: # 100 warmup steps
class_path: ecdm.lr_scheduler.LambdaLinearSchedulerWithMultiStep
init_args:
warm_up_steps: [100]
cycle_lengths: [10000000000000]
f_start: [0.5e-6]
f_max: [1.]
f_min: [1.]
milestones: [7500,15000]
gamma: 0.1
unet_config:
class_path: ecdm.modules.diffusionmodules.simple_unet.SimpleUNet
init_args:
image_size: [512,640]
in_channels: 6
out_channels: 3
model_channels: 128
attention_resolutions: [36]
num_res_blocks: 2
channel_mult: [1, 2, 2, 2]
dropout: 0.1
discriminator_config:
norm_type: "instance"
return_binary: True
num_channels: 3
num_channels_d: 64
kernel_size_d: 4
initializer: "normal"
sample_config:
sample_method: "dpm"
sample_args:
noise_schedule_type: 'discrete'
timesteps: 5
dpm_solver_order: 3
skip_type: "time_uniform"
dpm_solver_method: "singlestep"
lower_order_final: False
dpm_solver_type: "taylor"
dpm_solver_atol: 0.0078
dpm_solver_rtol: 0.05
scale: 0.5
denoise: False
thresholding: False
# sample_config:
# sample_method: "ddim"
# steps: 200
data:
class_path: ecdm.data.dataset.DataModuleFromConfig
init_args:
batch_size: 1
num_workers: 23
wrap: True
train:
target: ecdm.data.dataset.LLVIPDataset
params:
root: data/LLVIP
size:
- 512
- 640
validation:
target: ecdm.data.dataset.LLVIPDataset
params:
root: data/LLVIP
mode: test
size:
- 512
- 640
test:
target: ecdm.data.dataset.LLVIPDataset
params:
root: data/LLVIP
mode: test
size:
- 512
- 640
trainer:
num_sanity_val_steps: 0
benchmark: True
# limit_train_batches: 0.1
# limit_val_batches: 0.1
# limit_test_batches: 0.01
max_epochs: 30
callbacks:
- class_path: ecdm.callbacks.ImageLogger
init_args:
train_batch_frequency: 1000
val_batch_frequency: 200
max_images: 4
increase_log_steps: False
log_on_batch_idx: True
logger:
class_path: lightning.pytorch.loggers.tensorboard.TensorBoardLogger
init_args:
save_dir: logs
checkpoint_callback:
monitor: 'val/loss_simple_ema'