02-m2model_by_lossclass.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# M2モデルによる半教師あり学習\n",
    "では，Model APIに用意されたクラスでは表現することの難しい，より複雑なモデルはどう実装すればよいでしょうか．\n",
    "\n",
    "Pixyzでは，複雑なモデルに対応するためにLoss APIが用意されています．\n",
    "\n",
    "出典: https://github.com/masa-su/pixyz/blob/master/examples/m2.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import torch\n",
    "import torch.utils.data\n",
    "from torch import nn, optim\n",
    "from torch.nn import functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import save_image\n",
    "from tensorboardX import SummaryWriter\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "batch_size = 128\n",
    "epochs = 10\n",
    "seed = 1\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "else:\n",
    "    device = \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/wohlert/semi-supervised-pytorch/blob/master/examples/notebooks/datautils.py\n",
    "\n",
    "from functools import reduce\n",
    "from operator import __or__\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from torchvision.datasets import MNIST\n",
    "import numpy as np\n",
    "from itertools import cycle\n",
    "\n",
    "labels_per_class = 10\n",
    "n_labels = 10\n",
    "\n",
    "root = 'data'\n",
    "transform = transforms.Compose([transforms.ToTensor(),\n",
    "                                transforms.Lambda(lambd=lambda x: x.view(-1))])\n",
    "\n",
    "mnist_train = MNIST(root=root, train=True, download=True, transform=transform)\n",
    "mnist_valid = MNIST(root=root, train=False, transform=transform)\n",
    "\n",
    "def get_sampler(labels, n=None):\n",
    "    # Only choose digits in n_labels\n",
    "    (indices,) = np.where(reduce(__or__, [labels == i for i in np.arange(n_labels)]))\n",
    "\n",
    "    # Ensure uniform distribution of labels\n",
    "    np.random.shuffle(indices)\n",
    "    indices = np.hstack([list(filter(lambda idx: labels[idx] == i, indices))[:n] for i in range(n_labels)])\n",
    "\n",
    "    indices = torch.from_numpy(indices)\n",
    "    sampler = SubsetRandomSampler(indices)\n",
    "    return sampler\n",
    "\n",
    "# Dataloaders for MNIST\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True}\n",
    "labelled = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size,\n",
    "                                       sampler=get_sampler(mnist_train.train_labels.numpy(), labels_per_class),\n",
    "                                       **kwargs)\n",
    "unlabelled = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size,\n",
    "                                         sampler=get_sampler(mnist_train.train_labels.numpy()), **kwargs)\n",
    "validation = torch.utils.data.DataLoader(mnist_valid, batch_size=batch_size,\n",
    "                                         sampler=get_sampler(mnist_valid.test_labels.numpy()), **kwargs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0) Pixyzのimport\n",
    "今回は，複雑なモデルの代表例としてm2モデルをLoss APIを使って実装します．\n",
    "\n",
    "m2モデルのlossは\n",
    "$$\n",
    "-\\sum _ { x , y \\sim p _ { d a t a } ( x , y ) } \\left[ E _ { q ( z | x , y ) } \\left[ \\log \\frac { p ( x , z | y ) } { q ( z | x , y ) } \\right] + \\alpha \\log q ( y | x ) \\right] - \\sum _ { x _ { u } \\sim p _ { d a t a } \\left( x _ { u } \\right) } \\left[ E _ { q ( z | x _ { u } , y ) q ( y | x _ { u } ) } \\left[ \\log \\frac { p \\left( x _ { u } , z | y \\right) } { q ( z | x _ { u } , y ) q ( y | x _ { u } ) } \\right] \\right]\n",
    "$$\n",
    "\n",
    "で表され，第1項は，ラベル$y$が入手できる場合のloss，第2項はラベルが入手できない場合のlossになっています．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "事前分布とエンコーダには正規分布，デコーダにはベルヌーイ分布を利用します．\n",
    "- 事前分布は平均が0，標準偏差が1の標準正規分布を用います(なのでlocとscaleが定数)\n",
    "\n",
    "さらに，m2モデルでは，カテゴリ分布を用いて，識別器 $p(y|x)$ も訓練します．\n",
    "- 今回は，`RelaxedCategorical`を用いることにします．\n",
    "\n",
    "VAEのときと同様に，分布の中身を書いていきましょう．各分布のforwardの返り値が，分布のパラメータをdictにしたものになるようにします．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pixyz.distributions import Normal, Bernoulli, RelaxedCategorical\n",
    "from pixyz.models import Model\n",
    "from pixyz.losses import ELBO, NLL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_dim = 784\n",
    "y_dim = 10\n",
    "z_dim = 64\n",
    "\n",
    "\n",
    "# inference model q(z|x,y)\n",
    "class Inference(Normal):\n",
    "    def __init__(self):\n",
    "        super().__init__(cond_var=[\"x\",\"y\"], var=[\"z\"], name=\"q\")\n",
    "\n",
    "        self.fc1 = nn.Linear(x_dim+y_dim, 512)\n",
    "        self.fc21 = nn.Linear(512, z_dim)\n",
    "        self.fc22 = nn.Linear(512, z_dim)\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        h = F.relu(self.fc1(torch.cat([x, y], 1)))\n",
    "        return {\"loc\": self.fc21(h), \"scale\": F.softplus(self.fc22(h))}\n",
    "\n",
    "    \n",
    "# generative model p(x|z,y)    \n",
    "class Generator(Bernoulli):\n",
    "    def __init__(self):\n",
    "        super().__init__(cond_var=[\"z\",\"y\"], var=[\"x\"], name=\"p\")\n",
    "\n",
    "        self.fc1 = nn.Linear(z_dim+y_dim, 512)\n",
    "        self.fc2 = nn.Linear(512, x_dim)\n",
    "\n",
    "    def forward(self, z, y):\n",
    "        h = F.relu(self.fc1(torch.cat([z, y], 1)))\n",
    "        return {\"probs\": torch.sigmoid(self.fc2(h))}\n",
    "\n",
    "# classifier p(y|x)\n",
    "class Classifier(RelaxedCategorical):\n",
    "    def __init__(self):\n",
    "        super(Classifier, self).__init__(cond_var=[\"x\"], var=[\"y\"], temperature=0.1, name=\"p\")\n",
    "        self.fc1 = nn.Linear(x_dim, 512)\n",
    "        self.fc2 = nn.Linear(512, y_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = F.relu(self.fc1(x))\n",
    "        h = F.softmax(self.fc2(h), dim=1)\n",
    "        return {\"probs\": h}\n",
    "    \n",
    "# prior model p(z)\n",
    "loc = torch.tensor(0.).to(device)\n",
    "scale = torch.tensor(1.).to(device)\n",
    "prior = Normal(loc=loc, scale=scale, var=[\"z\"], dim=z_dim, name=\"prior\")    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "分布のインスタンスを作って，指定のdeviceに載せます"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# distributions for supervised learning\n",
    "p = Generator().to(device)\n",
    "q = Inference().to(device)\n",
    "f = Classifier().to(device)\n",
    "p_joint = p * prior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pixyzでは，分布をprintすることで，その分布の中身のネットワークがどうなっているかを確認できるのでしたね．\n",
    "\n",
    "実際にあっているか確認してみましょう．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distribution:\n",
      "  p(x,z|y) = p(x|z,y)prior(z)\n",
      "Network architecture:\n",
      "  prior(z) (Normal): Normal()\n",
      "  p(x|z,y) (Bernoulli): Generator(\n",
      "    (fc1): Linear(in_features=74, out_features=512, bias=True)\n",
      "    (fc2): Linear(in_features=512, out_features=784, bias=True)\n",
      "  )\n",
      "Distribution:\n",
      "  q(z|x,y) (Normal)\n",
      "Network architecture:\n",
      "  Inference(\n",
      "    (fc1): Linear(in_features=794, out_features=512, bias=True)\n",
      "    (fc21): Linear(in_features=512, out_features=64, bias=True)\n",
      "    (fc22): Linear(in_features=512, out_features=64, bias=True)\n",
      "  )\n",
      "Distribution:\n",
      "  p(y|x) (RelaxedCategorical)\n",
      "Network architecture:\n",
      "  Classifier(\n",
      "    (fc1): Linear(in_features=784, out_features=512, bias=True)\n",
      "    (fc2): Linear(in_features=512, out_features=10, bias=True)\n",
      "  )\n"
     ]
    }
   ],
   "source": [
    "print(p_joint)\n",
    "print(q)\n",
    "print(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "意図した通りになっていそうですね．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "次に，ラベルがない場合のときのLossを書いていきましょう．\n",
    "\n",
    "Pixyzでは，同じネットワークのパラメータを共有する分布を `replace_var`メソッドを使って作ることができます．\n",
    "- この時は，以下のように，元の確率変数を新たな確率変数で置き換えるように書きます．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# distributions for unsupervised learning\n",
    "q_u = q.replace_var(x=\"x_u\", y=\"y_u\")\n",
    "p_u = p.replace_var(x=\"x_u\", y=\"y_u\")\n",
    "f_u = f.replace_var(x=\"x_u\", y=\"y_u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ラベルのない場合には，$x_u$が与えられたときの$y$と$z$の同時分布$q(z,y_u|x_u)$と，$y_u$が与えられたときの$x_u$と$z$の同時分布$p(x_u,z|y_u)$が必要ですので，これらの同時分布を作りましょう．\n",
    "\n",
    "Pixyzでは同時分布は分布同士の掛け算で表現できるのでした．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_u = q_u * f_u\n",
    "p_joint_u = p_u * prior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "確認してみましょう"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distribution:\n",
      "  p(x_u,z|y_u) = p(x_u|z,y_u)prior(z)\n",
      "Network architecture:\n",
      "  prior(z) (Normal): Normal()\n",
      "  p(x_u|z,y_u) (Bernoulli): Generator(\n",
      "    (fc1): Linear(in_features=74, out_features=512, bias=True)\n",
      "    (fc2): Linear(in_features=512, out_features=784, bias=True)\n",
      "  )\n",
      "Distribution:\n",
      "  p(z,y_u|x_u) = q(z|x_u,y_u)p(y_u|x_u)\n",
      "Network architecture:\n",
      "  p(y_u|x_u) (RelaxedCategorical): Classifier(\n",
      "    (fc1): Linear(in_features=784, out_features=512, bias=True)\n",
      "    (fc2): Linear(in_features=512, out_features=10, bias=True)\n",
      "  )\n",
      "  q(z|x_u,y_u) (Normal): Inference(\n",
      "    (fc1): Linear(in_features=794, out_features=512, bias=True)\n",
      "    (fc21): Linear(in_features=512, out_features=64, bias=True)\n",
      "    (fc22): Linear(in_features=512, out_features=64, bias=True)\n",
      "  )\n",
      "Distribution:\n",
      "  p(y_u|x_u) (RelaxedCategorical)\n",
      "Network architecture:\n",
      "  Classifier(\n",
      "    (fc1): Linear(in_features=784, out_features=512, bias=True)\n",
      "    (fc2): Linear(in_features=512, out_features=10, bias=True)\n",
      "  )\n"
     ]
    }
   ],
   "source": [
    "print(p_joint_u)\n",
    "print(q_u)\n",
    "print(f_u)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "意図した通り， \n",
    "$$p(x_u,z|y_u) = p(x_u|z,y_u)prior(z)$$\n",
    "$$ p(z,y_u|x_u) = q(z|x_u,y_u)p(y_u|x_u)$$\n",
    "となっており，2つの分布の積として因数分解される同時分布になっていることが確認できました．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2) 目的関数の設定\n",
    "m2モデルのlossは\n",
    "$$\n",
    "-\\sum _ { x , y \\sim p _ { d a t a } ( x , y ) } \\left[ E _ { q ( z | x , y ) } \\left[ \\log \\frac { p ( x , z | y ) } { q ( z | x , y ) } \\right] + \\alpha \\log q ( y | x ) \\right] - \\sum _ { x _ { u } \\sim p _ { d a t a } \\left( x _ { u } \\right) } \\left[ E _ { q ( z | x _ { u } , y ) q ( y | x _ { u } ) } \\left[ \\log \\frac { p \\left( x _ { u } , z | y \\right) } { q ( z | x _ { u } , y ) q ( y | x _ { u } ) } \\right] \\right]\n",
    "$$\n",
    "で表されるのでした．\n",
    "\n",
    "それではこのlossを表現していきましょう．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "lossの式を見ると，$\n",
    "\\sum _ { x , y \\sim p _ { d a t a } ( x , y ) } \\left[ E _ { q ( z | x , y ) } \\left[ \\log \\frac { p ( x , z | y ) } { q ( z | x , y ) } \\right]  \\right]$と，\n",
    "$\\sum _ { x _ { u } \\sim p _ { d a t a } \\left( x _ { u } \\right) } \\left[ E _ { q ( z | x _ { u } , y ) q ( y | x _ { u } ) } \\left[ \\log \\frac { p \\left( x _ { u } , z | y \\right) } { q ( z | x _ { u } , y ) q ( y | x _ { u } ) } \\right] \\right]$の部分は，それぞれラベルがあるとき，ないときのELBOになっていることがわかります．\n",
    "\n",
    "そのため，この2つには，`ELBO` Lossを使ってあげればよいでしょう．\n",
    "- ELBO Lossの詳細は， https://docs.pixyz.io/en/latest/losses.html#lower-bound にあります．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "elbo_u = ELBO(p_joint_u, q_u)\n",
    "elbo = ELBO(p_joint, q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "残りの項で先頭のマイナスを，シグマの内側に入れたもの$\\sum _ { x , y \\sim p _ { d a t a } ( x , y ) } \\left[ -\\alpha \\log q ( y | x ) \\right]$は，$q ( y | x ) \\$の負の対数尤度(NLL)となっています．\n",
    "そのためこの項には`NLL` Lossを使ってあげればよいでしょう．\n",
    "- NLL Lossの詳細は， https://docs.pixyz.io/en/latest/losses.html#nll にあります．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "nll = NLL(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "全体のLossは以上の3つのLoss(`elbo_u`，`elbo`,`nll`)を組み合わせたものになっています．\n",
    "\n",
    "これらの四則演算によって，最終的なLossを作り出しましょう．\n",
    "\n",
    "【ポイント！】 PixyzではLoss同士の四則演算ができます．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "rate = 1 * (len(unlabelled) + len(labelled)) / len(labelled)\n",
    "\n",
    "loss_cls = -elbo_u.mean() -elbo.mean() + (rate * nll).mean() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "それでは，このモデルを学習する前に，Lossが本当に欲しいものかどうかチェックしてみましょう．\n",
    "\n",
    "【ポイント！】 Lossをprintするとそのロスの式を確認できます．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-(mean(E_p(z,y_u|x_u)[log p(x_u,z|y_u)/p(z,y_u|x_u)])) - mean(E_q(z|x,y)[log p(x,z|y)/q(z|x,y)]) + mean(-log p(y|x) * 470.0)\n"
     ]
    }
   ],
   "source": [
    "print(loss_cls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "意図したものになっていますね．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lossクラスを使って，自分でLossを定義した時は，ModelクラスのModelに必要なものを渡します．\n",
    "- 以下のように，学習時のloss，テスト時のloss，モデルを構成する分布を渡します.\n",
    "- モデルクラスの詳細は， https://docs.pixyz.io/en/latest/models.html#model にあります，"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(loss_cls,test_loss=nll.mean(),\n",
    "              distributions=[p, q, f], optimizer=optim.Adam, optimizer_params={\"lr\":1e-3})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "構築したモデルの中身をチェックしてみます．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distributions (for training): \n",
      "  p(x|z,y), q(z|x,y), p(y|x) \n",
      "Loss function: \n",
      "  -(mean(E_p(z,y_u|x_u)[log p(x_u,z|y_u)/p(z,y_u|x_u)])) - mean(E_q(z|x,y)[log p(x,z|y)/q(z|x,y)]) + mean(-log p(y|x) * 470.0) \n",
      "Optimizer: \n",
      "  Adam (\n",
      "  Parameter Group 0\n",
      "      amsgrad: False\n",
      "      betas: (0.9, 0.999)\n",
      "      eps: 1e-08\n",
      "      lr: 0.001\n",
      "      weight_decay: 0\n",
      "  )\n"
     ]
    }
   ],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "大丈夫そうですね．\n",
    "\n",
    "では，学習を始めましょう．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 学習してみる\n",
    "それでは実際に学習してみましょう．\n",
    "\n",
    "モデルクラスには，`train()`メソッドが用意されており，その引数に入力を与えるだけで，内部でlossを`backward()`して,パラメータを更新(`optimizer.step()`)してくれます．\n",
    "\n",
    "`test()`メソッドでは，引数として与えられた入力に対するLossを返します．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    train_loss = 0\n",
    "    for (x, y), (x_u, y_u) in tqdm(zip(cycle(labelled), unlabelled), total=len(unlabelled)):\n",
    "        x = x.to(device)\n",
    "        y = torch.eye(10)[y].to(device)\n",
    "        x_u = x_u.to(device)        \n",
    "        loss = model.train({\"x\": x, \"y\": y, \"x_u\": x_u})\n",
    "        train_loss += loss\n",
    "        \n",
    "    train_loss = train_loss * unlabelled.batch_size / len(unlabelled.dataset)\n",
    "    print('Epoch: {} Train loss: {:.4f}'.format(epoch, train_loss))\n",
    "    \n",
    "    return train_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(epoch):\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    total = 0    \n",
    "    for x, y in validation:\n",
    "        x = x.to(device)\n",
    "        y = torch.eye(10)[y].to(device)        \n",
    "        loss = model.test({\"x\": x, \"y\": y})\n",
    "        test_loss += loss\n",
    "        \n",
    "        pred_y = f.sample_mean({\"x\": x})\n",
    "        total += y.size(0)\n",
    "        correct += (pred_y.argmax(dim=1) == y.argmax(dim=1)).sum().item()      \n",
    "\n",
    "    test_loss = test_loss * validation.batch_size / len(validation.dataset)\n",
    "    test_accuracy = 100 * correct / total\n",
    "    print('Test loss: {:.4f}, Test accuracy: {:.4f}'.format(test_loss, test_accuracy))\n",
    "    return test_loss, test_accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:47<00:00, 10.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 Train loss: 328.9215\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 1.0270, Test accuracy: 76.5000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:47<00:00, 11.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 2 Train loss: 213.4401\n",
      "Test loss: 1.4781, Test accuracy: 74.7000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:48<00:00,  9.68it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 3 Train loss: 197.3201\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 1.5356, Test accuracy: 75.7700\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:47<00:00,  9.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 4 Train loss: 190.3922\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 1.6502, Test accuracy: 76.7500\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:49<00:00,  9.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 5 Train loss: 186.4787\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 1.8583, Test accuracy: 76.9300\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:47<00:00,  9.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 6 Train loss: 183.9057\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 1.9586, Test accuracy: 76.6900\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:49<00:00, 11.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 7 Train loss: 181.9903\n",
      "Test loss: 2.1133, Test accuracy: 76.7800\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:47<00:00,  9.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 8 Train loss: 180.5518\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 2.4012, Test accuracy: 76.1500\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:49<00:00,  9.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 9 Train loss: 179.3192\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 2.1873, Test accuracy: 77.9000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 469/469 [00:48<00:00,  9.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 10 Train loss: 178.3697\n",
      "Test loss: 2.7381, Test accuracy: 75.2500\n"
     ]
    }
   ],
   "source": [
    "writer = SummaryWriter()\n",
    "\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train_loss = train(epoch)\n",
    "    test_loss, test_accuracy = test(epoch)\n",
    "\n",
    "    writer.add_scalar('train_loss', train_loss.item(), epoch)\n",
    "    writer.add_scalar('test_loss', test_loss.item(), epoch)\n",
    "    writer.add_scalar('test_accuracy', test_accuracy, epoch)    \n",
    "    \n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}