class DigitCNN(nn.Module):\n",
" \"\"\"Four convolution/pooling stages for MNIST NCHW images.\n",
"\n",
" Zero padding expands 28x28 to 32x32 before four reductions to 16, 8,\n",
" 4, and 2 pixels per side. Earlier two-stage checkpoints are incompatible.\n",
" \"\"\"\n",
"\n",
" def __init__(self, width=16, classes=10, *, rngs):\n",
" super().__init__()\n",
" self.conv1 = nn.Conv2d(1, width, 3, padding=1, rngs=rngs)\n",
" self.conv2 = nn.Conv2d(width, width * 2, 3, padding=1, rngs=rngs)\n",
" self.conv3 = nn.Conv2d(width * 2, width * 4, 3, padding=1, rngs=rngs)\n",
" self.conv4 = nn.Conv2d(width * 4, width * 8, 3, padding=1, rngs=rngs)\n",
" self.pool = nn.MaxPool2d(2, stride=2)\n",
" self.dropout = nn.Dropout(0.2)\n",
" self.head = nn.Linear(width * 8 * 2 * 2, classes, rngs=rngs)\n",
"\n",
" def forward(self, x, *, rngs=None):\n",
" x = jnp.pad(x, ((0, 0), (0, 0), (2, 2), (2, 2)))\n",
" for conv in (self.conv1, self.conv2, self.conv3, self.conv4):\n",
" x = self.pool(F.relu(conv(x)))\n",
" return self.head(self.dropout(x.reshape(x.shape[0], -1), rngs=rngs))\n",
"@jit\n",
"def cnn_step(model, optimizer, batch, rngs):\n",
" x, y = batch\n",
"\n",
" def objective(candidate):\n",
" logits = candidate(x, rngs=rngs)\n",
" return F.cross_entropy(logits, y), (logits.argmax(-1) == y).sum()\n",
"\n",
" (loss, correct), grads = value_and_grad(objective, has_aux=True)(model)\n",
" optimizer.step(grads)\n",
" return {\"loss\": loss, \"correct\": correct, \"count\": y.size}\n",
"