cayley_adam/grassmann_optimizer.py
def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            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:
            momentum = group['momentum']
            grassmann = group['grassmann']

            if grassmann:
                grad_clip = group['grad_clip']
                omega = group['omega']

                for p in group['params']:
                    if p.grad is None:
                        continue

                    unity,_ = unit(p.data.view(p.size()[0],-1))
                    g = p.grad.data.view(p.size()[0],-1)

                    if omega != 0:
                      # L=|Y'Y-I|^2/2=|YY'-I|^2/2+c
                      # dL/dY=2(YY'Y-Y)
                      g.add_(2*omega, torch.mm(torch.mm(unity, unity.t()), unity) - unity)

                    h = gproj(unity, g)

                    if grad_clip is not None:
                        h_hat = clip_by_norm(h, grad_clip)
                    else:
                        h_hat = h

                    param_state = self.state[p]
                    if 'momentum_buffer' not in param_state:
                        param_state['momentum_buffer'] = torch.zeros(h_hat.size())
                        if p.is_cuda:
                          param_state['momentum_buffer'] = param_state['momentum_buffer'].cuda()

                    mom = param_state['momentum_buffer']
                    mom_new = momentum*mom - group['lr']*h_hat

                    p.data.copy_(gexp(unity, mom_new).view(p.size()))
                    mom.copy_(gpt(unity, mom_new))

            else:
                # This routine is from https://github.com/pytorch/pytorch/blob/master/torch/optim/sgd.py
                weight_decay = group['weight_decay']
                dampening = group['dampening']
                nesterov = group['nesterov']
                for p in group['params']:
                    if p.grad is None:
                        continue
                    d_p = p.grad.data
                    if weight_decay != 0:
                        d_p.add_(weight_decay, p.data)
                    if momentum != 0:
                        param_state = self.state[p]
                        if 'momentum_buffer' not in param_state:
                            buf = param_state['momentum_buffer'] = d_p.clone()
                        else:
                            buf = param_state['momentum_buffer']
                            buf.mul_(momentum).add_(1 - dampening, d_p)
                        if nesterov:
                            d_p = d_p.add(momentum, buf)
                        else:
                            d_p = buf

                    p.data.add_(-group['lr'], d_p)

        return loss
Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        
Performs a single optimization step.
Performs a single optimization step .
def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()
        for group in self.param_groups:
            momentum = group['momentum']
            grassmann = group['grassmann']
            if grassmann:
                grad_clip = group['grad_clip']
                omega = group['omega']
                for p in group['params']:
                    if p.grad is None:
                        continue
                    unity,_ = unit(p.data.view(p.size()[0],-1))
                    g = p.grad.data.view(p.size()[0],-1)
                    if omega != 0:
                      g.add_(2*omega, torch.mm(torch.mm(unity, unity.t()), unity) - unity)
                    h = gproj(unity, g)
                    if grad_clip is not None:
                        h_hat = clip_by_norm(h, grad_clip)
                    else:
                        h_hat = h
                    param_state = self.state[p]
                    if 'momentum_buffer' not in param_state:
                        param_state['momentum_buffer'] = torch.zeros(h_hat.size())
                        if p.is_cuda:
                          param_state['momentum_buffer'] = param_state['momentum_buffer'].cuda()
                    mom = param_state['momentum_buffer']
                    mom_new = momentum*mom - group['lr']*h_hat
                    p.data.copy_(gexp(unity, mom_new).view(p.size()))
                    mom.copy_(gpt(unity, mom_new))
            else:
                weight_decay = group['weight_decay']
                dampening = group['dampening']
                nesterov = group['nesterov']
                for p in group['params']:
                    if p.grad is None:
                        continue
                    d_p = p.grad.data
                    if weight_decay != 0:
                        d_p.add_(weight_decay, p.data)
                    if momentum != 0:
                        param_state = self.state[p]
                        if 'momentum_buffer' not in param_state:
                            buf = param_state['momentum_buffer'] = d_p.clone()
                        else:
                            buf = param_state['momentum_buffer']
                            buf.mul_(momentum).add_(1 - dampening, d_p)
                        if nesterov:
                            d_p = d_p.add(momentum, buf)
                        else:
                            d_p = buf
                    p.data.add_(-group['lr'], d_p)
        return loss
def step ( self , closure = None ) : loss = None if closure is not None : loss = closure ( ) for group in self . param_groups : momentum = group [ 'momentum' ] grassmann = group [ 'grassmann' ] if grassmann : grad_clip = group [ 'grad_clip' ] omega = group [ 'omega' ] for p in group [ 'params' ] : if p . grad is None : continue unity , _ = unit ( p . data . view ( p . size ( ) [ 0 ] , - 1 ) ) g = p . grad . data . view ( p . size ( ) [ 0 ] , - 1 ) if omega != 0 : g . add_ ( 2 * omega , torch . mm ( torch . mm ( unity , unity . t ( ) ) , unity ) - unity ) h = gproj ( unity , g ) if grad_clip is not None : h_hat = clip_by_norm ( h , grad_clip ) else : h_hat = h param_state = self . state [ p ] if 'momentum_buffer' not in param_state : param_state [ 'momentum_buffer' ] = torch . zeros ( h_hat . size ( ) ) if p . is_cuda : param_state [ 'momentum_buffer' ] = param_state [ 'momentum_buffer' ] . cuda ( ) mom = param_state [ 'momentum_buffer' ] mom_new = momentum * mom - group [ 'lr' ] * h_hat p . data . copy_ ( gexp ( unity , mom_new ) . view ( p . size ( ) ) ) mom . copy_ ( gpt ( unity , mom_new ) ) else : weight_decay = group [ 'weight_decay' ] dampening = group [ 'dampening' ] nesterov = group [ 'nesterov' ] for p in group [ 'params' ] : if p . grad is None : continue d_p = p . grad . data if weight_decay != 0 : d_p . add_ ( weight_decay , p . data ) if momentum != 0 : param_state = self . state [ p ] if 'momentum_buffer' not in param_state : buf = param_state [ 'momentum_buffer' ] = d_p . clone ( ) else : buf = param_state [ 'momentum_buffer' ] buf . mul_ ( momentum ) . add_ ( 1 - dampening , d_p ) if nesterov : d_p = d_p . add ( momentum , buf ) else : d_p = buf p . data . add_ ( - group [ 'lr' ] , d_p ) return loss
Performs a single optimization step.
Performs a single optimization step .
"""Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """ # L=|Y'Y-I|^2/2=|YY'-I|^2/2+c
 # dL/dY=2(YY'Y-Y)
 # This routine is from https://github.com/pytorch/pytorch/blob/master/torch/optim/sgd.py

self
closure
A closure that reevaluates the model
and returns the loss.
A closure that reevaluates the model and returns the loss .