saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update LRScheduler to use scheduling similar to HuggingFace (#4880)

Browse source
This commit is contained in:
Thiago Crepaldi 2020-08-21 10:24:04 -07:00 committed by GitHub
parent ef19916d07
commit 5427a7e9af
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 181 additions and 155 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

142
orttraining/orttraining/python/experimental/optim/lr_scheduler.py
View file

@ -19,6 +19,15 @@ class _LRScheduler(object):
def __init__(self):
self._last_lr = []
def _step(self, train_step_info):
r"""Internal method called to compute learning rate"""
# Store last lr for future inquiry
new_lr = self.get_lr(train_step_info)
self._last_lr = new_lr
return new_lr
def get_lr(self, train_step_info):
r"""Returns a list of learning rate
@ -37,26 +46,18 @@ class _LRScheduler(object):
r""" Return last computed learning rate by LR Scheduler"""
return self._last_lr
def step(self, train_step_info):
r"""Public method called to update learning rate
NOTE: This class is used internally.
"""
# Store last lr for future inquiry
new_lr = self.get_lr(train_step_info)
self._last_lr = new_lr
# Update ORTTrainer's optimizer config instance
train_step_info.optimizer_config.lr = new_lr[0]
class ConstantWarmupLRScheduler(_LRScheduler):
r"""Constant warmup strategy for learning rate update
r"""Constant warmup strategy for learning rate update based on HuggingFace's Transformers implementation
Creates a schedule with constant learning rate preceded by a warmup period during which the learning rate
increases linearly between 0 and the initial lr set in the optimizer.
Learning rate update strategy:
lr = base_lr * (step / total_steps) / warmup, when step / total_steps < warmup
lr = base_lr, when step / total_steps >= warmup
When current_step < warmup
lr = base_lr * (current_step / max(1, num_warmup_steps))
Otherwise,
lr = base_lr
Args:
total_steps (int): total training steps for learning.
@ -90,28 +91,35 @@ class ConstantWarmupLRScheduler(_LRScheduler):
self.total_steps = total_steps
self.warmup = warmup
self._num_warmup_steps = warmup * total_steps
def _warmup_constant(self, train_step_info):
# Adds 1 to train_step_info.optimization_step and self.total_steps to prevent zero'ing lr
x = (train_step_info.optimization_step + 1) / (self.total_steps + 1)
if x < self.warmup:
return x/self.warmup
if train_step_info.optimization_step < self._num_warmup_steps:
return float(train_step_info.optimization_step) / float(max(1, self._num_warmup_steps))
return 1.0
def get_lr(self, train_step_info):
warmup = self._warmup_constant(train_step_info)
return [train_step_info.optimizer_config.lr * warmup]
return [train_step_info.optimizer_config.lr * self._warmup_constant(train_step_info)]
class CosineWarmupLRScheduler(_LRScheduler):
r"""Cosine warmup strategy for learning rate update
r"""Cosine warmup strategy for learning rate update based on HuggingFace's Transformers implementation
Creates a schedule with learning rate that decreases following the values of the cosine function between the
initial lr set in the :py:class`.optim._OptimizerConfig` to 0, after a warmup period during which it increases
linearly between 0 and the initial lr set in the :py:class`.optim._OptimizerConfig`.
Learning rate update strategy:
lr = base_lr * (step / total_steps) / warmup, when step / total_steps < warmup
lr = base_lr * 0.5 * (1.0 + cosine(pi * (step / total_steps))), when step / total_steps >= warmup
When current_step < warmup
lr = base_lr * (current_step / max(1, num_warmup_steps)), when
Otherwise
lr = base_lr * max(0, 0.5 * (1.0 + math.cos(math.pi * cycles * 2.0 * progress))), where
progress = current_step - num_warmup_steps / max(1, total_steps - num_warmup_steps)
Args:
total_steps (int): total training steps for learning.
cycles (float, default is 0.5): number of waves in the cosine schedule.
The default decreases from max value to 0, following a half-cosine
warmup (float, default is 0.002): portion of total steps for warmup. Range is (0, 1]
Example:
@ -131,35 +139,44 @@ class CosineWarmupLRScheduler(_LRScheduler):
outputs = ort_trainer.train_step(**inputs)
"""
def __init__(self, total_steps, warmup=0.002):
def __init__(self, total_steps, cycles=0.5, warmup=0.002):
super().__init__()
assert isinstance(total_steps, int) and total_steps > 0,\
"total_steps must be a strict positive number"
assert isinstance(cycles, float) and cycles > 0,\
"cycles must be a positive float"
assert isinstance(warmup, float) and warmup >= 0 and warmup < 1,\
"warmup must be a float between (0, 1]"
assert total_steps > warmup,\
"total_steps must be greater than warmup"
self.total_steps = total_steps
self.cycles = cycles
self.warmup = warmup
self._num_warmup_steps = warmup * total_steps
def _warmup_cosine(self, train_step_info):
# Adds 1 to train_step_info.optimization_step and self.total_steps to prevent zero'ing lr
x = (train_step_info.optimization_step + 1) / (self.total_steps + 1)
if x < self.warmup:
return x/self.warmup
return 0.5 * (1.0 + math.cos(math.pi * x))
if train_step_info.optimization_step < self._num_warmup_steps:
return float(train_step_info.optimization_step) / float(max(1, self._num_warmup_steps))
progress = float(train_step_info.optimization_step - self._num_warmup_steps) / float(max(1, self.total_steps - self._num_warmup_steps))
return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(self.cycles) * 2.0 * progress)))
def get_lr(self, train_step_info):
return [train_step_info.optimizer_config.lr * self._warmup_cosine(train_step_info)]
class LinearWarmupLRScheduler(_LRScheduler):
r"""Linear warmup strategy for learning rate update
r"""Linear warmup strategy for learning rate update based on HuggingFace's Transformers implementation
Creates a schedule with a learning rate that decreases linearly from the initial lr
set in the :py:class`.optim._OptimizerConfig` to 0, after a warmup period during which
it increases linearly from 0 to the initial lr set in the :py:class`.optim._OptimizerConfig`.
Learning rate update strategy:
lr = base_lr * (step / total_steps) / warmup, when step / total_steps < warmup
lr = base_lr * max(((step / total_steps) - 1.) / (warmup - 1.), 0.), when step / total_steps >= warmup
When current_step < warmup
lr = base_lr * (current_step / max(1, num_warmup_steps))
Otherwise
lr = base_lr * (max(0, total_steps - current_step) / max(1, total_steps - num_warmup_steps)))
Args:
total_steps (int): total training steps for learning.
@ -193,29 +210,36 @@ class LinearWarmupLRScheduler(_LRScheduler):
self.total_steps = total_steps
self.warmup = warmup
self._num_warmup_steps = warmup * total_steps
def _warmup_linear(self, train_step_info):
# Adds 1 to train_step_info.optimization_step and self.total_steps to prevent zero'ing lr
x = (train_step_info.optimization_step + 1) / (self.total_steps + 1)
if x < self.warmup:
return x / self.warmup
return max((x - 1.) / (self.warmup - 1.), 0.)
if train_step_info.optimization_step < self._num_warmup_steps:
return float(train_step_info.optimization_step) / float(max(1, self._num_warmup_steps))
return max(0.0, float(self.total_steps - train_step_info.optimization_step) / float(max(1, self.total_steps - self._num_warmup_steps)))
def get_lr(self, train_step_info):
return [train_step_info.optimizer_config.lr * self._warmup_linear(train_step_info)]
class PolyWarmupLRScheduler(_LRScheduler):
r"""Polynomial warmup strategy for learning rate update
r"""Polynomial warmup strategy for learning rate update based on HuggingFace's Transformers implementation
Creates a schedule with a learning rate that decreases as a polynomial decay
from the initial lr set in the :py:class`.optim._OptimizerConfig` to lr_end,
after a warmup period during which it increases linearly from 0 to the
initial lr set in the :py:class`.optim._OptimizerConfig`
Learning rate update strategy:
lr = base_lr * (step / total_steps) / warmup, when step / total_steps < warmup
lr = base_lr * (1 − step / total_steps ) ^ degree, when step / total_steps >= warmup
When current_step < warmup
lr = base_lr * (current_step / max(1, num_warmup_steps))
Otherwise
Args:
total_steps (int): total training steps for learning.
lr_end (float, default 1e-7): final learning rate value.
Applies to the default lr and parameter groups in :py:class:`.optim._OptimizerConfig`
power (float, default is 1.0): polynomial factor
warmup (float, default is 0.002): portion of total steps for warmup. Range is (0, 1]
degree (float, default is 0.5): polynomial power
Example:
.. code-block:: python
@ -234,27 +258,41 @@ class PolyWarmupLRScheduler(_LRScheduler):
outputs = ort_trainer.train_step(**inputs)
"""
def __init__(self, total_steps, warmup=0.002, degree=0.5):
def __init__(self, total_steps, lr_end=1e-7, power=1.0, warmup=0.002):
super().__init__()
assert isinstance(total_steps, int) and total_steps > 0,\
"total_steps must be a strict positive number"
assert isinstance(lr_end, float) and lr_end >= 0,\
"lr_end must be a positive float"
assert isinstance(warmup, float) and warmup >= 0 and warmup < 1,\
"warmup must be a float between (0, 1]"
assert isinstance(power, float) and power >= 0,\
"power must be a positive float"
assert total_steps > warmup,\
"total_steps must be greater than warmup"
assert isinstance(degree, float) and warmup >= 0,\
"degree must be a positive float"
self.total_steps = total_steps
self.lr_end = lr_end
self.power = power
self.warmup = warmup
self.degree = degree
self._num_warmup_steps = warmup * total_steps
def _warmup_poly(self, train_step_info):
# Adds 1 to train_step_info.optimization_step and self.total_steps to prevent zero'ing lr
x = (train_step_info.optimization_step + 1) / (self.total_steps + 1)
if x < self.warmup:
return x/self.warmup
return (1.0 - x)**self.degree
assert train_step_info.optimizer_config.lr > self.lr_end,\
f"lr_end ({lr_end}) must be be smaller than initial lr ({train_step_info.optimizer_config.lr})"
if train_step_info.optimization_step < self._num_warmup_steps:
return float(train_step_info.optimization_step) / float(max(1, self._num_warmup_steps))
elif train_step_info.optimization_step > self.total_steps:
return self.lr_end / train_step_info.optimizer_config.lr
else:
lr_range = train_step_info.optimizer_config.lr - self.lr_end
decay_steps = self.total_steps - self._num_warmup_steps
pct_remaining = 1 - (train_step_info.optimization_step - self._num_warmup_steps) / decay_steps
decay = lr_range * pct_remaining ** self.power + self.lr_end
return decay / train_step_info.optimizer_config.lr
def get_lr(self, train_step_info):
return [train_step_info.optimizer_config.lr * self._warmup_poly(train_step_info)]

9
orttraining/orttraining/python/experimental/orttrainer.py
View file

@ -328,8 +328,9 @@ class ORTTrainer(object):
outputs_desc = self._model_desc_outputs_with_all_finite
# Update Learning Rate if Necessary
lr = self.optim_config.lr
if self.options.lr_scheduler:
self.options.lr_scheduler.step(self._train_step_info)
lr = self.options.lr_scheduler._step(self._train_step_info)[0]
# Loss Scale for mixed precision
loss_scale = None
@ -340,7 +341,7 @@ class ORTTrainer(object):
inputs_desc = self._model_desc_inputs_with_lr_and_loss_scale
# Get data. CombineTorchModelLossFn takes label as last input and outputs loss first
input = self._prepare_model_input(inputs_desc, self.optim_config.lr, loss_scale, *args, **kwargs)
input = self._prepare_model_input(inputs_desc, lr, loss_scale, *args, **kwargs)
# Normalize input
if not isinstance(args, (list, tuple)):
@ -658,13 +659,13 @@ class ORTTrainer(object):
# Append learning rate
extra_inputs = 0
if lr:
if lr is not None:
lr = torch.tensor([lr])
input += (lr,)
extra_inputs += 1
# Append loss scale
if loss_scale:
if loss_scale is not None:
assert self.options.mixed_precision.enabled, "Loss scale cannot be used without mixed precision"
loss_scale = loss_scale.clone().detach()
input += (loss_scale, )

128
orttraining/orttraining/test/python/orttraining_test_orttrainer_bert_toy_onnx.py
View file

@ -79,7 +79,7 @@ def optimizer_parameters(model):
if any(key in initializer.name for key in no_decay_keys):
no_decay_param_group.append(initializer.name)
params = [{'params': no_decay_param_group, "alpha": 0.9, "beta": 0.999, "lambda_coef": 0.0, "epsilon": 1e-6}]
return params
@ -158,61 +158,63 @@ def legacy_constant_lr_scheduler_5(global_step):
def legacy_constant_lr_scheduler(global_step, initial_lr):
warmup = 0.5
total_steps = 10
lr = initial_lr
for i in range(global_step+1):
x = (i+1) / (total_steps+1)
if x < warmup:
warmup_val = x/warmup
else:
warmup_val =1
lr *= warmup_val
return lr
num_warmup_steps = warmup * total_steps
if global_step < num_warmup_steps:
new_lr = initial_lr * float(global_step) / float(max(1, num_warmup_steps))
else:
new_lr = initial_lr
return new_lr
def legacy_cosine_lr_scheduler(global_step):
initial_lr = 1.0
warmup = 0.5
total_steps = 10
lr = initial_lr
for i in range(global_step+1):
x = (i+1) / (total_steps+1)
if x < warmup:
warmup_val = x/warmup
else:
warmup_val = 0.5 * (1.0 + math.cos(math.pi * x))
lr *= warmup_val
return lr
cycles = 0.5
num_warmup_steps = warmup * total_steps
if global_step < num_warmup_steps:
new_lr = initial_lr * float(global_step) / float(max(1, num_warmup_steps))
else:
progress = float(global_step - num_warmup_steps) / float(max(1, total_steps - num_warmup_steps))
new_lr = initial_lr * max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(cycles) * 2.0 * progress)))
return new_lr
def legacy_linear_lr_scheduler(global_step):
initial_lr = 1.0
warmup = 0.5
total_steps = 10
lr = initial_lr
for i in range(global_step+1):
x = (i+1) / (total_steps+1)
if x < warmup:
warmup_val = x/warmup
else:
warmup_val = max((x - 1.0) / (warmup - 1.0), 0.0)
lr *= warmup_val
return lr
num_warmup_steps = warmup * total_steps
if global_step < num_warmup_steps:
new_lr = initial_lr * float(global_step) / float(max(1, num_warmup_steps))
else:
new_lr = max(0.0, float(total_steps - global_step) / float(max(1, total_steps - num_warmup_steps)))
return new_lr
def legacy_poly_lr_scheduler(global_step):
initial_lr = 1.0
warmup = 0.5
total_steps = 10
degree = 0.5
lr = initial_lr
for i in range(global_step+1):
x = (i+1) / (total_steps+1)
if x < warmup:
warmup_val = x/warmup
else:
warmup_val = (1.0 - x) ** degree
lr *= warmup_val
return lr
lr_end = 1e-7
power = 1.0
num_warmup_steps = warmup * total_steps
if global_step < num_warmup_steps:
new_lr = initial_lr * float(global_step) / float(max(1, num_warmup_steps))
elif global_step > total_steps:
new_lr = lr_end / initial_lr
else:
lr_range = initial_lr - lr_end
decay_steps = total_steps - num_warmup_steps
pct_remaining = 1 - (global_step - num_warmup_steps) / decay_steps
decay = lr_range * pct_remaining ** power + lr_end
new_lr = decay / initial_lr
return new_lr
def legacy_optim_params_a(name):
@ -294,26 +296,26 @@ def testToyBERTDeterministicCheck(expected_losses):
@pytest.mark.parametrize("initial_lr, lr_scheduler, expected_learning_rates, expected_losses", [
(1.0, optim.lr_scheduler.ConstantWarmupLRScheduler, [0.18181818181818182, 0.06611570247933884, 0.03606311044327573, 0.026227716686018716, 0.02384337880547156,\
0.02384337880547156, 0.02384337880547156, 0.02384337880547156, 0.02384337880547156, 0.02384337880547156],
[10.988012313842773, 11.637386322021484, 11.099013328552246, 11.055734634399414, 11.145816802978516,\
10.974218368530273, 10.971613883972168, 11.203381538391113, 11.131250381469727, 11.017223358154297]),
(0.5, optim.lr_scheduler.ConstantWarmupLRScheduler, [0.09090909090909091, 0.03305785123966942, 0.018031555221637866, 0.013113858343009358, 0.01192168940273578,\
0.01192168940273578, 0.01192168940273578, 0.01192168940273578, 0.01192168940273578, 0.01192168940273578],
[10.988012313842773, 11.310077667236328, 11.025278091430664, 10.988797187805176, 11.125761032104492,\
10.958372116088867, 10.980047225952148, 11.175304412841797, 11.147686958312988, 11.10694694519043]),
(1.0, optim.lr_scheduler.CosineWarmupLRScheduler, [0.18181818181818182, 0.06611570247933884, 0.03606311044327573, 0.026227716686018716, 0.02384337880547156,\
0.010225056103441101, 0.0029887071446425494, 0.0005157600951772063, 4.093754650801759e-05, 8.291291382790071e-07],
[10.988012313842773, 11.637386322021484, 11.099013328552246, 11.05573558807373, 11.145816802978516,\
10.974218368530273, 10.964020729064941, 11.190014839172363, 11.16644287109375, 11.150431632995605]),
(1.0, optim.lr_scheduler.LinearWarmupLRScheduler, [0.18181818181818182, 0.06611570247933884, 0.03606311044327573, 0.026227716686018716, 0.02384337880547156,\
0.021675798914065056, 0.015764217392047315, 0.008598664032025808, 0.0031267869207366565, 0.0005685067128612105],
[10.988012313842773, 11.637386322021484, 11.099013328552246, 11.05573558807373, 11.145816802978516,\
10.974218368530273, 10.970070838928223, 11.198983192443848, 11.134098052978516, 11.067017555236816]),
(1.0, optim.lr_scheduler.PolyWarmupLRScheduler, [0.18181818181818182, 0.06611570247933884, 0.03606311044327573, 0.026227716686018716, 0.02384337880547156,\
0.01607520271130791, 0.009693711967693117, 0.005062375970537139, 0.0021586043667598935, 0.0006508437050332076],
[10.988012313842773, 11.637386322021484, 11.099013328552246, 11.055734634399414, 11.145816802978516,\
10.974217414855957, 10.96664810180664, 11.193868637084961, 11.14560604095459, 11.097070693969727])
(1.0, optim.lr_scheduler.ConstantWarmupLRScheduler,\
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.0, 1.0, 1.0, 1.0],
[10.988012313842773, 10.99213981628418, 18.11327362060547, 35.164119720458984, 40.68449401855469,\
24.585613250732422, 90.66416931152344, 97.23558044433594, 123.73894500732422, 185.28268432617188]),
(0.5, optim.lr_scheduler.ConstantWarmupLRScheduler,\
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.5, 0.5, 0.5, 0.5],
[10.988012313842773, 10.99213981628418, 14.942025184631348, 32.938804626464844, 15.188745498657227,\
31.730798721313477, 38.70293045043945, 29.556488037109375, 34.914039611816406, 43.36301040649414]),
(1.0, optim.lr_scheduler.CosineWarmupLRScheduler,\
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.9045084971874737, 0.6545084971874737, 0.34549150281252633, 0.09549150281252633],
[10.988012313842773, 10.99213981628418, 18.11327362060547, 35.164119720458984, 40.68449401855469,\
24.585613250732422, 90.66416931152344, 95.54305267333984, 77.43124389648438, 111.71074676513672]),
(1.0, optim.lr_scheduler.LinearWarmupLRScheduler,\
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.8, 0.6, 0.4, 0.2],\
[10.988012313842773, 10.99213981628418, 18.11327362060547, 35.164119720458984, 40.68449401855469, 24.585613250732422,\
90.66416931152344, 94.88548278808594, 77.90852355957031, 116.56438446044922]),
(1.0, optim.lr_scheduler.PolyWarmupLRScheduler,\
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.80000002, 0.60000004, 0.40000006000000005, 0.20000007999999997],
[10.988012313842773, 10.99213981628418, 18.11327362060547, 35.164119720458984, 40.68449401855469, 24.585613250732422,\
90.66416931152344, 94.88548278808594, 77.90852355957031, 116.56439971923828])
])
def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rates, expected_losses):
# Common setup
@ -326,7 +328,7 @@ def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rate
# Modeling
model_desc = bert_model_description()
model = load_bert_onnx_model()
optim_config = optim.LambConfig(lr=initial_lr)
optim_config = optim.AdamConfig(lr=initial_lr)
opts = orttrainer.ORTTrainerOptions({
'debug' : {
'deterministic_compute': True
@ -620,9 +622,9 @@ def testORTTrainerFrozenWeights(model_params):
},
}
opts = orttrainer.ORTTrainerOptions(opts_dict)
torch.manual_seed(seed)
set_seed(seed)
onnxruntime.set_seed(seed)
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
for i in range(total_steps):
@ -762,7 +764,7 @@ def testToyBERTModelLegacyExperimentalLRScheduler(initial_lr, lr_scheduler, lega
model = load_bert_onnx_model()
torch.manual_seed(seed)
onnxruntime.set_seed(seed)
optim_config = optim.LambConfig(lr=initial_lr)
optim_config = optim.AdamConfig(lr=initial_lr)
opts = orttrainer.ORTTrainerOptions({
'debug' : {
'deterministic_compute': True
@ -784,7 +786,7 @@ def testToyBERTModelLegacyExperimentalLRScheduler(initial_lr, lr_scheduler, lega
onnxruntime.set_seed(seed)
device = torch.device(device)
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(initial_lr)
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "AdamOptimizer",
None,
learning_rate_description,
device,

37
orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py
View file

@ -502,14 +502,16 @@ def testLinearLRSchedulerCreation():
@pytest.mark.parametrize("lr_scheduler,expected_values", [
(optim.lr_scheduler.ConstantWarmupLRScheduler, [0.181818, 0.066116, 0.036063, 0.026228, 0.023843,
0.023843, 0.023843, 0.023843, 0.023843, 0.023843]),
(optim.lr_scheduler.CosineWarmupLRScheduler, [0.181818, 0.066116, 0.036063, 0.026228, 0.023843,
0.010225, 0.002989, 0.0005158, 0.000040937, 0.0000008291]),
(optim.lr_scheduler.LinearWarmupLRScheduler, [0.181818, 0.066116, 0.036063, 0.026228, 0.023843,
0.021675, 0.0157636, 0.0085983, 0.0031266, 0.00056847]),
(optim.lr_scheduler.PolyWarmupLRScheduler, [0.181818, 0.066116, 0.036063, 0.026228, 0.023843,
0.0160749, 0.0096935, 0.0050622, 0.0021585, 0.000650833])
(optim.lr_scheduler.ConstantWarmupLRScheduler,
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.0, 1.0, 1.0, 1.0]),
(optim.lr_scheduler.CosineWarmupLRScheduler,
[0.0, 0.9763960957919413, 0.9059835861602854, 0.7956724530494887, 0.6563036824392345,\
0.5015739416158049, 0.34668951940611276, 0.2068719061737831, 0.09586187986225325, 0.0245691111902418]),
(optim.lr_scheduler.LinearWarmupLRScheduler,
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.8, 0.6, 0.4, 0.2]),
(optim.lr_scheduler.PolyWarmupLRScheduler,
[0.0, 0.9509018036072144, 0.9008016032064128, 0.8507014028056112, 0.8006012024048097,\
0.750501002004008, 0.7004008016032064, 0.6503006012024048, 0.6002004008016032, 0.5501002004008015])
])
def testLRSchedulerUpdateImpl(lr_scheduler, expected_values):
# Test tolerance
@ -527,7 +529,7 @@ def testLRSchedulerUpdateImpl(lr_scheduler, expected_values):
# Emulate ORTTRainer.train_step() call that updates its train_step_info
train_step_info = TrainStepInfo(optimizer_config=optimizer_config, optimization_step=optimization_step)
lr_scheduler.step(train_step_info)
lr_scheduler._step(train_step_info)
lr_list = lr_scheduler.get_last_lr()
assert len(lr_list) == 1
assert_allclose(lr_list[0],
@ -537,14 +539,15 @@ def testLRSchedulerUpdateImpl(lr_scheduler, expected_values):
@pytest.mark.parametrize("step_fn, lr_scheduler, expected_lr_values, device", [
('train_step', None, None, 'cuda'),
('eval_step', None, None, 'cpu'),
('train_step', optim.lr_scheduler.ConstantWarmupLRScheduler, [0.181818, 0.066116, 0.036063, 0.026228, 0.023843,
0.023843, 0.023843, 0.023843, 0.023843, 0.023843], 'cpu'),
('train_step', optim.lr_scheduler.CosineWarmupLRScheduler, [0.181818, 0.066116, 0.036063, 0.026228, 0.023843,
0.010225, 0.002989, 0.0005158, 0.000040937, 0.0000008291], 'cuda'),
('train_step', optim.lr_scheduler.LinearWarmupLRScheduler, [0.181818, 0.066116, 0.036063, 0.026228, 0.023843,
0.021675, 0.0157636, 0.0085983, 0.0031266, 0.00056847], 'cpu'),
('train_step', optim.lr_scheduler.PolyWarmupLRScheduler, [0.181818, 0.066116, 0.036063, 0.026228, 0.023843,
0.0160749, 0.0096935, 0.0050622, 0.0021585, 0.000650833], 'cuda')
('train_step', optim.lr_scheduler.ConstantWarmupLRScheduler,
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.0, 1.0, 1.0, 1.0], 'cpu'),
('train_step', optim.lr_scheduler.CosineWarmupLRScheduler,
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.9045084971874737, 0.6545084971874737, 0.34549150281252633, 0.09549150281252633],
'cuda'),
('train_step', optim.lr_scheduler.LinearWarmupLRScheduler,
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.8, 0.6, 0.4, 0.2], 'cpu'),
('train_step', optim.lr_scheduler.PolyWarmupLRScheduler,
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.80000002, 0.60000004, 0.40000006000000005, 0.20000007999999997], 'cuda')
])
def testInstantiateORTTrainer(step_fn, lr_scheduler, expected_lr_values, device):
total_steps = 1

20
orttraining/orttraining/test/python/orttraining_transformer_trainer.py
View file

@ -87,24 +87,6 @@ def get_linear_schedule_with_warmup(num_warmup_steps, num_training_steps, base_l
return lambda_lr_get_lr
class CustomLRSchedulerLinearWarmup(_LRScheduler):
def __init__(self, num_warmup_steps, num_training_steps, base_lr):
super().__init__()
self.num_warmup_steps = num_warmup_steps
self.num_training_steps = num_training_steps
self.base_lr = base_lr
def lr_lambda_linear(self, current_step):
if current_step < self.num_warmup_steps:
return float(current_step) / float(max(1, self.num_warmup_steps))
return max(0.0, float(self.num_training_steps - current_step) / float(max(1, self.num_training_steps - self.num_warmup_steps)))
def get_lr(self, train_step_info):
# LambdaLR increment self.last_epoch at every step()
return [self.base_lr * self.lr_lambda_linear(train_step_info.optimization_step)]
class ORTTransformerTrainer:
"""
"""
@ -198,7 +180,7 @@ class ORTTransformerTrainer:
'outputs': [('loss', [], True),
('logits', ['batch', 2])]}
lr_scheduler = CustomLRSchedulerLinearWarmup(self.args.warmup_steps, t_total, self.args.learning_rate)
lr_scheduler = orttrainer.optim.LinearWarmupLRScheduler(t_total, self.args.warmup_steps/float(t_total))
loss_scaler = amp.DynamicLossScaler() if self.args.fp16 else None
device = self.args.device.type