File size: 1,733 Bytes
ae81e0f |
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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
"""
Optimizer and schedulers
"""
import torch
import torch.nn as nn
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler
def get_optimizer(optim: str, model: nn.Module, **kwargs: any) -> Optimizer:
"""
Return training optimizer
"""
if optim == 'sgd':
return torch.optim.SGD(model.parameters(), **kwargs)
elif optim == 'adam':
return torch.optim.Adam(model.parameters(), **kwargs)
elif optim in ['adamw', 'adamw_torch']:
return torch.optim.AdamW(model.parameters(), **kwargs)
elif optim == 'adamw_torch_fused':
return torch.optim.AdamW(model.parameters(), **kwargs, fused=True)
elif optim == 'adafactor':
from transformers import Adafactor
kwargs['relative_step'] = False # for now
return Adafactor(model.parameters(), **kwargs)
else:
raise NotImplementedError(f"{optim} optimizer not implemented sorry.")
def get_scheduler(lr_scheduler_type: str, optimizer: Optimizer,
**kwargs: any) -> LRScheduler:
"""
Return learning rate scheduler
"""
if lr_scheduler_type in ['plateau', 'reduce_lr_on_plateau']:
from torch.optim.lr_scheduler import ReduceLROnPlateau
return ReduceLROnPlateau(optimizer=optimizer, **kwargs)
elif lr_scheduler_type == 'cosine_warmup':
from transformers import get_cosine_schedule_with_warmup
return get_cosine_schedule_with_warmup(optimizer=optimizer, **kwargs)
elif lr_scheduler_type in ['linear_warmup', 'linear']:
from transformers import get_linear_schedule_with_warmup
return get_linear_schedule_with_warmup(optimizer=optimizer, **kwargs)
else:
return None |