mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-10-31 03:20:40 +00:00
refactor: remove lr_scheduler.py
This commit is contained in:
Bryce
Bryce Drennan
parent
01e32ff3f6
commit
0c456cd52a
@ -1,70 +0,0 @@
|
|||||||
model:
|
|
||||||
base_learning_rate: 1.0e-04
|
|
||||||
target: imaginairy.modules.diffusion.ddpm.LatentDiffusion
|
|
||||||
params:
|
|
||||||
linear_start: 0.00085
|
|
||||||
linear_end: 0.0120
|
|
||||||
num_timesteps_cond: 1
|
|
||||||
log_every_t: 200
|
|
||||||
timesteps: 1000
|
|
||||||
first_stage_key: "edited"
|
|
||||||
cond_stage_key: "edit"
|
|
||||||
image_size: 16
|
|
||||||
channels: 4
|
|
||||||
cond_stage_trainable: false
|
|
||||||
conditioning_key: hybrid
|
|
||||||
monitor: val/loss_simple_ema
|
|
||||||
scale_factor: 0.18215
|
|
||||||
use_ema: false
|
|
||||||
|
|
||||||
scheduler_config:
|
|
||||||
target: imaginairy.lr_scheduler.LambdaLinearScheduler
|
|
||||||
params:
|
|
||||||
warm_up_steps: [ 0 ]
|
|
||||||
cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
|
|
||||||
f_start: [ 1.e-6 ]
|
|
||||||
f_max: [ 1. ]
|
|
||||||
f_min: [ 1. ]
|
|
||||||
|
|
||||||
unet_config:
|
|
||||||
target: imaginairy.modules.diffusion.openaimodel.UNetModel
|
|
||||||
params:
|
|
||||||
use_checkpoint: True
|
|
||||||
image_size: 32 # unused
|
|
||||||
in_channels: 8
|
|
||||||
out_channels: 4
|
|
||||||
model_channels: 320
|
|
||||||
attention_resolutions: [ 4, 2, 1 ]
|
|
||||||
num_res_blocks: 2
|
|
||||||
channel_mult: [ 1, 2, 4, 4 ]
|
|
||||||
num_heads: 8
|
|
||||||
use_spatial_transformer: True
|
|
||||||
transformer_depth: 1
|
|
||||||
context_dim: 768
|
|
||||||
legacy: False
|
|
||||||
|
|
||||||
first_stage_config:
|
|
||||||
target: imaginairy.modules.autoencoder.AutoencoderKL
|
|
||||||
params:
|
|
||||||
embed_dim: 4
|
|
||||||
monitor: val/rec_loss
|
|
||||||
ddconfig:
|
|
||||||
double_z: true
|
|
||||||
z_channels: 4
|
|
||||||
resolution: 256
|
|
||||||
in_channels: 3
|
|
||||||
out_ch: 3
|
|
||||||
ch: 128
|
|
||||||
ch_mult:
|
|
||||||
- 1
|
|
||||||
- 2
|
|
||||||
- 4
|
|
||||||
- 4
|
|
||||||
num_res_blocks: 2
|
|
||||||
attn_resolutions: []
|
|
||||||
dropout: 0.0
|
|
||||||
lossconfig:
|
|
||||||
target: torch.nn.Identity
|
|
||||||
|
|
||||||
cond_stage_config:
|
|
||||||
target: imaginairy.modules.clip_embedders.FrozenCLIPEmbedder
|
|
||||||
@ -1,71 +0,0 @@
|
|||||||
model:
|
|
||||||
base_learning_rate: 7.5e-05
|
|
||||||
target: imaginairy.modules.diffusion.ddpm.LatentInpaintDiffusion
|
|
||||||
params:
|
|
||||||
linear_start: 0.00085
|
|
||||||
linear_end: 0.0120
|
|
||||||
num_timesteps_cond: 1
|
|
||||||
log_every_t: 200
|
|
||||||
timesteps: 1000
|
|
||||||
first_stage_key: "jpg"
|
|
||||||
cond_stage_key: "txt"
|
|
||||||
image_size: 64
|
|
||||||
channels: 4
|
|
||||||
cond_stage_trainable: false # Note: different from the one we trained before
|
|
||||||
conditioning_key: hybrid # important
|
|
||||||
monitor: val/loss_simple_ema
|
|
||||||
scale_factor: 0.18215
|
|
||||||
finetune_keys: null
|
|
||||||
use_ema: False
|
|
||||||
|
|
||||||
scheduler_config: # 10000 warm-up steps
|
|
||||||
target: ldm.lr_scheduler.LambdaLinearScheduler
|
|
||||||
params:
|
|
||||||
warm_up_steps: [ 2500 ] # NOTE for resuming. use 10000 if starting from scratch
|
|
||||||
cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
|
|
||||||
f_start: [ 1.e-6 ]
|
|
||||||
f_max: [ 1. ]
|
|
||||||
f_min: [ 1. ]
|
|
||||||
|
|
||||||
unet_config:
|
|
||||||
target: imaginairy.modules.diffusion.openaimodel.UNetModel
|
|
||||||
params:
|
|
||||||
image_size: 32 # unused
|
|
||||||
in_channels: 9 # 4 data + 4 downscaled image + 1 mask
|
|
||||||
out_channels: 4
|
|
||||||
model_channels: 320
|
|
||||||
attention_resolutions: [ 4, 2, 1 ]
|
|
||||||
num_res_blocks: 2
|
|
||||||
channel_mult: [ 1, 2, 4, 4 ]
|
|
||||||
num_heads: 8
|
|
||||||
use_spatial_transformer: True
|
|
||||||
transformer_depth: 1
|
|
||||||
context_dim: 768
|
|
||||||
use_checkpoint: True
|
|
||||||
legacy: False
|
|
||||||
|
|
||||||
first_stage_config:
|
|
||||||
target: imaginairy.modules.autoencoder.AutoencoderKL
|
|
||||||
params:
|
|
||||||
embed_dim: 4
|
|
||||||
monitor: val/rec_loss
|
|
||||||
ddconfig:
|
|
||||||
double_z: true
|
|
||||||
z_channels: 4
|
|
||||||
resolution: 256
|
|
||||||
in_channels: 3
|
|
||||||
out_ch: 3
|
|
||||||
ch: 128
|
|
||||||
ch_mult:
|
|
||||||
- 1
|
|
||||||
- 2
|
|
||||||
- 4
|
|
||||||
- 4
|
|
||||||
num_res_blocks: 2
|
|
||||||
attn_resolutions: []
|
|
||||||
dropout: 0.0
|
|
||||||
lossconfig:
|
|
||||||
target: torch.nn.Identity
|
|
||||||
|
|
||||||
cond_stage_config:
|
|
||||||
target: imaginairy.modules.clip_embedders.FrozenCLIPEmbedder
|
|
||||||
@ -1,70 +0,0 @@
|
|||||||
model:
|
|
||||||
base_learning_rate: 1.0e-4
|
|
||||||
target: imaginairy.modules.diffusion.ddpm.LatentDiffusion
|
|
||||||
params:
|
|
||||||
linear_start: 0.00085
|
|
||||||
linear_end: 0.0120
|
|
||||||
num_timesteps_cond: 1
|
|
||||||
log_every_t: 200
|
|
||||||
timesteps: 1000
|
|
||||||
first_stage_key: "image"
|
|
||||||
cond_stage_key: "txt"
|
|
||||||
image_size: 64
|
|
||||||
channels: 4
|
|
||||||
cond_stage_trainable: false
|
|
||||||
conditioning_key: crossattn
|
|
||||||
monitor: val/loss_simple_ema
|
|
||||||
scale_factor: 0.18215
|
|
||||||
use_ema: False
|
|
||||||
|
|
||||||
scheduler_config: # 10000 warm-up steps
|
|
||||||
target: imaginairy.lr_scheduler.LambdaLinearScheduler
|
|
||||||
params:
|
|
||||||
warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch
|
|
||||||
cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
|
|
||||||
f_start: [ 1.e-6 ]
|
|
||||||
f_max: [ 1. ]
|
|
||||||
f_min: [ 1. ]
|
|
||||||
|
|
||||||
unet_config:
|
|
||||||
target: imaginairy.modules.diffusion.openaimodel.UNetModel
|
|
||||||
params:
|
|
||||||
use_checkpoint: True
|
|
||||||
image_size: 32 # unused
|
|
||||||
in_channels: 4
|
|
||||||
out_channels: 4
|
|
||||||
model_channels: 320
|
|
||||||
attention_resolutions: [ 4, 2, 1 ]
|
|
||||||
num_res_blocks: 2
|
|
||||||
channel_mult: [ 1, 2, 4, 4 ]
|
|
||||||
num_heads: 8
|
|
||||||
use_spatial_transformer: True
|
|
||||||
transformer_depth: 1
|
|
||||||
context_dim: 768
|
|
||||||
legacy: False
|
|
||||||
|
|
||||||
first_stage_config:
|
|
||||||
target: imaginairy.modules.autoencoder.AutoencoderKL
|
|
||||||
params:
|
|
||||||
embed_dim: 4
|
|
||||||
monitor: val/rec_loss
|
|
||||||
ddconfig:
|
|
||||||
double_z: true
|
|
||||||
z_channels: 4
|
|
||||||
resolution: 256
|
|
||||||
in_channels: 3
|
|
||||||
out_ch: 3
|
|
||||||
ch: 128
|
|
||||||
ch_mult:
|
|
||||||
- 1
|
|
||||||
- 2
|
|
||||||
- 4
|
|
||||||
- 4
|
|
||||||
num_res_blocks: 2
|
|
||||||
attn_resolutions: []
|
|
||||||
dropout: 0.0
|
|
||||||
lossconfig:
|
|
||||||
target: torch.nn.Identity
|
|
||||||
|
|
||||||
cond_stage_config:
|
|
||||||
target: imaginairy.modules.clip_embedders.FrozenCLIPEmbedder
|
|
||||||
@ -1,132 +0,0 @@
|
|||||||
"""Classes for learning rate scheduling"""
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
|
|
||||||
class LambdaWarmUpCosineScheduler:
|
|
||||||
"""
|
|
||||||
note: use with a base_lr of 1.0.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
warm_up_steps,
|
|
||||||
lr_min,
|
|
||||||
lr_max,
|
|
||||||
lr_start,
|
|
||||||
max_decay_steps,
|
|
||||||
verbosity_interval=0,
|
|
||||||
):
|
|
||||||
self.lr_warm_up_steps = warm_up_steps
|
|
||||||
self.lr_start = lr_start
|
|
||||||
self.lr_min = lr_min
|
|
||||||
self.lr_max = lr_max
|
|
||||||
self.lr_max_decay_steps = max_decay_steps
|
|
||||||
self.last_lr = 0.0
|
|
||||||
self.verbosity_interval = verbosity_interval
|
|
||||||
|
|
||||||
def schedule(self, n, **kwargs):
|
|
||||||
if self.verbosity_interval > 0 and n % self.verbosity_interval == 0:
|
|
||||||
print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
|
|
||||||
if n < self.lr_warm_up_steps:
|
|
||||||
lr = (
|
|
||||||
self.lr_max - self.lr_start
|
|
||||||
) / self.lr_warm_up_steps * n + self.lr_start
|
|
||||||
self.last_lr = lr
|
|
||||||
return lr
|
|
||||||
t = (n - self.lr_warm_up_steps) / (
|
|
||||||
self.lr_max_decay_steps - self.lr_warm_up_steps
|
|
||||||
)
|
|
||||||
t = min(t, 1.0)
|
|
||||||
lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (1 + np.cos(t * np.pi))
|
|
||||||
self.last_lr = lr
|
|
||||||
return lr
|
|
||||||
|
|
||||||
def __call__(self, n, **kwargs):
|
|
||||||
return self.schedule(n, **kwargs)
|
|
||||||
|
|
||||||
|
|
||||||
class LambdaWarmUpCosineScheduler2:
|
|
||||||
"""
|
|
||||||
supports repeated iterations, configurable via lists
|
|
||||||
note: use with a base_lr of 1.0.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0
|
|
||||||
):
|
|
||||||
assert (
|
|
||||||
len(warm_up_steps)
|
|
||||||
== len(f_min)
|
|
||||||
== len(f_max)
|
|
||||||
== len(f_start)
|
|
||||||
== len(cycle_lengths)
|
|
||||||
)
|
|
||||||
self.lr_warm_up_steps = warm_up_steps
|
|
||||||
self.f_start = f_start
|
|
||||||
self.f_min = f_min
|
|
||||||
self.f_max = f_max
|
|
||||||
self.cycle_lengths = cycle_lengths
|
|
||||||
self.cum_cycles = np.cumsum([0, *list(self.cycle_lengths)])
|
|
||||||
self.last_f = 0.0
|
|
||||||
self.verbosity_interval = verbosity_interval
|
|
||||||
|
|
||||||
def find_in_interval(self, n):
|
|
||||||
interval = 0
|
|
||||||
for cl in self.cum_cycles[1:]:
|
|
||||||
if n <= cl:
|
|
||||||
return interval
|
|
||||||
interval += 1
|
|
||||||
return None
|
|
||||||
|
|
||||||
def schedule(self, n, **kwargs):
|
|
||||||
cycle = self.find_in_interval(n)
|
|
||||||
n = n - self.cum_cycles[cycle]
|
|
||||||
if self.verbosity_interval > 0 and n % self.verbosity_interval == 0:
|
|
||||||
print(
|
|
||||||
f"current step: {n}, recent lr-multiplier: {self.last_f}, "
|
|
||||||
f"current cycle {cycle}"
|
|
||||||
)
|
|
||||||
if n < self.lr_warm_up_steps[cycle]:
|
|
||||||
f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[
|
|
||||||
cycle
|
|
||||||
] * n + self.f_start[cycle]
|
|
||||||
self.last_f = f
|
|
||||||
return f
|
|
||||||
|
|
||||||
t = (n - self.lr_warm_up_steps[cycle]) / (
|
|
||||||
self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle]
|
|
||||||
)
|
|
||||||
t = min(t, 1.0)
|
|
||||||
f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
|
|
||||||
1 + np.cos(t * np.pi)
|
|
||||||
)
|
|
||||||
self.last_f = f
|
|
||||||
return f
|
|
||||||
|
|
||||||
def __call__(self, n, **kwargs):
|
|
||||||
return self.schedule(n, **kwargs)
|
|
||||||
|
|
||||||
|
|
||||||
class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
|
|
||||||
def schedule(self, n, **kwargs):
|
|
||||||
cycle = self.find_in_interval(n)
|
|
||||||
n = n - self.cum_cycles[cycle]
|
|
||||||
if self.verbosity_interval > 0 and n % self.verbosity_interval == 0:
|
|
||||||
print(
|
|
||||||
f"current step: {n}, recent lr-multiplier: {self.last_f}, "
|
|
||||||
f"current cycle {cycle}"
|
|
||||||
)
|
|
||||||
|
|
||||||
if n < self.lr_warm_up_steps[cycle]:
|
|
||||||
f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[
|
|
||||||
cycle
|
|
||||||
] * n + self.f_start[cycle]
|
|
||||||
self.last_f = f
|
|
||||||
return f
|
|
||||||
|
|
||||||
f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (
|
|
||||||
self.cycle_lengths[cycle] - n
|
|
||||||
) / (self.cycle_lengths[cycle])
|
|
||||||
self.last_f = f
|
|
||||||
return f
|
|
||||||
Loading…
Reference in New Issue
Block a user