elixir-nx / polaris Goto Github PK

From elixir-nx/axon#452

Most docs suggest using this library like so:

Polaris.Optimizers.adam(0.005)

However, this causes an error, because the args expect a keyword list

** (FunctionClauseError) no function clause matching in Keyword.pop/3    
    
    The following arguments were given to Keyword.pop/3:
    
        # 1
        0.005
    
        # 2
        :learning_rate
    
        # 3
        0.001
    
    Attempted function clauses (showing 1 out of 1):
    
        def pop(keywords, key, default) when is_list(keywords) and is_atom(key)
    
    (elixir 1.15.7) lib/keyword.ex:1289: Keyword.pop/3
    (polaris 0.1.0) lib/polaris/optimizers.ex:109: Polaris.Optimizers.adam/1

Looks like this was intentional: 228f829

So this is an issue to update the docs 😄

From elixir-nx/axon#451

Pytorch has the one cycle LR scheduler: https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.OneCycleLR.html

It seems to me that implementing this kind of scheduler is not possible at the moment because scheduler in Polaris have an API that takes as input the step and returns a learning rate.

In the one cycle LR scenario, the scheduler should also be able to mutate the momentum of the optimizer (so affect the optimizer state).

Am I missing something ?