class Autoencoder(nn.Module):\n",
" \"\"\"MNIST autoencoder with named reconstruction and latent outputs.\"\"\"\n",
"\n",
" def __init__(self, latent_size=16, width=128, *, rngs):\n",
" super().__init__()\n",
" self.encoder = MLP(784, [width], latent_size, rngs=rngs)\n",
" self.decoder = MLP(latent_size, [width], 784, rngs=rngs)\n",
"\n",
" def encode(self, images):\n",
" return self.encoder(images.reshape(images.shape[0], 784))\n",
"\n",
" def decode(self, latent):\n",
" return F.sigmoid(self.decoder(latent)).reshape(-1, 1, 28, 28)\n",
"\n",
" def forward(self, images):\n",
" z = self.encode(images)\n",
" return {\"reconstruction\": self.decode(z), \"latent\": z}\n",
"@jit\n",
"def autoencoder_step(model, optimizer, images):\n",
" loss, grads = value_and_grad(lambda m: F.mse_loss(m(images)[\"reconstruction\"], images))(model)\n",
" optimizer.step(grads)\n",
" return {\"loss\": loss, \"count\": images.shape[0]}\n",
"