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 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 conv along one spatial axis.

Matches the standard PyTorch formula (for dilation = 1, groups = 1):

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

@[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.outH1 (inH kH stride1 padding1 : ℕ) : ℕ

Output height after the first convolution stage.

@[reducible, inline]

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

Output width after the first convolution stage.

@[reducible, inline]

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

Output height after the first pooling stage.

@[reducible, inline]

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

Output width after the first pooling stage.

@[reducible, inline]

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

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

@[reducible, inline]

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

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

@[reducible, inline]

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

Output height after the second pooling stage.

@[reducible, inline]

abbrev Models.CNN.poolW2 (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.MultiChannelImage inC inH inW α) : Spec.Tensor α (Spec.Shape.dim outC Spec.Shape.scalar)

Run cnn_spec forward on a single input image.

This is a convenience wrapper around SpecChain.forward.

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.

This section provides 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.MultiChannelImage inC inH inW α) : 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.MultiChannelImage inC inH inW α) (grad_output : Spec.Tensor α (Spec.Shape.dim cfg.outDim Spec.Shape.scalar)) : CNN2Grads cfg inC inH inW α × Spec.MultiChannelImage inC inH inW α

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).