wehub-resource-sync
93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
import paddle
|
|
import paddle.nn as nn
|
|
|
|
|
|
class OptimizerBenchmarkBase(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
@staticmethod
|
|
def add_args(args, parser):
|
|
raise NotImplementedError
|
|
|
|
def build_optimizer(self, args, learning_rate, model, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
|
|
class SGDBenchmark(OptimizerBenchmarkBase):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@staticmethod
|
|
def add_args(args, parser):
|
|
parser.add_argument("--max_grad_norm", type=float, default=None, help="Norm clip. ")
|
|
|
|
def build_optimizer(self, args, learning_rate, model, **kwargs):
|
|
if getattr(args, "max_grad_norm", None) is not None:
|
|
grad_clip = nn.ClipGradByGlobalNorm(args.max_grad_norm)
|
|
else:
|
|
grad_clip = None
|
|
|
|
optimizer = paddle.optimizer.SGD(
|
|
learning_rate=learning_rate, parameters=model.parameters(), grad_clip=grad_clip
|
|
)
|
|
|
|
return optimizer
|
|
|
|
|
|
class AdamBenchmark(OptimizerBenchmarkBase):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@staticmethod
|
|
def add_args(args, parser):
|
|
parser.add_argument("--max_grad_norm", type=float, default=None, help="Norm clip. ")
|
|
|
|
def build_optimizer(self, args, learning_rate, model, **kwargs):
|
|
if getattr(args, "max_grad_norm", None) is not None:
|
|
grad_clip = nn.ClipGradByGlobalNorm(args.max_grad_norm)
|
|
else:
|
|
grad_clip = None
|
|
|
|
optimizer = paddle.optimizer.Adam(
|
|
learning_rate=learning_rate, parameters=model.parameters(), grad_clip=grad_clip
|
|
)
|
|
|
|
return optimizer
|
|
|
|
|
|
class AdamWBenchmark(OptimizerBenchmarkBase):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@staticmethod
|
|
def add_args(args, parser):
|
|
parser.add_argument("--beta1", type=float, default=0.9, help=". ")
|
|
parser.add_argument("--beta2", type=float, default=0.999, help=". ")
|
|
parser.add_argument("--epsilon", type=float, default=1e-8, help=". ")
|
|
parser.add_argument("--max_grad_norm", type=float, default=None, help=". ")
|
|
parser.add_argument("--weight_decay", type=float, default=0.0, help=". ")
|
|
|
|
def build_optimizer(self, args, learning_rate, model, **kwargs):
|
|
if getattr(args, "max_grad_norm", None) is not None:
|
|
grad_clip = nn.ClipGradByGlobalNorm(args.max_grad_norm)
|
|
else:
|
|
grad_clip = None
|
|
|
|
decay_params = [
|
|
p.name for n, p in model.named_parameters() if not any(nd in n for nd in ["bias", "layer_norm"])
|
|
]
|
|
|
|
optimizer = paddle.optimizer.AdamW(
|
|
learning_rate=learning_rate,
|
|
beta1=args.beta1,
|
|
beta2=args.beta2,
|
|
epsilon=args.epsilon,
|
|
parameters=model.parameters(),
|
|
grad_clip=grad_clip,
|
|
weight_decay=args.weight_decay,
|
|
apply_decay_param_fun=lambda x: x in decay_params,
|
|
)
|
|
|
|
return optimizer
|