Source code for imaginaire.optimizers.fromage

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out
# import torch
import math
from torch.optim.optimizer import Optimizer, required
[docs]class Fromage(Optimizer): r"""Fromage optimizer implementation (""" def __init__(self, params, lr=required, momentum=0): if lr is not required and lr < 0.0: raise ValueError("Invalid learning rate: {}".format(lr)) defaults = dict(lr=lr, momentum=momentum) super(Fromage, self).__init__(params, defaults)
[docs] def step(self, closure=None): r"""Performs a single optimization step. Args: closure (callable, optional): A closure that reevaluates the model and returns the loss. """ loss = None if closure is not None: loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue d_p = d_p_norm = p.grad.norm() p_norm = p.norm() if p_norm > 0.0 and d_p_norm > 0.0:['lr'], d_p * (p_norm / d_p_norm)) else:['lr'], d_p) /= math.sqrt(1 + group['lr'] ** 2) return loss