File size: 885 Bytes
1a4c139 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
#!/usr/bin/python3
# -*- coding: utf-8 -*-
from typing import Any, Dict, List, Tuple
from allennlp.training.optimizers import Registrable, Optimizer, make_parameter_groups
from pytorch_pretrained_bert.optimization import BertAdam
import torch
@Optimizer.register("bert_adam")
class BertAdamOptimizer(Optimizer, BertAdam):
def __init__(
self,
model_parameters: List[Tuple[str, torch.nn.Parameter]],
parameter_groups: List[Tuple[List[str], Dict[str, Any]]] = None,
lr: float = 5e-5,
warmup: float = 0.1,
t_total: int = 50000,
schedule: str = 'warmup_linear',
):
super().__init__(
params=make_parameter_groups(model_parameters, parameter_groups),
lr=lr,
warmup=warmup,
t_total=t_total,
schedule=schedule,
)
if __name__ == '__main__':
pass
|