ProjectedGradient.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 学習型射影勾配法\n",
    "\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wadayama/MIKA2019/blob/master/ProjectedGradient.ipynb)\n",
    "\n",
    "学習型射影勾配法のサンプルコードです。ステップサイズを固定した方式と比較しています。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 必要なパッケージのインポート"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import math\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## グローバル定数の設定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 500 # 信号長\n",
    "mbs = 100 # ミニバッチサイズ\n",
    "noise_std = 0.5 # 雑音の標準偏差\n",
    "alpha = 8.0 # 射影関数のソフトさを表す\n",
    "A = torch.normal(mean=torch.zeros(n,n), std = 1.0)\n",
    "max_itr = 10 # 射影勾配法の反復回数\n",
    "adam_lr = 0.00005 # Adamの学習率"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ミニバッチ生成関数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_minibatch():\n",
    "    x = 1.0 - 2.0 * torch.randint(0, 2, (mbs, n)).float()\n",
    "    w = torch.normal(mean=torch.zeros(mbs, n), std = noise_std)\n",
    "    y = torch.mm(x, A.t()) + w\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 学習型射影勾配法 TPG クラス (Trainable Projected Gradient)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TPG(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TPG, self).__init__()\n",
    "        self.beta = nn.Parameter(0.01*torch.ones(max_itr)) # 学習可能ステップサイズパラメータ\n",
    "    def forward(self, y, num_itr):\n",
    "        s = torch.zeros(mbs, n) # 初期探索点\n",
    "        for i in range(num_itr):\n",
    "            r = s + self.beta[i] * torch.mm(y - torch.mm(s, A.t()), A) # 勾配ステップ\n",
    "            s = torch.tanh(alpha * r) # 射影ステップ\n",
    "        return s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 訓練ループ(インクリメンタルトレーニング）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = TPG()\n",
    "opt   = optim.Adam(model.parameters(), lr=adam_lr)\n",
    "loss_func = nn.MSELoss()\n",
    "for gen in range(max_itr):\n",
    "    for i in range(100):\n",
    "        x, y = gen_minibatch()\n",
    "        opt.zero_grad()\n",
    "        x_hat = model(y, gen + 1)\n",
    "        loss  = loss_func(x_hat, x)\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "    print(gen, loss.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 学習しない射影勾配法"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PG(nn.Module): # 普通の射影勾配法\n",
    "    def __init__(self):\n",
    "        super(PG, self).__init__()\n",
    "    def forward(self, y, num_itr, gamma):\n",
    "        s = torch.zeros(mbs, n) # 初期探索点\n",
    "        for i in range(num_itr):\n",
    "            r = s + gamma * torch.mm(y - torch.mm(s, A.t()), A)  # 勾配ステップ　\n",
    "            s = torch.tanh(alpha * r) # 射影ステップ\n",
    "        return s\n",
    "pg_model = PG()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 誤差のプロット"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mbs = 1\n",
    "x, y = gen_minibatch()\n",
    "\n",
    "norm_list = []\n",
    "itr_list = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for i in range(max_itr):\n",
    "        err = (torch.norm(x - model(y, i))**2).item()/n\n",
    "        norm_list.append(10.0*math.log10(err))\n",
    "        itr_list.append(i)\n",
    "plt.plot(itr_list, norm_list, color=\"red\", label=\"Trainable-PG\")\n",
    "\n",
    "norm_list = []\n",
    "itr_list = []\n",
    "\n",
    "for i in range(max_itr):\n",
    "    err = (torch.norm(x - pg_model(y, i, 0.0010))**2).item()/n\n",
    "    norm_list.append(10.0*math.log10(err))\n",
    "    itr_list.append(i)\n",
    "plt.plot(itr_list, norm_list, color=\"blue\", label=\"PG-1\")\n",
    "\n",
    "norm_list = []\n",
    "itr_list = []\n",
    "for i in range(max_itr):\n",
    "    err = (torch.norm(x - pg_model(y, i, 0.0015))**2).item()/n\n",
    "    norm_list.append(10.*math.log10(err))\n",
    "    itr_list.append(i)\n",
    "plt.plot(itr_list, norm_list, color=\"black\", label=\"PG-2\")\n",
    "\n",
    "norm_list = []\n",
    "itr_list = []\n",
    "for i in range(max_itr):\n",
    "    err = (torch.norm(x - pg_model(y, i, 0.0005))**2).item()/n\n",
    "    norm_list.append(10.0*math.log10(err))\n",
    "    itr_list.append(i)\n",
    "plt.plot(itr_list, norm_list, color=\"green\", label=\"PG-3\")\n",
    "\n",
    "plt.grid()\n",
    "plt.xlabel(\"iteration\")\n",
    "plt.ylabel(\"squared error (dB)\")\n",
    "plt.legend()\n",
    "plt.ylim([-30, 5])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('beta = ', model.beta)"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}