CNN (spec model)

This file wires together a small CNN in two styles:

1. A SpecChain description, which is great for "model wiring" and exporters:

• cnn_spec: Conv2D → MaxPool2D → Conv2D → MaxPool2D → Flatten → Linear
• cnn_with_relu_spec: same, but with ReLU inserted after each conv

2. A fully explicit forward/backward pair (Models.Full.CNN2Spec) for the classic training setup:

Conv → ReLU → MaxPool → Conv → ReLU → MaxPool → Flatten → Linear

PyTorch mental model (single image, no batch):

nn.Sequential(
  nn.Conv2d(inC, c1, (kH,kW), stride=stride1, padding=padding1),
  nn.ReLU(),
  nn.MaxPool2d((poolKH,poolKW), stride=poolStride1),
  nn.Conv2d(c1, c2, (kH,kW), stride=stride2, padding=padding2),
  nn.ReLU(),
  nn.MaxPool2d((poolKH,poolKW), stride=poolStride2),
  nn.Flatten(),
  nn.Linear(c2 * H2 * W2, outDim),
)

All shapes are tracked at the type level; the feature dimension for the final LinearSpec is computed as Spec.Shape.size of the post-pooling feature map.

This is intended as a reference/specification of model structure, not as a tuned implementation.

@[reducible, inline]

abbrev Models.CNN.convOut (input kernel stride padding : ℕ) : ℕ

Output size for a convolution along one spatial axis.

The shared shape helper agrees with the usual floor formula when the kernel and stride are valid, and returns zero for invalid geometry instead of manufacturing a one-cell output through truncated natural-number subtraction.

@[reducible, inline]

abbrev Models.CNN.poolOut (input kernel stride : ℕ) : ℕ

Output size for a pooling op along one spatial axis (no padding).

Matches the standard formula:

out = (in - k) / stride + 1.

@[reducible, inline]

abbrev Models.CNN.firstConvOutHeight (inH kH stride1 padding1 : ℕ) : ℕ

Output height after the first convolution stage.

@[reducible, inline]

abbrev Models.CNN.firstConvOutWidth (inW kW stride1 padding1 : ℕ) : ℕ

Output width after the first convolution stage.

@[reducible, inline]

abbrev Models.CNN.firstPoolOutHeight (inH kH stride1 padding1 poolKH poolStride1 : ℕ) : ℕ

Output height after the first pooling stage.

@[reducible, inline]

abbrev Models.CNN.firstPoolOutWidth (inW kW stride1 padding1 poolKW poolStride1 : ℕ) : ℕ

Output width after the first pooling stage.

@[reducible, inline]

abbrev Models.CNN.secondConvOutHeight (inH kH stride1 padding1 stride2 padding2 poolKH poolStride1 : ℕ) : ℕ

Output height after the second convolution stage (after pool1).

@[reducible, inline]

abbrev Models.CNN.secondConvOutWidth (inW kW stride1 padding1 stride2 padding2 poolKW poolStride1 : ℕ) : ℕ

Output width after the second convolution stage (after pool1).

@[reducible, inline]

abbrev Models.CNN.secondPoolOutHeight (inH kH stride1 padding1 stride2 padding2 poolKH poolStride1 poolStride2 : ℕ) : ℕ

Output height after the second pooling stage.

@[reducible, inline]

abbrev Models.CNN.secondPoolOutWidth (inW kW stride1 padding1 stride2 padding2 poolKW poolStride1 poolStride2 : ℕ) : ℕ

Output width after the second pooling stage.

@[reducible, inline]

abbrev Models.CNN.featShape (c2 inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolStride1 poolStride2 : ℕ) : Spec.Shape

Feature-map shape after the second pooling stage: (c2, H2, W2).

@[reducible, inline]

abbrev Models.CNN.featSize (c2 inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolStride1 poolStride2 : ℕ) : ℕ

Flattened feature size after the second pooling stage: c2 * H2 * W2.

def Models.cnnSpec {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {inC outC inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolstride1 poolstride2 : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} {h4 : outC ≠ 0} {h5 : poolKH ≠ 0} {h6 : poolKW ≠ 0} {hPoolStride1 : poolstride1 ≠ 0} {hPoolStride2 : poolstride2 ≠ 0} (conv1_spec : Spec.Conv2DSpec inC outC kH kW stride1 padding1 α h1 h2 h3) (conv2_spec : Spec.Conv2DSpec outC outC kH kW stride2 padding2 α h4 h2 h3) (pool1_spec : Spec.MaxPool2DSpec poolKH poolKW poolstride1 h5 h6 hPoolStride1) (pool2_spec : Spec.MaxPool2DSpec poolKH poolKW poolstride2 h5 h6 hPoolStride2) (linearSpec : Spec.LinearSpec α (CNN.featSize outC inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolstride1 poolstride2) outC) : ModSpec.SpecChain α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) (Spec.Shape.dim outC Spec.Shape.scalar)

CNN SpecChain wiring (no activations): Conv2D -> MaxPool2D -> Conv2D -> MaxPool2D -> Flatten -> Linear.

If you want the "classic" ReLU-after-conv variant, use cnn_with_relu_spec.

def Models.cnnWithReluSpec {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {inC outC inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolstride1 poolstride2 : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} {h4 : outC ≠ 0} {h5 : poolKH ≠ 0} {h6 : poolKW ≠ 0} {hPoolStride1 : poolstride1 ≠ 0} {hPoolStride2 : poolstride2 ≠ 0} (conv1_spec : Spec.Conv2DSpec inC outC kH kW stride1 padding1 α h1 h2 h3) (conv2_spec : Spec.Conv2DSpec outC outC kH kW stride2 padding2 α h4 h2 h3) (pool1_spec : Spec.MaxPool2DSpec poolKH poolKW poolstride1 h5 h6 hPoolStride1) (pool2_spec : Spec.MaxPool2DSpec poolKH poolKW poolstride2 h5 h6 hPoolStride2) (linearSpec : Spec.LinearSpec α (CNN.featSize outC inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolstride1 poolstride2) outC) : ModSpec.SpecChain α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar))) (Spec.Shape.dim outC Spec.Shape.scalar)

A slightly more "classic" CNN SpecChain with ReLU after each convolution.

PyTorch analogy: insert nn.ReLU() between the conv and pool blocks.

def Models.cnnForward {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {inC outC inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolstride1 poolstride2 : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} {h4 : outC ≠ 0} {h5 : poolKH ≠ 0} {h6 : poolKW ≠ 0} {hPoolStride1 : poolstride1 ≠ 0} {hPoolStride2 : poolstride2 ≠ 0} (conv1_spec : Spec.Conv2DSpec inC outC kH kW stride1 padding1 α h1 h2 h3) (conv2_spec : Spec.Conv2DSpec outC outC kH kW stride2 padding2 α h4 h2 h3) (pool1_spec : Spec.MaxPool2DSpec poolKH poolKW poolstride1 h5 h6 hPoolStride1) (pool2_spec : Spec.MaxPool2DSpec poolKH poolKW poolstride2 h5 h6 hPoolStride2) (linearSpec : Spec.LinearSpec α (CNN.featSize outC inH inW kH kW stride1 padding1 stride2 padding2 poolKH poolKW poolstride1 poolstride2) outC) (x : Spec.Tensor α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : Spec.Tensor α (Spec.Shape.dim outC Spec.Shape.scalar)

Run cnn_spec forward on a single input image through the assembled SpecChain.

A fully explicit CNN spec with backward pass

The SpecChain wiring above is convenient for model diagrams. For training/verification workflows we also want an explicit reverse-mode spec that returns parameter gradients.

The code below defines a small CNN in the classic:

Conv → ReLU → MaxPool → Conv → ReLU → MaxPool → Flatten → Linear

form, with a complete backward pass using the per-layer backward specs:

• conv2d_backward_spec,
• max_pool2d_multi_backward_spec,
• linear_backward_spec,
• elementwise gating for ReLU.

Configuration

The explicit CNN2Spec below is intended to be a readable, end-to-end training reference. To keep call sites ergonomic (and avoid hiding numeric architecture choices in long argument lists), we bundle the architectural hyperparameters into a config record.

structure Models.Full.CNN2Config : Type

Hyperparameters for the small 2-layer CNN (CNN2Spec).

• c1 : ℕ

  Channels after the first conv.

• c2 : ℕ

  Channels after the second conv.

• outDim : ℕ

  Output dimension of the linear head (e.g. number of classes).

• kH : ℕ

  Convolution kernel height.

• kW : ℕ

  Convolution kernel width.

• stride1 : ℕ

  Stride of the first conv.

• padding1 : ℕ

  Padding of the first conv.

• stride2 : ℕ

  Stride of the second conv.

• padding2 : ℕ

  Padding of the second conv.

• poolKH : ℕ

  Pooling kernel height.

• poolKW : ℕ

  Pooling kernel width.

• poolStride1 : ℕ

  Stride of the first pooling op.

• poolStride2 : ℕ

  Stride of the second pooling op.

structure Models.Full.CNN2Config.WF (cfg : CNN2Config) : Prop

Well-formedness conditions for CNN2Config (the nonzero facts needed by some layer specs).

• c1_ne0 : cfg.c1 ≠ 0
• c2_ne0 : cfg.c2 ≠ 0
• outDim_ne0 : cfg.outDim ≠ 0
• kH_ne0 : cfg.kH ≠ 0
• kW_ne0 : cfg.kW ≠ 0
• poolKH_ne0 : cfg.poolKH ≠ 0
• poolKW_ne0 : cfg.poolKW ≠ 0
• poolStride1_ne0 : cfg.poolStride1 ≠ 0
• poolStride2_ne0 : cfg.poolStride2 ≠ 0

def Models.Full.cnn2DefaultConfig : CNN2Config

Default CNN2Config (a small "classic" CNN shape).

theorem Models.Full.cnn2DefaultConfig_wf : cnn2DefaultConfig.WF

cnn2DefaultConfig is well-formed.

structure Models.Full.CNN2Spec (cfg : CNN2Config) (inC inH inW : ℕ) (α : Type) (h_inC : inC ≠ 0) (hCfg : cfg.WF) : Type

A small 2-block CNN with an explicit linear head.

Parameters:

• conv1 : Conv2d(inC -> c1)
• conv2 : Conv2d(c1 -> c2)
• pool1, pool2 : MaxPool2d (no padding in this spec)
• head : Linear(c2 * H2 * W2 -> outDim)

PyTorch mental model:

Conv → ReLU → MaxPool → Conv → ReLU → MaxPool → Flatten → Linear.

This structure is used by forward and backward below.

• conv1 : Spec.Conv2DSpec inC cfg.c1 cfg.kH cfg.kW cfg.stride1 cfg.padding1 α h_inC ⋯ ⋯
• conv2 : Spec.Conv2DSpec cfg.c1 cfg.c2 cfg.kH cfg.kW cfg.stride2 cfg.padding2 α ⋯ ⋯ ⋯
• pool1 : Spec.MaxPool2DSpec cfg.poolKH cfg.poolKW cfg.poolStride1 ⋯ ⋯ ⋯
• pool2 : Spec.MaxPool2DSpec cfg.poolKH cfg.poolKW cfg.poolStride2 ⋯ ⋯ ⋯
• head : Spec.LinearSpec α (CNN.featSize cfg.c2 inH inW cfg.kH cfg.kW cfg.stride1 cfg.padding1 cfg.stride2 cfg.padding2 cfg.poolKH cfg.poolKW cfg.poolStride1 cfg.poolStride2) cfg.outDim

structure Models.Full.CNN2Grads (cfg : CNN2Config) (inC inH inW : ℕ) (α : Type) : Type

Gradients for CNN2Spec parameters (returned by CNN2Spec.backward).

• d_conv1_kernel : Spec.Tensor α (Spec.Shape.dim cfg.c1 (Spec.Shape.dim inC (Spec.Shape.dim cfg.kH (Spec.Shape.dim cfg.kW Spec.Shape.scalar))))
• d_conv1_bias : Spec.Tensor α (Spec.Shape.dim cfg.c1 Spec.Shape.scalar)
• d_conv2_kernel : Spec.Tensor α (Spec.Shape.dim cfg.c2 (Spec.Shape.dim cfg.c1 (Spec.Shape.dim cfg.kH (Spec.Shape.dim cfg.kW Spec.Shape.scalar))))
• d_conv2_bias : Spec.Tensor α (Spec.Shape.dim cfg.c2 Spec.Shape.scalar)
• d_head_W : Spec.Tensor α (Spec.Shape.dim cfg.outDim (Spec.Shape.dim (CNN.featSize cfg.c2 inH inW cfg.kH cfg.kW cfg.stride1 cfg.padding1 cfg.stride2 cfg.padding2 cfg.poolKH cfg.poolKW cfg.poolStride1 cfg.poolStride2) Spec.Shape.scalar))
• d_head_b : Spec.Tensor α (Spec.Shape.dim cfg.outDim Spec.Shape.scalar)

def Models.Full.CNN2Spec.forward {α : Type} [Context α] {cfg : CNN2Config} {inC inH inW : ℕ} {h_inC : inC ≠ 0} {hCfg : cfg.WF} (m : CNN2Spec cfg inC inH inW α h_inC hCfg) (x : Spec.Tensor α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : Spec.Tensor α (Spec.Shape.dim cfg.outDim Spec.Shape.scalar)

Forward pass for CNN2Spec.

This is the "full implementation" version of the model, written directly in terms of the layer-level specs, rather than via SpecChain.

def Models.Full.CNN2Spec.backward {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {cfg : CNN2Config} {inC inH inW : ℕ} {h_inC : inC ≠ 0} {hCfg : cfg.WF} (m : CNN2Spec cfg inC inH inW α h_inC hCfg) (x : Spec.Tensor α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) (grad_output : Spec.Tensor α (Spec.Shape.dim cfg.outDim Spec.Shape.scalar)) : CNN2Grads cfg inC inH inW α × Spec.Tensor α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))

Backward pass (reverse-mode / VJP) for CNN2Spec.

Returns parameter gradients plus the input gradient.

Each step uses the corresponding layer backward spec:

• linear_backward_spec
• max_pool2d_multi_backward_spec
• conv2d_backward_spec

and ReLU uses the standard pointwise gate dY = dR ⊙ relu'(preact).