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 } { 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 } )
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 } )
275+
276+ return schedule_func (optimizer , num_warmup_steps = num_warmup_steps , num_training_steps = num_training_steps )
0 commit comments