{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quick MNIST"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*by Jeremy Howard, fast.ai.*\n",
    "\n",
    "In this notebook we'll see how easy it is to use the simple functions we created along with `torch.nn` and friends to train MNIST. We're going to create the same convolution net that we created at the end of the previous notebook, but this time we'll use the functions we've already written, and we'll skip the explanatory text so you can see just the final code.\n",
    "\n",
    "You can use this same notebook to train other neural nets on other datasets with minimal changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle, gzip\n",
    "from pathlib import Path\n",
    "from mnist_sample import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = Path('data')/'mnist'\n",
    "\n",
    "with gzip.open(PATH/'mnist.pkl.gz', 'rb') as f:\n",
    "    ((train_x, train_y), (valid_x, valid_y), _) = pickle.load(f, encoding='latin-1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs=64\n",
    "lr=0.1\n",
    "epochs=2\n",
    "\n",
    "dev = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(x,y): return x.view(-1,1,28,28).to(dev),y.to(dev)\n",
    "\n",
    "def get_dataloader(x,y,bs,shuffle):\n",
    "    ds = TensorDataset(*map(tensor, (x,y)))\n",
    "    dl = DataLoader(ds, batch_size=bs, shuffle=shuffle)\n",
    "    return WrappedDataLoader(dl, preprocess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl = get_dataloader(train_x, train_y, bs,   shuffle=False)\n",
    "valid_dl = get_dataloader(valid_x, valid_y, bs*2, shuffle=True )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    nn.Conv2d(1,  16, kernel_size=3, stride=2, padding=1), nn.ReLU(),\n",
    "    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1), nn.ReLU(),\n",
    "    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1), nn.ReLU(),\n",
    "    nn.AdaptiveAvgPool2d(1),\n",
    "    Lambda(lambda x: x.view(x.size(0),-1))\n",
    ").to(dev)\n",
    "\n",
    "opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit(epochs, model, F.cross_entropy, opt, train_dl, valid_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}