Skip to content

Commit 71289ba

Browse files
anton-lanton-l
committed
add lr schedule utils
1 parent 0417baf commit 71289ba

3 files changed

Lines changed: 297 additions & 2 deletions

File tree

‎examples/train_unconditional.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
ToTensor,
2222
)
2323
from tqdm.auto import tqdm
24-
from transformers import get_linear_schedule_with_warmup
24+
from diffusers.optimization import get_scheduler
2525

2626

2727
logger = logging.get_logger(__name__)
@@ -60,7 +60,8 @@ def transforms(examples):
6060
dataset.set_transform(transforms)
6161
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
6262

63-
lr_scheduler = get_linear_schedule_with_warmup(
63+
lr_scheduler = get_scheduler(
64+
"linear",
6465
optimizer=optimizer,
6566
num_warmup_steps=args.warmup_steps,
6667
num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,
@@ -107,11 +108,13 @@ def transforms(examples):
107108
output = model(noisy_images, timesteps)
108109
# predict the noise residual
109110
loss = F.mse_loss(output, noise_samples)
111+
loss = loss / args.gradient_accumulation_steps
110112
accelerator.backward(loss)
111113
else:
112114
output = model(noisy_images, timesteps)
113115
# predict the noise residual
114116
loss = F.mse_loss(output, noise_samples)
117+
loss = loss / args.gradient_accumulation_steps
115118
accelerator.backward(loss)
116119
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
117120
optimizer.step()

‎src/diffusers/hub_utils.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,19 @@
1+
# coding=utf-8
2+
# Copyright 2022 The HuggingFace Inc. team.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
117
import os
218
import shutil
319
from pathlib import Path

‎src/diffusers/optimization.py‎

Lines changed: 276 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,276 @@
1+
# coding=utf-8
2+
# Copyright 2022 The HuggingFace Inc. team.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
"""PyTorch optimization for diffusion models."""
16+
17+
import math
18+
from enum import Enum
19+
from typing import Optional, Union
20+
21+
import torch
22+
from torch.optim import Optimizer
23+
from torch.optim.lr_scheduler import LambdaLR
24+
25+
from .utils import logging
26+
27+
28+
logger = logging.get_logger(__name__)
29+
30+
31+
class SchedulerType(Enum):
32+
LINEAR = "linear"
33+
COSINE = "cosine"
34+
COSINE_WITH_RESTARTS = "cosine_with_restarts"
35+
POLYNOMIAL = "polynomial"
36+
CONSTANT = "constant"
37+
CONSTANT_WITH_WARMUP = "constant_with_warmup"
38+
39+
40+
def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
41+
"""
42+
Create a schedule with a constant learning rate, using the learning rate set in optimizer.
43+
44+
Args:
45+
optimizer ([`~torch.optim.Optimizer`]):
46+
The optimizer for which to schedule the learning rate.
47+
last_epoch (`int`, *optional*, defaults to -1):
48+
The index of the last epoch when resuming training.
49+
50+
Return:
51+
`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
52+
"""
53+
return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)
54+
55+
56+
def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):
57+
"""
58+
Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
59+
increases linearly between 0 and the initial lr set in the optimizer.
60+
61+
Args:
62+
optimizer ([`~torch.optim.Optimizer`]):
63+
The optimizer for which to schedule the learning rate.
64+
num_warmup_steps (`int`):
65+
The number of steps for the warmup phase.
66+
last_epoch (`int`, *optional*, defaults to -1):
67+
The index of the last epoch when resuming training.
68+
69+
Return:
70+
`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
71+
"""
72+
73+
def lr_lambda(current_step: int):
74+
if current_step < num_warmup_steps:
75+
return float(current_step) / float(max(1.0, num_warmup_steps))
76+
return 1.0
77+
78+
return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)
79+
80+
81+
def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
82+
"""
83+
Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after
84+
a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
85+
86+
Args:
87+
optimizer ([`~torch.optim.Optimizer`]):
88+
The optimizer for which to schedule the learning rate.
89+
num_warmup_steps (`int`):
90+
The number of steps for the warmup phase.
91+
num_training_steps (`int`):
92+
The total number of training steps.
93+
last_epoch (`int`, *optional*, defaults to -1):
94+
The index of the last epoch when resuming training.
95+
96+
Return:
97+
`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
98+
"""
99+
100+
def lr_lambda(current_step: int):
101+
if current_step < num_warmup_steps:
102+
return float(current_step) / float(max(1, num_warmup_steps))
103+
return max(
104+
0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
105+
)
106+
107+
return LambdaLR(optimizer, lr_lambda, last_epoch)
108+
109+
110+
def get_cosine_schedule_with_warmup(
111+
optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
112+
):
113+
"""
114+
Create a schedule with a learning rate that decreases following the values of the cosine function between the
115+
initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
116+
initial lr set in the optimizer.
117+
118+
Args:
119+
optimizer ([`~torch.optim.Optimizer`]):
120+
The optimizer for which to schedule the learning rate.
121+
num_warmup_steps (`int`):
122+
The number of steps for the warmup phase.
123+
num_training_steps (`int`):
124+
The total number of training steps.
125+
num_cycles (`float`, *optional*, defaults to 0.5):
126+
The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
127+
following a half-cosine).
128+
last_epoch (`int`, *optional*, defaults to -1):
129+
The index of the last epoch when resuming training.
130+
131+
Return:
132+
`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
133+
"""
134+
135+
def lr_lambda(current_step):
136+
if current_step < num_warmup_steps:
137+
return float(current_step) / float(max(1, num_warmup_steps))
138+
progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
139+
return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
140+
141+
return LambdaLR(optimizer, lr_lambda, last_epoch)
142+
143+
144+
def get_cosine_with_hard_restarts_schedule_with_warmup(
145+
optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
146+
):
147+
"""
148+
Create a schedule with a learning rate that decreases following the values of the cosine function between the
149+
initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
150+
linearly between 0 and the initial lr set in the optimizer.
151+
152+
Args:
153+
optimizer ([`~torch.optim.Optimizer`]):
154+
The optimizer for which to schedule the learning rate.
155+
num_warmup_steps (`int`):
156+
The number of steps for the warmup phase.
157+
num_training_steps (`int`):
158+
The total number of training steps.
159+
num_cycles (`int`, *optional*, defaults to 1):
160+
The number of hard restarts to use.
161+
last_epoch (`int`, *optional*, defaults to -1):
162+
The index of the last epoch when resuming training.
163+
164+
Return:
165+
`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
166+
"""
167+
168+
def lr_lambda(current_step):
169+
if current_step < num_warmup_steps:
170+
return float(current_step) / float(max(1, num_warmup_steps))
171+
progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
172+
if progress >= 1.0:
173+
return 0.0
174+
return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))
175+
176+
return LambdaLR(optimizer, lr_lambda, last_epoch)
177+
178+
179+
def get_polynomial_decay_schedule_with_warmup(
180+
optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1
181+
):
182+
"""
183+
Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the
184+
optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the
185+
initial lr set in the optimizer.
186+
187+
Args:
188+
optimizer ([`~torch.optim.Optimizer`]):
189+
The optimizer for which to schedule the learning rate.
190+
num_warmup_steps (`int`):
191+
The number of steps for the warmup phase.
192+
num_training_steps (`int`):
193+
The total number of training steps.
194+
lr_end (`float`, *optional*, defaults to 1e-7):
195+
The end LR.
196+
power (`float`, *optional*, defaults to 1.0):
197+
Power factor.
198+
last_epoch (`int`, *optional*, defaults to -1):
199+
The index of the last epoch when resuming training.
200+
201+
Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT
202+
implementation at
203+
https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
204+
205+
Return:
206+
`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
207+
208+
"""
209+
210+
lr_init = optimizer.defaults["lr"]
211+
if not (lr_init > lr_end):
212+
raise ValueError(f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})")
213+
214+
def lr_lambda(current_step: int):
215+
if current_step < num_warmup_steps:
216+
return float(current_step) / float(max(1, num_warmup_steps))
217+
elif current_step > num_training_steps:
218+
return lr_end / lr_init # as LambdaLR multiplies by lr_init
219+
else:
220+
lr_range = lr_init - lr_end
221+
decay_steps = num_training_steps - num_warmup_steps
222+
pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps
223+
decay = lr_range * pct_remaining**power + lr_end
224+
return decay / lr_init # as LambdaLR multiplies by lr_init
225+
226+
return LambdaLR(optimizer, lr_lambda, last_epoch)
227+
228+
229+
TYPE_TO_SCHEDULER_FUNCTION = {
230+
SchedulerType.LINEAR: get_linear_schedule_with_warmup,
231+
SchedulerType.COSINE: get_cosine_schedule_with_warmup,
232+
SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
233+
SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
234+
SchedulerType.CONSTANT: get_constant_schedule,
235+
SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
236+
}
237+
238+
239+
def get_scheduler(
240+
name: Union[str, SchedulerType],
241+
optimizer: Optimizer,
242+
num_warmup_steps: Optional[int] = None,
243+
num_training_steps: Optional[int] = None,
244+
):
245+
"""
246+
Unified API to get any scheduler from its name.
247+
248+
Args:
249+
name (`str` or `SchedulerType`):
250+
The name of the scheduler to use.
251+
optimizer (`torch.optim.Optimizer`):
252+
The optimizer that will be used during training.
253+
num_warmup_steps (`int`, *optional*):
254+
The number of warmup steps to do. This is not required by all schedulers (hence the argument being
255+
optional), the function will raise an error if it's unset and the scheduler type requires it.
256+
num_training_steps (`int``, *optional*):
257+
The number of training steps to do. This is not required by all schedulers (hence the argument being
258+
optional), the function will raise an error if it's unset and the scheduler type requires it.
259+
"""
260+
name = SchedulerType(name)
261+
schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
262+
if name == SchedulerType.CONSTANT:
263+
return schedule_func(optimizer)
264+
265+
# All other schedulers require `num_warmup_steps`
266+
if num_warmup_steps is None:
267+
raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
268+
269+
if name == SchedulerType.CONSTANT_WITH_WARMUP:
270+
return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
271+
272+
# All other schedulers require `num_training_steps`
273+
if num_training_steps is None:
274+
raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
275+
276+
return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)

0 commit comments

Comments
 (0)