def register_lr_scheduler(cls, name=None):
return register(cls, MEAD_LAYERS_LR_SCHEDULERS, name, "lr_scheduler")
PiecewiseDecaySchedulerTensorFlow = PiecewiseDecaySchedulerTensorFlow1
ExponentialDecaySchedulerTensorFlow = ExponentialDecaySchedulerTensorFlow1
InverseTimeDecaySchedulerTensorFlow = InverseTimeDecaySchedulerTensorFlow1
ZarembaDecaySchedulerTensorFlow = ZarembaDecaySchedulerTensorFlow1
else:
ConstantSchedulerTensorFlow = ConstantSchedulerTensorFlow2
WarmupLinearSchedulerTensorFlow = WarmupLinearSchedulerTensorFlow2
CyclicLRSchedulerTensorFlow = CyclicLRSchedulerTensorFlow2
SGDRSchedulerTensorFlow = SGDRSchedulerTensorFlow2
CompositeLRSchedulerTensorFlow = CompositeLRSchedulerTensorFlow2
PiecewiseDecaySchedulerTensorFlow = PiecewiseDecaySchedulerTensorFlow2
ExponentialDecaySchedulerTensorFlow = ExponentialDecaySchedulerTensorFlow2
InverseTimeDecaySchedulerTensorFlow = InverseTimeDecaySchedulerTensorFlow2
ZarembaDecaySchedulerTensorFlow = ZarembaDecaySchedulerTensorFlow2
register(ConstantSchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS, "default",
"lr_scheduler")
register(WarmupLinearSchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS,
"warmup_linear", "lr_scheduler")
register(CyclicLRSchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS, "clr",
"lr_scheduler")
register(SGDRSchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS, "sgdr",
"lr_scheduler")
register(CompositeLRSchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS,
"composite", "lr_scheduler")
register(PiecewiseDecaySchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS,
"piecewise", "lr_scheduler")
register(ZarembaDecaySchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS, "zaremba",
"lr_scheduler")
register(CosineDecaySchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS, "cosine",
"lr_scheduler")
register(InverseTimeDecaySchedulerTensorFlow, MEAD_LAYERS_LR_SCHEDULERS,
def register_decoder(cls, name=None):
"""Register a function as a plug-in"""
return register(cls, BASELINE_SEQ2SEQ_DECODERS, name, 'decoder')
def register_progress(cls, name=None):
return register(cls, MEAD_LAYERS_PROGRESS, name, 'progress')
def register_arc_policy(cls, name=None):
"""Register a function as a plug-in"""
return register(cls, BASELINE_SEQ2SEQ_ARC_POLICY, name, 'decoder')
def register_cache_backend(cls, name):
"""Register a function as a plug-in"""
return register(cls, ODIN_CACHE_BACKENDS, name, 'cache')
def register_store_backend(cls, name):
"""Register a function as a plug-in"""
return register(cls, ODIN_STORE_BACKENDS, name, 'backend')
def register_tokenizer(cls, name=None):
r = register(cls, TOKENIZER_DEFS, name, "tokenizer defs")
return r
def register_masking(cls, name=None):
"""Register a class as a handler for masking rules by key name"""
r = register(cls, MASKING_RULE_DEFS, name, "masking rule defs")
return r