Skip to main content

mnist_lenet5

Function mnist_lenet5 

Source
pub fn mnist_lenet5<D: Float>() -> impl Fn(SymTensor) -> SymTensor
Expand description

LeNet-5 (valid-convolution variant) — no padding, no bias, returns raw logits. Use this for training with a cross-entropy loss kernel that computes log-softmax internally.

Input         [N,  1, 28, 28]
Conv2d(1→6,   5×5, pad=0) →  [N,  6, 24, 24]
ReLU
AvgPool2d(2×2, stride=2)  →  [N,  6, 12, 12]
Conv2d(6→16,  5×5, pad=0) →  [N, 16,  8,  8]
ReLU
AvgPool2d(2×2, stride=2)  →  [N, 16,  4,  4]
Flatten                   →  [N, 256]
Linear(256→120)
ReLU
Linear(120→84)
ReLU
Linear(84→10)             →  [N, 10]  raw logits