optim.py

import torch
import torch.optim as optim


def get_optimizer(config, model):
    param_biases = [p for p in model.parameters() if p.ndim == 1]
    param_weights = [p for p in model.parameters() if p.ndim != 1]

    parameters = [
        {"params": param_weights, "lr": config.lr_weights},
        {"params": param_biases, "lr": config.lr_biases},
    ]
    if config.optimizer.upper() == "LARS":
        optimizer = LARS(
            parameters,
            lr=0,
            weight_decay=config.weight_decay,
            weight_decay_filter=True,
            lars_adaptation_filter=True,
        )
    elif config.optimizer.upper() == "SGD":
        optimizer = optim.SGD(
            parameters,
            lr=config.lr_weights,
            momentum=config.momentum,
            weight_decay=config.weight_decay,
        )
    else:
        raise ValueError(f"Unknown optimizer {config.optimizer}")
    return optimizer


class LARS(optim.Optimizer):
    def __init__(
        self,
        params,
        lr,
        weight_decay=0,
        momentum=0.9,
        eta=0.001,
        weight_decay_filter=False,
        lars_adaptation_filter=False,
    ):
        defaults = dict(
            lr=lr,
            weight_decay=weight_decay,
            momentum=momentum,
            eta=eta,
            weight_decay_filter=weight_decay_filter,
            lars_adaptation_filter=lars_adaptation_filter,
        )
        super().__init__(params, defaults)

    def exclude_bias_and_norm(self, p):
        return p.ndim == 1

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            for p in g["params"]:
                dp = p.grad

                if dp is None:
                    continue

                if not g["weight_decay_filter"] or not self.exclude_bias_and_norm(p):
                    dp = dp.add(p, alpha=g["weight_decay"])

                if not g["lars_adaptation_filter"] or not self.exclude_bias_and_norm(p):
                    param_norm = torch.norm(p)
                    update_norm = torch.norm(dp)
                    one = torch.ones_like(param_norm)
                    q = torch.where(
                        param_norm > 0.0,
                        torch.where(
                            update_norm > 0, (g["eta"] * param_norm / update_norm), one
                        ),
                        one,
                    )
                    dp = dp.mul(q)

                param_state = self.state[p]
                if "mu" not in param_state:
                    param_state["mu"] = torch.zeros_like(p)
                mu = param_state["mu"]
                mu.mul_(g["momentum"]).add_(dp)

                p.add_(mu, alpha=-g["lr"])