gizml:nn

what training a network needs on top of gizml:grad. Layers are functions of a Graph and their parameters (plain Tensors the caller keeps), losses are Graph values, and optimizers update the parameters in fused passes. Files of parameters are gizml:tensor's (save, load).

import * as Nn from "gizml:nn";

Types

Sgd

type Sgd

Stochastic gradient descent with momentum, as torch.optim.SGD: v = momentum v + g, p -= rate v. Without momentum there is no v, and p - rate g is one pass (or the epilogue of the product that makes g).

• rate: f64
• momentum: f64

Adam

type Adam

Adam (Kingma and Ba) with decoupled weight decay (AdamW when decay > 0), as torch.optim.AdamW.

• rate: f64
• beta1: f64
• beta2: f64
• epsilon: f64
• decay: f64

GptConfig

type GptConfig = { vocab: i64; context: i64; dim: i64; heads: i64; layers: i64; dropout: f64; }

A decoder-only transformer as GPT-2 builds one: pre-norm blocks of causal multi-head attention and a 4x GELU MLP, learned positions, a final norm and an unbiased head.

• vocab: i64
• context: i64
• dim: i64
• heads: i64
• layers: i64
• dropout: f64

KvCache

type KvCache = { keys: Vec<Tensor>; values: Vec<Tensor>; length: i64; }

Inference state: per layer the keys and values of the tokens so far, (heads, tokens, hd) each, so a new token attends without recomputing them.

• keys: Vec<M32.Tensor>
• values: Vec<M32.Tensor>
• length: i64

Functions

linearInit

function linearInit(k: Key, inputs: i64, outputs: i64): Vec<Tensor>

A dense layer's parameters, [weight (inputs, outputs), bias (outputs)], uniform in +-1/sqrt(inputs) as PyTorch's nn.Linear starts.

Effects: fail.

linear

function linear(g: mut Graph, x: Var, weight: Var, bias: Var): Var

x (..., inputs) times weight plus bias.

Effects: fail, read, write, console, io.

inputs

function inputs(g: mut Graph, params: Vec<Tensor>): Vec<Var>

Each tensor of PARAMS as an input of G, in order.

Effects: fail.

gradients

function gradients(grads: Grads, vars: Vec<Var>): Vec<T.Lazy>

The gradient of each of VARS, in order (for an optimizer's step), not yet computed: the step's update runs as each is stored.

Effects: fail, read, write, console, io.

mse

function mse(g: mut Graph, y: Var, target: Var): Var

The mean of (y - target)^2.

Effects: fail, read, write, console, io.

oneHot

function oneHot(labels: Vec<i64>, classes: i64): Tensor

Rows of 0s with a 1 at each label: (labels, classes).

Effects: fail.

accuracy

function accuracy(logits: Tensor, labels: Vec<i64>): f64

The share of rows whose largest logit is at the label.

Effects: fail, read, write, console, io.

sgd

function sgd(rate: f64, momentum: f64): Sgd

A stochastic gradient descent optimizer with this learning rate and momentum.

adam

function adam(rate: f64): Adam

An Adam optimizer with this learning rate.

adamW

function adamW(rate: f64, decay: f64): Adam

An AdamW optimizer with this learning rate and weight decay.

gptInit

function gptInit(k: Key, c: GptConfig): Vec<Tensor>

Parameters in order: token and position embeddings; per layer ln1 (w, b), qkv (w, b), proj (w, b), ln2 (w, b), fc (w, b), fc2 (w, b); final norm (w, b); head. Weights normal(0, 0.02), biases 0, norms 1 and 0 (GPT-2's start).

Effects: fail, read, write, console, io.

gpt

function gpt(g: mut Graph, c: GptConfig, p: Vec<Var>, ids: Vec<i64>, batch: i64, t: i64, k: Option<Key>): Var

Logits (batch * t, vocab) of the token ids (batch rows of t). key draws the dropout masks when training (c.dropout > 0).

Effects: fail, read, write, console, io.

kvCache

function kvCache(): KvCache

An empty cache.

gptStep

function gptStep(c: GptConfig, p: Vec<Tensor>, cache: mut KvCache, id: i64): Tensor

The logits (vocab) after token id at the cache's next position, extending the cache.

Effects: fail, read, write, console, io.

generate

function generate(c: GptConfig, p: Vec<Tensor>, prompt: Vec<i64>, count: i64): Vec<i64>

count tokens after prompt, each the most likely (greedy), with a KV cache.

Effects: fail, read, write, console, io.

conv2dInit

function conv2dInit(k: Key, inputs: i64, outputs: i64, size: i64): Vec<Tensor>

A convolution's parameters, [weight (outputs, inputs, k, k), bias (outputs)], uniform in +-1/sqrt(inputs k k) as torch's nn.Conv2d starts.

Effects: fail.

onGpu

function onGpu(params: Vec<Tensor>, gpu: Gpu): Vec<Tensor>

The parameters on gpu.

Effects: fail, read, write, console, io.

Methods

Sgd.step

function Sgd.step(self: mut Sgd, params: mut Vec<Tensor>, grads: Vec<T.Lazy>): void

One step: moves PARAMS against GRADS.

Effects: fail, read, write, console, io.

Adam.step

function Adam.step(self: mut Adam, params: mut Vec<Tensor>, grads: Vec<T.Lazy>): void

One step: moves PARAMS against GRADS.

Effects: fail, read, write, console, io.