examples/pytorch_experiment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## trixi PyTorch Experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to use the trixi `PytorchExperiment` with the Pytorch MNIST example to classify mnist digits."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before running call:  \n",
    "`python -m visdom.server -p 8080`  \n",
    "This starts a visdom server which is used to visualize the training's progress in real-time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "from trixi.util import Config\n",
    "from trixi.experiment import PytorchExperiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's prepare the experiment dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir experiment_dir/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls experiment_dir/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf experiment_dir/20*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!du -sh experiment_dir/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's create a Config. \n",
    "A config is basically a dict (which can be accessed with the \".\" operator).\n",
    "All objects in the dict will be initialized when the experiment starts.\n",
    "Additonally all config keywords/elements can be parsed over the command line (e.g. --batch_size=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = Config()\n",
    "\n",
    "c.batch_size = 64\n",
    "c.batch_size_test = 1000\n",
    "c.n_epochs = 10\n",
    "c.learning_rate = 0.01\n",
    "c.momentum = 0.9\n",
    "if torch.cuda.is_available():\n",
    "    c.use_cuda = True\n",
    "else:\n",
    "    c.use_cuda = False\n",
    "c.rnd_seed = 1\n",
    "c.log_interval = 200\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define the Model we use for classification "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build a simple cnn model\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        x = x.view(-1, 320)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now setup a PytorchExperiment, therefore we create a class which inherits the PytorchExperiment class.\n",
    "We than overwrite the setup, train and validate method. \n",
    "When we finally call the experiement.run() method it will call the setup, and then for the defined number of n_epochs it will call the train and validate method in an alternating fashion.\n",
    "\n",
    "The PytorchExperiment has serveral benefits:\n",
    "* It automatically creates a experimentLogger (elog) which will create a defined folder structure and can be used to store all the results\n",
    "* It automatically creates a visdomLogger (vlog) which can be used to show the results live on a visdom server\n",
    "* It automatically creates a combinedLogger (clog) which has the same interface as the experimentLogger and visdomLogger and logs to both in defined interval (e.g. you can see each result in visdom, while only saving every 10th on your hard disk)\n",
    "* After each epoch and at the end, if an error occours, it automatically stores a checkpoint of your experiment\n",
    "* You can simply resume an ended experiment by provinding its experiment folder when creating a new one\n",
    "* You can use the add_result method to compare your experiments in the trixi browser, backtrace all your result updates, and see all the results on your visdom server (at the same time)\n",
    "* Save your config and your code (if given globs=globals()) for full reproducibility)\n",
    "* Many more ;-)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNIST_experiment(PytorchExperiment):\n",
    "    def setup(self):\n",
    "        \n",
    "        self.elog.print(\"Config:\")\n",
    "        self.elog.print(self.config)\n",
    "        \n",
    "        ### Get Dataset\n",
    "        transf = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.1307,), (0.3081,))\n",
    "        ])\n",
    "        self.dataset_train = datasets.MNIST(root=\"experiment_dir/data/\", download=True,\n",
    "                                                        transform=transf, train=True)\n",
    "        self.dataset_test = datasets.MNIST(root=\"experiment_dir/data/\", download=True,\n",
    "                                                       transform=transf, train=False)\n",
    "\n",
    "        data_loader_kwargs = {'num_workers': 1, 'pin_memory': True} if self.config.use_cuda else {}\n",
    "        \n",
    "        self.train_data_loader = torch.utils.data.DataLoader(self.dataset_train, batch_size=self.config.batch_size,\n",
    "                                                        shuffle=True, **data_loader_kwargs)\n",
    "        self.test_data_loader = torch.utils.data.DataLoader(self.dataset_test, batch_size=self.config.batch_size,\n",
    "                                                       shuffle=True, **data_loader_kwargs)\n",
    "\n",
    "        \n",
    "        self.device = torch.device(\"cuda\" if self.config.use_cuda else \"cpu\")\n",
    "        \n",
    "        self.model = Net()\n",
    "        self.model.to(self.device)\n",
    "\n",
    "        self.optimizer = optim.SGD(self.model.parameters(), lr=self.config.learning_rate,\n",
    "                                               momentum=self.config.momentum)\n",
    "        \n",
    "        self.save_checkpoint(name=\"checkpoint_start\")\n",
    "        self.vlog.plot_model_structure(self.model,\n",
    "                                       [self.config.batch_size, 1, 28, 28], \n",
    "                                       name='Model Structure')\n",
    "        \n",
    "        self.batch_counter = 0        \n",
    "        self.elog.print('Experiment set up.')\n",
    "        \n",
    "    \n",
    "    def train(self, epoch):\n",
    "        \n",
    "        self.model.train()\n",
    "        \n",
    "        for batch_idx, (data, target) in enumerate(self.train_data_loader):\n",
    "            \n",
    "            self.batch_counter += 1\n",
    "            \n",
    "            if self.config.use_cuda:\n",
    "                data, target = data.cuda(), target.cuda()\n",
    "                \n",
    "            self.optimizer.zero_grad()\n",
    "            \n",
    "            output = self.model(data)\n",
    "            self.loss = F.nll_loss(output, target)\n",
    "            self.loss.backward()\n",
    "            \n",
    "            self.optimizer.step()\n",
    "            \n",
    "            if batch_idx % self.config.log_interval == 0:\n",
    "                # plot train loss (mathematically mot 100% correct, just so that lisa can sleep at night (if no one is breathing next to her ;-P) )\n",
    "                self.add_result(value=self.loss.item(), name='Train_Loss',\n",
    "                                     counter=epoch + batch_idx / len(self.train_data_loader), label='Loss')\n",
    "                # log train batch loss and progress\n",
    "                self.clog.show_text(\n",
    "                    'Train Epoch: {} [{}/{} samples ({:.0f}%)]\\t Batch Loss: {:.6f}'\n",
    "                    .format(epoch, batch_idx * len(data),\n",
    "                            len(self.train_data_loader.dataset),\n",
    "                            100. * batch_idx / len(self.train_data_loader),\n",
    "                            self.loss.item()), name=\"log\")\n",
    "                \n",
    "                self.clog.show_image_grid(data, name=\"mnist_training\", n_iter=epoch + batch_idx / len(self.train_data_loader), iter_format=\"{:0.02f}\")\n",
    "                \n",
    "                self.save_checkpoint(name=\"checkpoint\", n_iter=batch_idx)\n",
    "                \n",
    "    def validate(self, epoch):\n",
    "        self.model.eval()\n",
    "        \n",
    "        validation_loss = 0\n",
    "        correct = 0\n",
    "        \n",
    "        for data, target in self.test_data_loader:\n",
    "            if self.config.use_cuda:\n",
    "                data, target = data.cuda(), target.cuda()\n",
    "            output = self.model(data)\n",
    "            validation_loss += F.nll_loss(output, target, size_average=False).item()\n",
    "            pred = output.data.max(1, keepdim=True)[1]\n",
    "            correct += pred.eq(target.data.view_as(pred)).cpu().sum().item()\n",
    "        validation_loss /= len(self.test_data_loader.dataset)\n",
    "        # plot the test loss\n",
    "        self.add_result(value=validation_loss, name='Validation_Loss',\n",
    "                             counter=epoch + 1, label='Loss')\n",
    "        # plot the test accuracy\n",
    "        acc = 100. * correct / len(self.test_data_loader.dataset)\n",
    "        self.add_result(value=acc, name='ValidationAccurracy',\n",
    "                             counter=epoch + 1, label='Accurracy' )\n",
    "        \n",
    "        # log validation loss and accuracy\n",
    "        self.elog.print(\n",
    "            '\\nValidation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'\n",
    "            .format(validation_loss, correct, len(self.test_data_loader.dataset),\n",
    "                    100. * correct / len(self.test_data_loader.dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "exp = MNIST_experiment(config=c, name='experiment', n_epochs=c.n_epochs, \n",
    "                       seed=42, base_dir='./experiment_dir')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "exp.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "last_experiment = 'experiment_dir/' + sorted([d for d in os.listdir('experiment_dir/') if '20' in str(d)], reverse=True)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls experiment_dir/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now resume the last Experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "last_experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from trixi.experiment import PytorchExperiment\n",
    "exp_resume = MNIST_experiment(config=c, name='resume_experiment', \n",
    "                              n_epochs=5, seed=42, base_dir='./experiment_dir', \n",
    "                              resume=last_experiment, resume_save_types=('model',\n",
    "                                                                         'simple',\n",
    "                                                                         'th_vars',\n",
    "                                                                         'results'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "exp_resume.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls experiment_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also change a parameter in your experiment and simply run the same experiment again (this can of course also be done automatically)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c.lr = 0.0001\n",
    "exp2 = MNIST_experiment(config=c, name='experiment2', n_epochs=c.n_epochs, \n",
    "                       seed=42, base_dir='./experiment_dir')\n",
    "exp2.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets compare all our experiments. Therefore we simply start the trixi browser:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m trixi.browser $PWD/experiment_dir"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}