sru_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNISTでSRU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import math\n",
    "import numpy as np\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.datasets import fetch_mldata\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "from models import SRU, GRU, LSTM\n",
    "\n",
    "gpu = False\n",
    "# torch.cuda.set_device(2)\n",
    "seed = 42\n",
    "torch.manual_seed(seed)\n",
    "dir_path = './trained_models/ipynb'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. MNISTデータセット準備"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def load_mnist():\n",
    "    mnist = fetch_mldata('MNIST original')\n",
    "    mnist_X, mnist_y = shuffle(mnist.data, mnist.target, random_state=seed)\n",
    "    mnist_X = mnist_X / 255.0\n",
    "\n",
    "    # pytorch用に型変換\n",
    "    mnist_X, mnist_y = mnist_X.astype('float32'), mnist_y.astype('int64')\n",
    "\n",
    "    # 2次元の画像を、各行を互い違いにして1次元に変換\n",
    "    def flatten_img(images):\n",
    "        '''\n",
    "        images: shape => (n, rows, columns)\n",
    "        output: shape => (n, rows*columns)\n",
    "        '''\n",
    "        n_rows    = images.shape[1]\n",
    "        n_columns = images.shape[2]\n",
    "        for num in range(n_rows):\n",
    "            if num % 2 != 0:\n",
    "                images[:, num, :] = images[:, num, :][:, ::-1]\n",
    "        output = images.reshape(-1, n_rows*n_columns)\n",
    "        return output\n",
    "\n",
    "    mnist_X = mnist_X.reshape(-1, 28, 28)\n",
    "    mnist_X = flatten_img(mnist_X) # X.shape => (n_samples, seq_len)\n",
    "    mnist_X = mnist_X[:, :, np.newaxis] # X.shape => (n_samples, seq_len, n_features)\n",
    "\n",
    "    # 訓練、テストデータに分割\n",
    "    train_X, test_X, train_y, test_y = train_test_split(mnist_X, mnist_y,\n",
    "                                                        test_size=0.2,\n",
    "                                                        random_state=seed)\n",
    "    return train_X, test_X, train_y, test_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_X, test_X, train_y, test_y = load_mnist()\n",
    "# デモ用にサンプルを小さくする\n",
    "train_X, test_X, train_y, test_y = train_X[:256], test_X[:256], train_y[:256], test_y[:256]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 訓練の準備"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 計算時間を表示させる\n",
    "def timeSince(since):\n",
    "    now = time.time()\n",
    "    s = now - since\n",
    "    m = math.floor(s / 60)\n",
    "    s -= m * 60\n",
    "    return '%dm %ds' % (m, s)\n",
    "\n",
    "# batchあたりの訓練\n",
    "def train(model, inputs, labels, optimizer, criterion, clip):\n",
    "    batch_size = inputs.size(1)\n",
    "    model.initHidden(batch_size) # 隠れ変数の初期化\n",
    "    optimizer.zero_grad() # 勾配の初期化\n",
    "    outputs = model(inputs)\n",
    "    loss = criterion(outputs, labels)\n",
    "    torch.nn.utils.clip_grad_norm(model.parameters(), clip) # gradient clipping\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    acc = (torch.max(outputs, 1)[1] == labels).float().sum().data[0] / batch_size\n",
    "    return loss.data[0], acc\n",
    "\n",
    "# 検証\n",
    "def test(model, inputs, labels, criterion):\n",
    "    batch_size = inputs.size(1)\n",
    "    model.initHidden(batch_size)\n",
    "    outputs = model(inputs)\n",
    "    loss = criterion(outputs, labels)\n",
    "    acc = (torch.max(outputs, 1)[1] == labels).float().sum().data[0] / batch_size\n",
    "    return outputs, loss.data[0], acc\n",
    "\n",
    "# モデルの保存\n",
    "def checkpoint(model, optimizer, acc):\n",
    "    filename = os.path.join(dir_path, '%s_acc-%d' % (model.__class__.__name__, acc))\n",
    "    # modelの状態保存\n",
    "    torch.save(model.state_dict(), filename + '.model')\n",
    "    # optimizerの状態保存\n",
    "    torch.save(optimizer.state_dict(), filename + '.state')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. モデル、パラメータの設定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "input_size = train_X.shape[2]\n",
    "output_size = np.unique(train_y).size\n",
    "\n",
    "# パラメータの設定\n",
    "lr = 0.0005\n",
    "weight_decay = 0.0005\n",
    "dropout = 0.2\n",
    "clip = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def load_model(model_name):\n",
    "    if model_name == 'sru':\n",
    "        phi_size      = 200\n",
    "        r_size        = 60\n",
    "        cell_out_size = 200\n",
    "    elif model_name in ['gru', 'lstm']:\n",
    "        hidden_size = 200\n",
    "        num_layers  = 1\n",
    "        init_forget_bias = 1\n",
    "    \n",
    "    # モデルのインスタンス作成\n",
    "    if model_name == 'sru':\n",
    "        model = SRU(input_size, phi_size, r_size, cell_out_size, output_size, dropout=dropout, gpu=gpu)\n",
    "        model.initWeight()\n",
    "    elif model_name == 'gru':\n",
    "        model = GRU(input_size, hidden_size, output_size, num_layers, dropout, gpu=gpu)\n",
    "        model.initWeight(init_forget_bias)\n",
    "    elif model_name == 'lstm':\n",
    "        model = LSTM(input_size, hidden_size, output_size, num_layers, dropout, gpu=gpu)\n",
    "        model.initWeight(init_forget_bias)\n",
    "    if gpu == True:\n",
    "        model.cuda()\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model = load_model('sru')\n",
    "# loss, optimizerの定義\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 訓練"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:: 1, (0m 17s) train_cost: 2.328, test_cost: 2.300, train_acc: 0.086, test_acc: 0.094\n",
      "EPOCH:: 2, (0m 34s) train_cost: 2.290, test_cost: 2.297, train_acc: 0.168, test_acc: 0.117\n",
      "EPOCH:: 3, (0m 50s) train_cost: 2.276, test_cost: 2.309, train_acc: 0.160, test_acc: 0.137\n",
      "EPOCH:: 4, (1m 7s) train_cost: 2.262, test_cost: 2.320, train_acc: 0.152, test_acc: 0.129\n",
      "EPOCH:: 5, (1m 24s) train_cost: 2.237, test_cost: 2.310, train_acc: 0.195, test_acc: 0.188\n",
      "EPOCH:: 6, (1m 42s) train_cost: 2.206, test_cost: 2.337, train_acc: 0.211, test_acc: 0.195\n",
      "EPOCH:: 7, (1m 57s) train_cost: 2.202, test_cost: 2.280, train_acc: 0.188, test_acc: 0.152\n",
      "EPOCH:: 8, (2m 13s) train_cost: 2.162, test_cost: 2.245, train_acc: 0.219, test_acc: 0.176\n",
      "EPOCH:: 9, (2m 29s) train_cost: 2.312, test_cost: 2.329, train_acc: 0.219, test_acc: 0.125\n",
      "EPOCH:: 10, (2m 46s) train_cost: 2.254, test_cost: 2.264, train_acc: 0.137, test_acc: 0.082\n",
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "''' 訓練 '''\n",
    "\n",
    "n_epochs = 10\n",
    "batch_size = 64\n",
    "n_batches = train_X.shape[0]//batch_size\n",
    "n_batches_test = test_X.shape[0]//batch_size\n",
    "all_acc = []\n",
    "start_time = time.time()\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    train_cost, test_cost, train_acc, test_acc  = 0, 0, 0, 0\n",
    "    train_X, train_y = shuffle(train_X, train_y, random_state=seed)\n",
    "\n",
    "    # 訓練\n",
    "    model.train()\n",
    "    train_X_t = np.transpose(train_X, (1, 0, 2)) # X.shape => (seq_len, n_samples, n_features) に変換\n",
    "    for i in range(n_batches):\n",
    "        start = i * batch_size\n",
    "        end = start + batch_size\n",
    "        inputs, labels = train_X_t[:, start:end, :], train_y[start:end]\n",
    "        inputs, labels = Variable(torch.from_numpy(inputs)\n",
    "                         ), Variable(torch.from_numpy(labels))\n",
    "        if gpu == True:\n",
    "            inputs, labels = inputs.cuda(), labels.cuda()\n",
    "        cost, accuracy = train(model, inputs, labels, optimizer, criterion, clip)\n",
    "        train_cost += cost / n_batches\n",
    "        train_acc  += accuracy / n_batches\n",
    "\n",
    "    # 検証\n",
    "    model.eval()\n",
    "    test_X_t = np.transpose(test_X, (1, 0, 2))\n",
    "    for i in range(n_batches_test):\n",
    "        start = i * batch_size\n",
    "        end = start + batch_size\n",
    "        inputs, labels = test_X_t[:, start:end, :], test_y[start:end]\n",
    "        inputs, labels = Variable(torch.from_numpy(inputs)\n",
    "                         ), Variable(torch.from_numpy(labels))\n",
    "        if gpu == True:\n",
    "            inputs, labels = inputs.cuda(), labels.cuda()\n",
    "        _, cost, accuracy = test(model, inputs, labels, criterion)\n",
    "        test_cost += cost / n_batches_test\n",
    "        test_acc += accuracy / n_batches_test\n",
    "\n",
    "    print('EPOCH:: %i, (%s) train_cost: %.3f, test_cost: %.3f, train_acc: %.3f, test_acc: %.3f' % (epoch + 1,\n",
    "                       timeSince(start_time), train_cost, test_cost, train_acc, test_acc))\n",
    "\n",
    "#     # 過去のエポックのtest_accを上回った時だけモデルの保存\n",
    "#     if len(all_acc) == 0 or test_acc > max(all_acc):\n",
    "#         checkpoint(model, optimizer, test_acc*10000)\n",
    "#     all_acc.append(test_acc)\n",
    "\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 訓練済みモデルの読み込み"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model.load_state_dict(torch.load(dir_path + '/SRU_sample.model'))\n",
    "if gpu == True:\n",
    "    model.cuda()\n",
    "    model._gpu = True\n",
    "    model.A_mask = model.A_mask.cuda()\n",
    "    optimizer.load_state_dict(torch.load(dir_path + '/SRU_sample.state'))\n",
    "else:\n",
    "    model.cpu()\n",
    "    model._gpu = False\n",
    "    model.A_mask = model.A_mask.cpu()\n",
    "    # ToDo: optimizerをcpu, gpuともに対応させてloadできるようにする\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. テスト"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.96875\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "test_X_t = np.transpose(test_X, (1, 0, 2))\n",
    "inputs, labels = test_X_t, test_y\n",
    "inputs, labels = Variable(torch.from_numpy(inputs)\n",
    "                 ), Variable(torch.from_numpy(labels))\n",
    "if gpu == True:\n",
    "    inputs, labels = inputs.cuda(), labels.cuda()\n",
    "outputs, cost, accuracy = test(model, inputs, labels, criterion)\n",
    "print(accuracy)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf_py35",
   "language": "python",
   "name": "tf_py35"
  },
  "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}