Salam
Online editor

Standard library · Machine learning

nn

Neural network layers: linear, convolutional, recurrent, transformer, and checkpoints.

import nn
On this page
  1. Types
  2. Functions

This package has ready made layers for building a model. The weights of every layer live in one ParamStore; make each layer with a New... function and a unique name, and call forward on the tape at every step. Sequential runs several layers one after the other.

Once it is trained with optim, save the model with SaveCheckpoint and load it again later.

Types

Parameter

struct Parameter
One learnable parameter.
NameTypeDescription
namestr = ""The name.
valuetensor.TensorThe value.
gradtensor.TensorThe gradient.
nn-Layers.salam
import nn
import autograd
import optim
import tensor
import os
import compress

func main:
    mut store := nn.ParamStore {}
    mut rng := tensor.NewRng(3)
    fc1 := nn.NewLinear(store, "fc1", 2, 8, rng)
    fc2 := nn.NewLinearNoBias(store, "fc2", 8, 1, rng)
    mut model := nn.Sequential {}
    model.add(fc1)
    model.add(nn.TanhLayer {})
    model.add(fc2)
    println store.count()

    mut inputs := tensor.Reshape(tensor.Arange(8), tensor.Shape2(4, 2))
    tensor.ScaleInPlace(inputs, 0.25)
    mut target := tensor.Zeros2(4, 1)
    repeat 4 in i:
        target.set_at(i, inputs.at2(i, 0) + inputs.at2(i, 1))
    end
    mut adam := optim.Adam { lr = 0.05 }
    mut tape := autograd.NewTape(1)
    mut last := 0.0
    repeat 300:
        tape.reset()
        store.zero_grad()
        output := model.forward(tape, store, autograd.Input(tape, inputs))
        loss := autograd.MSELoss(tape, output, autograd.Constant(tape, target))
        autograd.Backward(tape, loss)
        adam.step(store)
        last = autograd.Value(tape, loss).at(0)
    end
    println last < 0.01
    mut p := store.at(fc1.w)
    println p.name, tensor.ShapeString(p.value), tensor.ShapeString(p.grad)
    weight_node := nn.Bind(tape, store, fc1.w)
    println tensor.ShapeString(autograd.Value(tape, weight_node))

    path := os.TempDir() + "/salam_mlp_en.ckpt"
    println nn.SaveCheckpoint(path, store), nn.LoadCheckpoint(path, store)
    mut buffer := nn.SaveCheckpointBuf(store)
    println nn.LoadCheckpointBuf(buffer, store)
    compress.BufFree(buffer)

    mut images := tensor.Ones(tensor.Shape4(2, 1, 8, 8))
    conv := nn.NewConv2D(store, "conv", 1, 4, 3, 1, 1, rng)
    mut batch_norm := nn.NewBatchNorm(store, "bn", 4)
    group_norm := nn.NewGroupNorm(store, "gn", 4, 2)
    pool := nn.MaxPool2D { size = 2, stride = 2 }
    average := nn.AvgPool2D {}
    h := conv.forward(tape, store, autograd.Input(tape, images))
    h2 := group_norm.forward(tape, store, batch_norm.forward(tape, store, h))
    h3 := average.forward(tape, pool.forward(tape, nn.ReluLayer {}.forward(tape, store, h2)))
    flat := nn.Flatten(tape, h3)
    println tensor.ShapeString(autograd.Value(tape, h)), tensor.ShapeString(autograd.Value(tape, flat))

    row := autograd.Input(tape, tensor.Ones(tensor.Shape2(2, 4)))
    layer_norm := nn.NewLayerNorm(store, "ln", 4)
    rms_norm := nn.NewRMSNorm(store, "rms", 4)
    drop := nn.Dropout { p = 0.0 }
    gelu := nn.GeluLayer {}
    out := gelu.forward(tape, store, drop.forward(tape, rms_norm.forward(tape, store, layer_norm.forward(tape, store, row))))
    println tensor.ShapeString(autograd.Value(tape, out))

    mut weight := tensor.Zeros2(64, 64)
    nn.XavierUniform(weight, 64, 64, rng)
    println tensor.MaxAll(tensor.Abs(weight)) <= 0.22
    nn.XavierNormal(weight, 64, 64, rng)
    nn.KaimingUniform(weight, 64, rng)
    println tensor.MaxAll(tensor.Abs(weight)) <= 0.31
    nn.KaimingNormal(weight, 64, rng)
    println tensor.MaxAll(tensor.Abs(weight)) > 0.0

    tape.free()
    adam.free()
    model.free()
    batch_norm.free()
    store.free()
end
Output3 true fc1.w 2x8 2x8 2x8 true true true 2x4x8x8 2x16 2x4 true true true

ParamStore

struct ParamStore
The parameter store; every weight of the model lives there and the layers point at them by index.
NameTypeDescription
itemsVectorThe parameters.

The example for this item is shown with the ones above: nn-Layers.salam

add

func add(name: str, t: tensor.Tensor): int
Adds a parameter and returns its index.

The example for this item is shown with the ones above: nn-Layers.salam

count

func count(): int
The number of parameters.

The example for this item is shown with the ones above: nn-Layers.salam

at

func at(h: int): Parameter
Parameter number h.

The example for this item is shown with the ones above: nn-Layers.salam

zero_grad

func zero_grad()
Zeroes the gradients; at the start of every step.

The example for this item is shown with the ones above: nn-Layers.salam

free

func free()
Frees all of them.

The example for this item is shown with the ones above: nn-Layers.salam

Linear

struct Linear
A linear layer: xW+b.
NameTypeDescription
wint = -1The weight index.
bint = -1The bias index.
in_featuresint = 0The input size.
out_featuresint = 0The output size.
has_biasbool = trueIt has a bias.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

LSTM

struct LSTM
A long short term memory network.
NameTypeDescription
wxLinearThe input weight.
whLinearThe state weight.
hiddenint = 0The state size.
nn-Sequence.salam
import nn
import autograd
import tensor

func shape(tape &: autograd.Tape, node: int): str:
    ret tensor.ShapeString(autograd.Value(tape, node))
end

func main:
    batch := 2
    length := 3
    dim := 8
    mut store := nn.ParamStore {}
    mut rng := tensor.NewRng(11)
    mut tape := autograd.NewTape(1)

    embedding := nn.NewEmbedding(store, "embed", 10, dim, rng)
    mut tokens := Vector {} as Vector<int>
    each t in [1, 2, 3, 4, 5, 6]:
        tokens.push(t)
    end
    h := embedding.forward(tape, store, tokens)
    mut positions := nn.SinusoidalPositions(batch, length, dim)
    h2 := autograd.Add(tape, h, autograd.Input(tape, positions))
    println shape(tape, h), tensor.ShapeString(positions)

    attention := nn.NewMultiHeadAttention(store, "attn", dim, 2, 0.0, rng)
    println shape(tape, attention.forward(tape, store, h2, batch, length))
    encoder := nn.NewTransformerEncoderLayer(store, "enc", dim, 2, 16, 0.0, rng)
    println shape(tape, encoder.forward(tape, store, h2, batch, length))
    decoder := nn.NewTransformerDecoderLayer(store, "dec", dim, 2, 16, 0.0, rng)
    println shape(tape, decoder.forward(tape, store, h2, batch, length))
    mut mask := nn.CausalMask(length)
    println tensor.ToString(tensor.LessScalar(mask, -1.0))

    mut sequence := tensor.Ones(tensor.Shape2(length * batch, 4))
    lstm := nn.NewLSTM(store, "lstm", 4, 6, rng)
    gru := nn.NewGRU(store, "gru", 4, 6, rng)
    sequence_input := autograd.Input(tape, sequence)
    states := lstm.forward(tape, store, sequence_input, length, batch)
    states2 := gru.forward(tape, store, sequence_input, length, batch)
    println shape(tape, states), shape(tape, states2), shape(tape, nn.LastStep(tape, states, length, batch))
    println store.count() > 10
    tape.free()
    store.free()
end
Output6x8 6x8 6x8 6x8 6x8 tensor(3x3)[0 1 1 0 0 1 0 0 0] 6x6 6x6 2x6 true

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int, seq: int, batch: int): int
Returns the state of every step.

The example for this item is shown with the ones above: nn-Sequence.salam

GRU

struct GRU
A gated recurrent unit.
NameTypeDescription
wxLinearThe input weight.
whLinearThe state weight.
hiddenint = 0The state size.

The example for this item is shown with the ones above: nn-Sequence.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int, seq: int, batch: int): int
Returns the state of every step.

The example for this item is shown with the ones above: nn-Sequence.salam

Layer

interface Layer
The layer interface: anything with a forward that can sit in a Sequential.

The example for this item is shown with the ones above: nn-Layers.salam

ReluLayer

struct ReluLayer
A rectifier layer for a Sequential.
NameTypeDescription
unusedint = 0Unused.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, _ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

GeluLayer

struct GeluLayer
A GELU layer for a Sequential.
NameTypeDescription
unusedint = 0Unused.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, _ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

TanhLayer

struct TanhLayer
A hyperbolic tangent layer for a Sequential.
NameTypeDescription
unusedint = 0Unused.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, _ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

Sequential

struct Sequential
A chain of layers that run one after the other.
NameTypeDescription
layersVectorThe layers.

The example for this item is shown with the ones above: nn-Layers.salam

add

func add(l: dyn Layer)
Adds a layer to the end.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int): int
Passes the input through every layer.

The example for this item is shown with the ones above: nn-Layers.salam

free

func free()
Frees the chain.

The example for this item is shown with the ones above: nn-Layers.salam

Conv2D

struct Conv2D
A two dimensional convolution layer.
NameTypeDescription
wint = -1The weight index.
bint = -1The bias index.
strideint = 1The stride.
padint = 0The padding.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

MaxPool2D

struct MaxPool2D
A two dimensional max pooling layer.
NameTypeDescription
sizeint = 2The window size.
strideint = 2The stride.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

AvgPool2D

struct AvgPool2D
A two dimensional average pooling layer.
NameTypeDescription
sizeint = 2The window size.
strideint = 2The stride.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

MultiHeadAttention

struct MultiHeadAttention
Multi head attention.
NameTypeDescription
wqLinearThe query layer.
wkLinearThe key layer.
wvLinearThe value layer.
woLinearThe output layer.
dimint = 0The size.
headsint = 1The number of heads.
dropoutf64 = 0The drop probability.
causalbool = falseEach position sees only the ones before it.

The example for this item is shown with the ones above: nn-Sequence.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int, batch: int, seq: int): int
Self attention over the input, with a batch and a length.

The example for this item is shown with the ones above: nn-Sequence.salam

TransformerEncoderLayer

struct TransformerEncoderLayer
One transformer layer: attention and a feed forward network, each with normalisation and a residual connection.
NameTypeDescription
attnMultiHeadAttentionThe attention.
ln1LayerNormThe first normalisation.
ln2LayerNormThe second normalisation.
ff1LinearThe first feed forward layer.
ff2LinearThe second feed forward layer.
dropoutf64 = 0The drop probability.

The example for this item is shown with the ones above: nn-Sequence.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int, batch: int, seq: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Sequence.salam

Dropout

struct Dropout
A dropout layer.
NameTypeDescription
pf64 = 0.1The drop probability.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

Embedding

struct Embedding
An embedding table: one vector per word index.
NameTypeDescription
wint = -1The table index.
vocabint = 0The vocabulary size.
dimint = 0The vector size.

The example for this item is shown with the ones above: nn-Sequence.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, indices &: Vector<int>): int
Takes the vectors of the given indices.

The example for this item is shown with the ones above: nn-Sequence.salam

LayerNorm

struct LayerNorm
Layer normalisation.
NameTypeDescription
gammaint = -1The gamma index.
betaint = -1The beta index.
epsf64 = 1e-05The small number for stability.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

BatchNorm

struct BatchNorm
Batch normalisation.
NameTypeDescription
gammaint = -1The gamma index.
betaint = -1The beta index.
running_meantensor.TensorThe running mean, for prediction.
running_vartensor.TensorThe running variance.
momentumf64 = 0.1The coefficient for updating the running statistics.
epsf64 = 1e-05The small number for stability.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

free

func free()
Frees the running statistics.

The example for this item is shown with the ones above: nn-Layers.salam

GroupNorm

struct GroupNorm
Group normalisation.
NameTypeDescription
gammaint = -1The gamma index.
betaint = -1The beta index.
groupsint = 1The number of groups.
epsf64 = 1e-05The small number for stability.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

RMSNorm

struct RMSNorm
Root mean square normalisation.
NameTypeDescription
gammaint = -1The gamma index.
epsf64 = 1e-05The small number for stability.

The example for this item is shown with the ones above: nn-Layers.salam

forward

func forward(tp &: autograd.Tape, ps &: ParamStore, x: int): int
Works out the output of the layer on the tape.

The example for this item is shown with the ones above: nn-Layers.salam

Functions

Bind

func Bind(tp &: autograd.Tape, ps &: ParamStore, h: int): int
Binds a parameter to the tape and returns its node.

The example for this item is shown with the ones above: nn-Layers.salam

NewLinear

func NewLinear(ps &: ParamStore, name: str, in_f: int, out_f: int, g &: tensor.Rng): Linear
A linear layer with Kaiming initial weights.

The example for this item is shown with the ones above: nn-Layers.salam

NewLinearNoBias

func NewLinearNoBias(ps &: ParamStore, name: str, in_f: int, out_f: int, g &: tensor.Rng): Linear
A linear layer with no bias.

The example for this item is shown with the ones above: nn-Layers.salam

NewLSTM

func NewLSTM(ps &: ParamStore, name: str, in_f: int, hidden: int, g &: tensor.Rng): LSTM
An LSTM with an input and a state size.

The example for this item is shown with the ones above: nn-Sequence.salam

NewGRU

func NewGRU(ps &: ParamStore, name: str, in_f: int, hidden: int, g &: tensor.Rng): GRU
A GRU with an input and a state size.

The example for this item is shown with the ones above: nn-Sequence.salam

LastStep

func LastStep(tp &: autograd.Tape, out: int, seq: int, batch: int): int
The state of the last step of each sequence; for sequence classification.

The example for this item is shown with the ones above: nn-Sequence.salam

NewConv2D

func NewConv2D(ps &: ParamStore, name: str, in_ch: int, out_ch: int, k: int, stride: int, pad: int, g &: tensor.Rng): Conv2D
A convolution layer with channels, a kernel size, a stride and padding.

The example for this item is shown with the ones above: nn-Layers.salam

Flatten

func Flatten(tp &: autograd.Tape, x: int): int
Folds every dimension but the first into one; between a convolution and a linear layer.

The example for this item is shown with the ones above: nn-Layers.salam

XavierUniform

func XavierUniform(t &: tensor.Tensor, fan_in: int, fan_out: int, g &: tensor.Rng)
Xavier uniform initialisation; for layers with tanh and sigmoid.

The example for this item is shown with the ones above: nn-Layers.salam

XavierNormal

func XavierNormal(t &: tensor.Tensor, fan_in: int, fan_out: int, g &: tensor.Rng)
Xavier normal initialisation.

The example for this item is shown with the ones above: nn-Layers.salam

KaimingUniform

func KaimingUniform(t &: tensor.Tensor, fan_in: int, g &: tensor.Rng)
Kaiming uniform initialisation; for layers with a rectifier.

The example for this item is shown with the ones above: nn-Layers.salam

KaimingNormal

func KaimingNormal(t &: tensor.Tensor, fan_in: int, g &: tensor.Rng)
Kaiming normal initialisation.

The example for this item is shown with the ones above: nn-Layers.salam

NewMultiHeadAttention

func NewMultiHeadAttention(ps &: ParamStore, name: str, dim: int, heads: int, dropout: f64, g &: tensor.Rng): MultiHeadAttention
Multi head attention with a size and a number of heads.

The example for this item is shown with the ones above: nn-Sequence.salam

NewTransformerEncoderLayer

func NewTransformerEncoderLayer(ps &: ParamStore, name: str, dim: int, heads: int, ff_dim: int, dropout: f64, g &: tensor.Rng): TransformerEncoderLayer
An encoder layer.

The example for this item is shown with the ones above: nn-Sequence.salam

CausalMask

func CausalMask(seq: int): tensor.Tensor
The causal mask: negative infinity above the diagonal.

The example for this item is shown with the ones above: nn-Sequence.salam

NewTransformerDecoderLayer

func NewTransformerDecoderLayer(ps &: ParamStore, name: str, dim: int, heads: int, ff_dim: int, dropout: f64, g &: tensor.Rng): TransformerEncoderLayer
A decoder layer: the same, with causal attention for a language model.

The example for this item is shown with the ones above: nn-Sequence.salam

NewEmbedding

func NewEmbedding(ps &: ParamStore, name: str, vocab: int, dim: int, g &: tensor.Rng): Embedding
An embedding table with a vocabulary and a size.

The example for this item is shown with the ones above: nn-Sequence.salam

SinusoidalPositions

func SinusoidalPositions(batch: int, seq: int, dim: int): tensor.Tensor
The sinusoidal position embedding of the original transformer, for a batch and a length.

The example for this item is shown with the ones above: nn-Sequence.salam

NewLayerNorm

func NewLayerNorm(ps &: ParamStore, name: str, dim: int): LayerNorm
Layer normalisation with a size.

The example for this item is shown with the ones above: nn-Layers.salam

NewBatchNorm

func NewBatchNorm(ps &: ParamStore, name: str, channels: int): BatchNorm
Batch normalisation with channels.

The example for this item is shown with the ones above: nn-Layers.salam

NewGroupNorm

func NewGroupNorm(ps &: ParamStore, name: str, channels: int, groups: int): GroupNorm
Group normalisation.

The example for this item is shown with the ones above: nn-Layers.salam

NewRMSNorm

func NewRMSNorm(ps &: ParamStore, name: str, dim: int): RMSNorm
Root mean square normalisation.

The example for this item is shown with the ones above: nn-Layers.salam

SaveCheckpointBuf

func SaveCheckpointBuf(ps &: ParamStore): compress.ByteBuf
The parameters into a buffer.

The example for this item is shown with the ones above: nn-Layers.salam

LoadCheckpointBuf

func LoadCheckpointBuf(buf: compress.ByteBuf, ps &: ParamStore): bool
The parameters from a buffer.

The example for this item is shown with the ones above: nn-Layers.salam

SaveCheckpoint

func SaveCheckpoint(path: str, ps &: ParamStore): bool
Saves every parameter to a file.

The example for this item is shown with the ones above: nn-Layers.salam

LoadCheckpoint

func LoadCheckpoint(path: str, ps &: ParamStore): bool
Loads the parameters from a file by name.

The example for this item is shown with the ones above: nn-Layers.salam