wehub-resource-sync
72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2019-11-11 18:44
|
|
import tensorflow as tf
|
|
from hanlp.optimizers.adamw.optimization import WarmUp, AdamWeightDecay
|
|
|
|
|
|
# from hanlp.optimization.adamw.optimizers_v2 import AdamW
|
|
# from hanlp.optimization.adamw.utils import get_weight_decays
|
|
|
|
|
|
# def create_optimizer(model, init_lr, num_train_steps, num_warmup_steps):
|
|
# """Creates an optimizer with learning rate schedule."""
|
|
# wd_dict = get_weight_decays(model)
|
|
#
|
|
# # Implements linear decay of the learning rate.
|
|
# learning_rate_fn = tf.keras.optimizers.schedules.PolynomialDecay(
|
|
# initial_learning_rate=init_lr,
|
|
# decay_steps=num_train_steps,
|
|
# end_learning_rate=0.0)
|
|
# if num_warmup_steps:
|
|
# learning_rate_fn = WarmUp(initial_learning_rate=init_lr,
|
|
# decay_schedule_fn=learning_rate_fn,
|
|
# warmup_steps=num_warmup_steps)
|
|
# optimizer = AdamW(
|
|
# learning_rate=learning_rate_fn,
|
|
# weight_decay_rate=0.01,
|
|
# beta_1=0.9,
|
|
# beta_2=0.999,
|
|
# epsilon=1e-6,
|
|
# exclude_from_weight_decay=['layer_norm', 'bias'])
|
|
# return optimizer
|
|
|
|
|
|
def create_optimizer(init_lr, num_train_steps, num_warmup_steps, weight_decay_rate=0.01, epsilon=1e-6, clipnorm=None):
|
|
"""Creates an optimizer with learning rate schedule.
|
|
|
|
Args:
|
|
init_lr:
|
|
num_train_steps:
|
|
num_warmup_steps:
|
|
weight_decay_rate: (Default value = 0.01)
|
|
epsilon: (Default value = 1e-6)
|
|
clipnorm: (Default value = None)
|
|
|
|
Returns:
|
|
|
|
"""
|
|
# Implements linear decay of the learning rate.
|
|
learning_rate_fn = tf.keras.optimizers.schedules.PolynomialDecay(
|
|
initial_learning_rate=init_lr,
|
|
decay_steps=num_train_steps,
|
|
end_learning_rate=0.0)
|
|
if num_warmup_steps:
|
|
learning_rate_fn = WarmUp(initial_learning_rate=init_lr,
|
|
decay_schedule_fn=learning_rate_fn,
|
|
warmup_steps=num_warmup_steps)
|
|
additional_args = {}
|
|
if clipnorm:
|
|
additional_args['clipnorm'] = clipnorm
|
|
optimizer = AdamWeightDecay(
|
|
learning_rate=learning_rate_fn,
|
|
weight_decay_rate=weight_decay_rate,
|
|
beta_1=0.9,
|
|
beta_2=0.999,
|
|
epsilon=epsilon,
|
|
exclude_from_weight_decay=['LayerNorm', 'bias'],
|
|
**additional_args
|
|
)
|
|
# {'LayerNorm/gamma:0', 'LayerNorm/beta:0'}
|
|
return optimizer
|