adversarial_VLNDUET/pretrain_src/optim/rangerlars.py
Shizhe Chen 747cf0587b init
2021-11-24 13:29:08 +01:00

15 lines
519 B
Python

import torch, math
from torch.optim.optimizer import Optimizer
import itertools as it
from .lookahead import *
from .ralamb import *
# RAdam + LARS + LookAHead
# Lookahead implementation from https://github.com/lonePatient/lookahead_pytorch/blob/master/optimizer.py
# RAdam + LARS implementation from https://gist.github.com/redknightlois/c4023d393eb8f92bb44b2ab582d7ec20
def RangerLars(params, alpha=0.5, k=6, *args, **kwargs):
ralamb = Ralamb(params, *args, **kwargs)
return Lookahead(ralamb, alpha, k)