{
"cells": [
{
"cell_type": "markdown",
"id": "00ee37b3",
"metadata": {},
"source": [
"# GAN: Fashion-MNIST Generation\n",
"\n",
"A convolutional generator and discriminator train on 55,000 Fashion-MNIST clothing images. A fixed 64-dimensional noise batch records progression across 30 epochs.\n",
"\n",
"The outputs below are saved from local GPU training. Documentation builds render these saved outputs without running cells."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "ea8bbe18",
"metadata": {
"execution": {
"iopub.execute_input": "2026-10-14T13:51:36.613234Z",
"iopub.status.busy": "2026-10-14T13:51:36.613142Z",
"iopub.status.idle": "2026-10-14T13:51:37.336404Z",
"shell.execute_reply": "2026-10-14T13:51:37.335668Z"
},
"tags": [
"hide-input"
]
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"JAXsaw 0.1.0 · JAX 0.5.3\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"[CudaDevice(id=0)]\n"
]
}
],
"source": [
"from pathlib import Path\n",
"import sys, json, inspect\n",
"ROOT = Path.cwd() if (Path.cwd() / 'examples').exists() else Path.cwd().parent\n",
"sys.path.insert(0,str(ROOT))\n",
"import jax\n",
"import numpy as np\n",
"import pandas as pd\n",
"import jaxsaw as js\n",
"from jaxsaw import models, training\n",
"from IPython.display import display, Image, Code\n",
"RUN = ROOT / 'runs' / 'gallery/gan'\n",
"print(f\"JAXsaw {js.__version__} · JAX {jax.__version__}\")\n",
"print(jax.devices())"
]
},
{
"cell_type": "markdown",
"id": "2d07f44b",
"metadata": {},
"source": [
"## Architecture\n",
"\n",
"The generator projects noise to a `7×7` feature map and upsamples twice with transposed convolutions and BatchNorm. The discriminator uses strided convolutions and emits logits. Real images are scaled to `[-1,1]`.\n",
"\n",
"The definition below is the actual package implementation used by this run."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "089c03c0",
"metadata": {
"execution": {
"iopub.execute_input": "2026-10-14T13:51:37.337678Z",
"iopub.status.busy": "2026-10-14T13:51:37.337428Z",
"iopub.status.idle": "2026-10-14T13:51:37.365395Z",
"shell.execute_reply": "2026-10-14T13:51:37.365030Z"
},
"tags": [
"hide-input"
]
},
"outputs": [
{
"data": {
"text/html": [
"class FashionGenerator(nn.Module):\n",
" \"\"\"Convolutional Fashion-MNIST generator, with explicit BatchNorm state.\"\"\"\n",
"\n",
" def __init__(self, latent_size=64, width=64, *, rngs):\n",
" super().__init__()\n",
" self.width = width\n",
" self.project = nn.Linear(latent_size, 2 * width * 7 * 7, rngs=rngs)\n",
" self.norm1 = nn.BatchNorm2d(2 * width)\n",
" self.up1 = nn.ConvTranspose2d(2 * width, width, 4, stride=2, padding=1, rngs=rngs)\n",
" self.norm2 = nn.BatchNorm2d(width)\n",
" self.up2 = nn.ConvTranspose2d(width, 1, 4, stride=2, padding=1, rngs=rngs)\n",
"\n",
" def forward(self, z):\n",
" x = self.project(z).reshape(-1, 2 * self.width, 7, 7)\n",
" x = F.relu(self.norm1(x))\n",
" x = F.relu(self.norm2(self.up1(x)))\n",
" return jnp.tanh(self.up2(x))\n",
"class FashionDiscriminator(nn.Module):\n",
" \"\"\"Convolutional discriminator; outputs logits without a sigmoid.\"\"\"\n",
"\n",
" def __init__(self, width=32, *, rngs):\n",
" super().__init__()\n",
" self.conv1 = nn.Conv2d(1, width, 4, stride=2, padding=1, rngs=rngs)\n",
" self.conv2 = nn.Conv2d(width, 2 * width, 4, stride=2, padding=1, rngs=rngs)\n",
" self.head = nn.Linear(2 * width * 7 * 7, 1, rngs=rngs)\n",
"\n",
" def forward(self, images):\n",
" x = self.conv1(images)\n",
" x = jnp.where(x >= 0, x, 0.2 * x)\n",
" x = self.conv2(x)\n",
" x = jnp.where(x >= 0, x, 0.2 * x)\n",
" return self.head(x.reshape(x.shape[0], -1))\n",
"def gan_step(generator, discriminator, g_optimizer, d_optimizer, real, rngs, latent_size=64):\n",
" \"\"\"One D update then one G update; inactive model buffers stay in eval mode.\"\"\"\n",
" generator.eval()\n",
" discriminator.train()\n",
" z = random.normal(rngs.next(\"sample\"), (real.shape[0], latent_size))\n",
" result = discriminator_step(generator, discriminator, d_optimizer, real, z)\n",
" generator.train()\n",
" discriminator.eval()\n",
" z = random.normal(rngs.next(\"sample\"), (real.shape[0], latent_size))\n",
" return {**result, **generator_step(generator, discriminator, g_optimizer, z)}\n",
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" epochs | \n",
" batch_size | \n",
" seed | \n",
" training_examples | \n",
" validation_examples | \n",
" test_examples | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" 30 | \n",
" 128 | \n",
" 0 | \n",
" 55000 | \n",
" 1000 | \n",
" 1000 | \n",
"
\n",
" \n",
"
\n",
"