GraphSpec ResNet-18

This file defines a ResNet-18–style convolutional network using GraphSpec's general DAG IR, with BasicBlocks, projection shortcuts, shape-indexed parameters, and TorchLean compilation support.

Why This Is A DAG Model

Classic ResNet blocks are not purely sequential: the input x flows down two paths:

1. a “main” path (Conv → BN → ReLU → Conv → BN),
2. a “skip” path (identity, or a learned projection when shapes change),

and then they are added. In a chain-only representation you either:

• recompute shared values, or
• add special-case “skip” combinators that complicate the core language.

GraphSpec’s DAG IR takes a different approach: it provides a small SSA-like term language (Term + Args) that can naturally express sharing. The semantics are:

• Term.eval: pure Spec interpreter (math-first).
• Term.compile: TorchLean program compilation (executable).

ResNet-18 here is written once, then we get:

• spec-side forward semantics (for proofs / reference),
• a backend-generic TorchLean Program (for execution / training).

Model Scope

This is a “CHW, no batch” variant (C×H×W), matching the rest of the Spec/TorchLean vision layers. It is faithful to the core ResNet-18 structure:

• stem: 7×7 conv stride 2 padding 3, BN, ReLU, 3×3 maxpool stride 2 padding 1
• stages: [2,2,2,2] BasicBlocks with channel widths [64,128,256,512]
  • first block of stages 2–4 downsamples (stride 2) and uses a 1×1 projection shortcut
• head: global average pool, linear classifier

The state records exactly the metadata needed for parameter allocation:

• Conv bias is included (our Conv2D spec has it), even though many PyTorch ResNets omit it.
• BatchNorm is “train-time” BN with learnable gamma/beta (no running mean/var state).

Shapes And Type-Level Arithmetic

The main practical challenge is typing the residual adds:

• for stride=1 blocks, we need the conv output shape to be exactly CHW c h w so we can add it to the skip input x : CHW c h w.
• for stride=2 blocks, both main-path and skip-path must agree on the downsampled shape.

We solve this by defining a small family of typed primitives that cast the “raw” conv/pool output shapes into a canonical downsample formula:

down2(h) = (h - 1) / 2 + 1

This is the standard stride-2 output formula for kernels with effective receptive field 1 or 3 when you choose padding the usual ResNet way:

• 7×7 s=2 p=3 → outH = (h - 1)/2 + 1
• 3×3 s=2 p=1 → outH = (h - 1)/2 + 1
• 1×1 s=2 p=0 → outH = (h - 1)/2 + 1
• 3×3 maxpool s=2 p=1 → outH = (h - 1)/2 + 1

By enforcing this “canonical” output shape, the residual add becomes definitional/typecheckable without introducing runtime reshapes.

References / citations:

• He et al. (2016), “Deep Residual Learning for Image Recognition” (ResNet-18, BasicBlock).
• Ioffe & Szegedy (2015), “Batch Normalization…” (BN).
• Lin et al. (2013), “Network In Network” (global average pooling).

@[reducible, inline]

abbrev NN.GraphSpec.Models.ResNet18.imgShape (c h w : ℕ) : Shape

Convenient local abbreviation for channel-first image tensors (C × H × W).

Note on parameter indexing

Inside model, the environment is Γ = params ++ [x], so parameters are accessed by index.

We avoid clever macros here. In practice, being explicit is faster to debug and (importantly) more stable under refactors: each use site carries its own proof that the index points at the expected shape.

Canonical stride-2 downsample formula

@[reducible, inline]

abbrev NN.GraphSpec.Models.ResNet18.down2 (h : ℕ) : ℕ

Canonical stride-2 output-size formula used throughout this file.

We write it once and reuse it for the stem, the downsampling residual blocks, and the max-pool so that all of those paths literally agree on the same type-level height/width expression.

theorem NN.GraphSpec.Models.ResNet18.down2_pos (h : ℕ) : down2 h > 0

down2 is always positive.

Small typed primitives for ResNet typing

theorem NN.GraphSpec.Models.ResNet18.conv3_same_out_eq {h : ℕ} (hh : h > 0) : (h + 2 * 1 - 3) / 1 + 1 = h

Stride-1, padding-1 3×3 conv preserves the spatial size for h > 0.

def NN.GraphSpec.Models.ResNet18.conv3x3Same (c h w : ℕ) (h_c : c ≠ 0) (hh : h > 0) (hw : w > 0) : DAG.PrimOp [Shape.OIHW c c 3 3, Shape.Vec c, imgShape c h w] (imgShape c h w)

Typed 3×3 convolution (stride 1, padding 1) whose output is cast to the canonical imgShape c h w.

def NN.GraphSpec.Models.ResNet18.conv7x7Down (inC outC h w : ℕ) (h_inC : inC ≠ 0) (_h_outC : outC ≠ 0) : DAG.PrimOp [Shape.OIHW outC inC 7 7, Shape.Vec outC, imgShape inC h w] (imgShape outC (down2 h) (down2 w))

Typed 7×7 stem convolution (stride 2, padding 3), cast to the canonical down2 spatial sizes.

def NN.GraphSpec.Models.ResNet18.conv3x3Down (inC outC h w : ℕ) (h_inC : inC ≠ 0) (_h_outC : outC ≠ 0) : DAG.PrimOp [Shape.OIHW outC inC 3 3, Shape.Vec outC, imgShape inC h w] (imgShape outC (down2 h) (down2 w))

Typed 3×3 convolution (stride 2, padding 1), cast to the canonical down2 spatial sizes.

def NN.GraphSpec.Models.ResNet18.conv1x1Down (inC outC h w : ℕ) (h_inC : inC ≠ 0) (_h_outC : outC ≠ 0) : DAG.PrimOp [Shape.OIHW outC inC 1 1, Shape.Vec outC, imgShape inC h w] (imgShape outC (down2 h) (down2 w))

Typed 1×1 projection convolution (stride 2), cast to the canonical down2 spatial sizes.

def NN.GraphSpec.Models.ResNet18.maxpool3x3Down (c h w : ℕ) : DAG.PrimOp [imgShape c h w] (imgShape c (down2 h) (down2 w))

Typed 3×3 max-pool (stride 2, padding 1), cast to the canonical down2 spatial sizes.

Parameter layout

@[reducible, inline]

abbrev NN.GraphSpec.Models.ResNet18.convParams (outC inC kH kW : ℕ) : List Shape

Parameter layout for a convolution layer: kernel first, then bias.

@[reducible, inline]

abbrev NN.GraphSpec.Models.ResNet18.bnParams (c : ℕ) : List Shape

Parameter layout for affine batch norm: (gamma, beta).

@[reducible, inline]

abbrev NN.GraphSpec.Models.ResNet18.basicBlockParams (inC outC : ℕ) (downsample : Bool) : List Shape

Parameter ABI for a single BasicBlock.

When downsample = true, this includes the projection shortcut (1×1 conv + BN) parameters after the two main-path conv/BN pairs.

@[reducible, inline]

abbrev NN.GraphSpec.Models.ResNet18.stageParams (inC outC : ℕ) : List Shape

Parameter ABI for one ResNet-18 stage.

A stage is two BasicBlocks. The first block downsamples exactly when inC ≠ outC; the second always keeps the same number of channels.

@[reducible, inline]

abbrev NN.GraphSpec.Models.ResNet18.params (inC numClasses : ℕ) : List Shape

Full parameter ABI for the GraphSpec ResNet-18 model.

This list is deliberately written in a flat, explicit order. The model body indexes parameters by closed numerals, and the flat ABI keeps those index proofs predictable for simp and easy to audit.

@[simp]

theorem NN.GraphSpec.Models.ResNet18.params_length (inC numClasses : ℕ) : (params inC numClasses).length = 82

The ResNet-18 GraphSpec ABI contains exactly 82 parameter tensors.

This is a small but useful structural theorem: the executable wrapper, deterministic initializer, and DAG body all rely on the same flat ABI. Keeping the count as a named theorem makes accidental parameter-layout edits visible during review instead of hiding them inside a local simp.

Deterministic initialization

def NN.GraphSpec.Models.ResNet18.initConv (outC inC kH kW seedK seedB : ℕ) : Tensor.Tensor Float (Shape.OIHW outC inC kH kW) × Tensor.Tensor Float (Shape.Vec outC)

Deterministically initialize a convolution kernel (uniform) and bias (zeros).

def NN.GraphSpec.Models.ResNet18.initBN (c seedGamma seedBeta : ℕ) : Tensor.Tensor Float (Shape.Vec c) × Tensor.Tensor Float (Shape.Vec c)

Deterministically initialize BatchNorm parameters (gamma, beta) as (ones, zeros).

def NN.GraphSpec.Models.ResNet18.initLinear (inDim outDim seedW seedB : ℕ) : Tensor.Tensor Float (Shape.Mat outDim inDim) × Tensor.Tensor Float (Shape.Vec outDim)

Deterministically initialize linear weights (uniform) and bias (zeros).

Model

def NN.GraphSpec.Models.ResNet18.model (inC h w numClasses : ℕ) (h_inC : inC > 0) (_h_h : h > 0) (_h_w : w > 0) (_h_cls : numClasses > 0) : DAG.Model (params inC numClasses) [imgShape inC h w] (Shape.Vec numClasses)

GraphSpec ResNet-18 model.

This is the public entrypoint for the DAG-authored ResNet in this directory. It packages together:

• the full typed parameter ABI (params inC numClasses),
• deterministic initialization for all 82 parameter tensors,
• and a DAG body whose pure semantics can be interpreted via DAG.Model.specFwd and compiled via DAG.Model.torchProgram.

The model is channel-first (CHW) and batch-free, matching the rest of TorchLean's vision-side Spec and runtime layers.