Module Ocannl.Demo_models

Demo models

Complete example architectures assembled from Nn_blocks: the models the tutorials, tests and benchmarks train end to end. They are demonstrations of composing the blocks rather than reusable components -- copy one and adapt it.

Attention and transformers

val multi_head_att_workshop : num_heads:Base.int -> d_k:Base.int -> d_v:Base.int -> unit -> Ocannl_tensor.Tensor.t -> Ocannl_tensor.Tensor.t

Multi-head self-attention reduced to its core -- no label, mask, dropout or position embedding -- for illustration; Nn_blocks.multi_head_attention is the reusable block.

val transformer : label:Base.string Base.list -> num_encoder_layers:int -> num_decoder_layers:int -> num_heads:Base__Int.t -> d_enc:Base.int -> d_dec:Base.int -> d_ff:Base.int -> ?epsilon:??? -> ?mask_fill:??? -> ?pos_embed:??? -> unit -> train_step:Ir.Indexing.static_symbol option -> src:Ocannl_tensor.Tensor.t -> tgt:Ocannl_tensor.Tensor.t -> mask:Ocannl_tensor.Tensor.t -> Ocannl_tensor.Tensor.t

An encoder-decoder transformer: token embeddings, position encoding per pos_embed, the Nn_blocks.transformer_encoder and Nn_blocks.transformer_decoder stacks, and an output projection.

val transformer_with_loss : label:'a -> model: (train_step:'b -> src:'c -> tgt:'d -> mask:'e -> Ocannl_tensor.Tensor.t) -> unit -> train_step:'b -> src:'c -> tgt_input:'d -> tgt_target:Ocannl_tensor.Tensor.t -> mask:'e -> Ocannl_tensor.Tensor.t * Ocannl_tensor.Tensor.t

Transformer with teacher forcing for autoregressive training.

TODO: Simplify once tensor shifting/slicing is better supported in shape inference. Currently requires pre-shifted tgt_input (all but last token) and tgt_target (all but first token). During training, the model learns to predict tgt_target given tgt_input.

Convolutional networks

val lenet : ?label:??? -> ?out_channels1:??? -> ?out_channels2:??? -> ?use_padding:??? -> unit -> train_step:'a -> Ocannl_tensor.Tensor.t -> Ocannl_tensor.Tensor.t

LeNet-style architecture for simple image classification (e.g., MNIST). Classic architecture: conv -> pool -> conv -> pool -> fc layers. Output shape is inferred from training data.

val sokoban_cnn : label:Base.string Base.list -> ?num_actions:??? -> unit -> train_step:'a option -> grid_state:Ocannl_tensor.Tensor.t -> Ocannl_tensor.Tensor.t * Ocannl_tensor.Tensor.t

Simple CNN for Sokoban-like grid environments. Processes grid states with three conv layers of 32, 64 and 64 channels and outputs action logits and a value estimate.

val scale_channels : width_mult:Base.Float.t -> int -> Base__Int.t

scale_channels ~width_mult channels is a channel width scaled by a MobileNet-style width multiplier: channels *. width_mult rounded to the NEAREST multiple of 8, a tie (an exact odd multiple of 4) rounding up, and never below 8. So scale_channels ~width_mult:0.75 32 = 24, scale_channels ~width_mult:0.6 64 = 40 (38.4 rounds up), scale_channels ~width_mult:0.6 32 = 16 (19.2 rounds down), scale_channels ~width_mult:1.125 32 = 40 (the tie 36 rounds up), and scale_channels ~width_mult:0.1 32 = 8 (the floor). Multiples of 8 keep channel axes aligned with SIMD lanes and tensor-core tiles. Unlike the MobileNet reference _make_divisible, there is no extra round-up when rounding down loses more than 10%: the nearest multiple is the whole rule. Raises Invalid_argument unless width_mult is finite and positive.

val mobile_cnn : label:Base.string Base.list -> ?num_classes:??? -> ?width_mult:??? -> unit -> train_step:'a option -> Ocannl_tensor.Tensor.t -> Ocannl_tensor.Tensor.t

Modern CNN with depthwise separable convolutions for efficiency, after MobileNet v1's first blocks. Suitable for mobile/edge deployment.

The channel widths at width_mult = 1.0 are 32 (the initial stride-2 conv), then 64, 128, 128 and 256 (the four depthwise separable blocks, at strides 1, 2, 1, 2). Every one of these internal widths is scaled by width_mult through scale_channels (nearest multiple of 8, ties up, at least 8), and the final batch norm and the classifier's input follow the last scaled width. The external dimensions are not scaled: the input channels are whatever the input has, and the output has num_classes logits. Raises Invalid_argument unless width_mult is finite and positive.