Standard library · Machine learning

optim

Neural network optimisers: SGD, Adam, RMSProp and learning rate schedules.

import optim

On this page

Types
Functions

An optimiser updates the model parameters from the gradients. A training loop always has this shape: reset the autograd tape, zero the gradients of the ParamStore, work out the loss, call Backward, and then the step of the optimiser.

The ...LR functions work out the learning rate for a given step; put it in the lr field of the optimiser before step.

Types

`Adam`

struct Adam

The Adam optimiser; a good default for most networks.

Name	Type	Description
`lr`	`f64 = 0.001`	The learning rate.
`beta1`	`f64 = 0.9`	The coefficient of the gradient average.
`beta2`	`f64 = 0.999`	The coefficient of the squared gradient average.
`eps`	`f64 = 1e-08`	A small number that keeps a division by zero from happening.
`weight_decay`	`f64 = 0`	Weight decay.
`decoupled`	`bool = false`	Apply the weight decay apart from the gradient, meaning AdamW.

optim-Train.salam

import optim
import nn
import autograd
import tensor

func train(method: int, steps: int, target &: tensor.Tensor): f64:
    mut store := nn.ParamStore {}
    handle := store.add("w", tensor.Zeros1(4))
    mut adam := optim.Adam { lr = 0.1 }
    mut adamw := optim.NewAdamW(0.1, 0.01)
    mut rms := optim.RMSProp { lr = 0.05 }
    mut sgd := optim.SGD { lr = 0.05, momentum = 0.9 }
    mut tape := autograd.NewTape(1)
    repeat steps:
        tape.reset()
        store.zero_grad()
        w := nn.Bind(tape, store, handle)
        t := autograd.Input(tape, target)
        loss := autograd.Sum(tape, autograd.Square(tape, autograd.Sub(tape, w, t)))
        autograd.Backward(tape, loss)
        _ := optim.ClipGradNorm(store, 10.0)
        optim.ClipGradValue(store, 5.0)
        if method == 0:
            adam.step(store)
        else method == 1:
            adamw.step(store)
        else method == 2:
            rms.step(store)
        else:
            sgd.step(store)
        end
    end
    mut p := store.at(handle)
    result := p.value.at(3)
    tape.free()
    adam.free()
    adamw.free()
    rms.free()
    sgd.free()
    store.free()
    ret result
end

func main:
    mut target := tensor.Arange(4)
    each method in [0, 1, 2, 3]:
        rounded := (train(method, 400, target) * 100.0 + 0.5) as int
        println method, rounded
    end
    tensor.Free(target)
    println optim.StepLR(1.0, 25, 10, 0.5), optim.CosineLR(1.0, 50, 100, 0.0), optim.WarmupCosine(1.0, 0, 10, 110, 0.0)
end

Output0 300 1 299 2 300 3 300 0.25 0.5 0.1