Mnist-Policies.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparation Instructions\n",
    "* Download mnist8m from \n",
    "https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/multiclass/mnist8m.bz2\n",
    "* Decompress the first 2 million samples and split them into two chunks with 1 million samples each.\n",
    "```\n",
    "bzip2 -dc mnist8m.bz2 | head -n 2000000 | split -l 1000000 -\n",
    "```\n",
    "* Execute all the cells below to produce policies h, pi, and reward predictor q and save the relevant parameters to an npz file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%reload_ext line_profiler\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "from joblib import Memory\n",
    "from sklearn.datasets import load_svmlight_file\n",
    "from sklearn.kernel_approximation import RBFSampler\n",
    "from scipy.spatial.distance import pdist\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "________________________________________________________________________________\n",
      "[Memory] Calling __main__redacted...\n",
      "get_data()\n",
      "_______________________________________________________get_data - 164.7s, 2.7min\n"
     ]
    }
   ],
   "source": [
    "mem = Memory(\"./mycache\")\n",
    "\n",
    "@mem.cache\n",
    "def get_data():\n",
    "    data = load_svmlight_file(\"xaa\")\n",
    "    return data[0], data[1]\n",
    "\n",
    "X, y = get_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = X.todense().A.astype(np.float32)\n",
    "del X\n",
    "x /= 255.\n",
    "xx = x.T@x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "d,v = np.linalg.eigh(xx)\n",
    "perm = np.argsort(-d)\n",
    "vv = v[:,perm[:50]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "xv = x@vv\n",
    "del x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.012065845514763526\n"
     ]
    }
   ],
   "source": [
    "class RFF():\n",
    "    def __init__(self, x, gamma=None, n_components=None, seed=90210):\n",
    "        n, d = x.shape\n",
    "        self.d = d\n",
    "        self.c = n_components if n_components is not None else int(np.round(np.sqrt(n)))\n",
    "        self.seed = seed\n",
    "        self.rs = np.random.RandomState(self.seed)\n",
    "        if gamma is None:\n",
    "            n_samples = min(n, 1000)\n",
    "            perm = self.rs.permutation(n)[:n_samples]\n",
    "            gamma = np.median(pdist(x[perm]).ravel())**(-2) \n",
    "        print(gamma)\n",
    "        self.gamma = gamma\n",
    "        self.omega = self.rs.normal(size=(d, self.c)).astype(np.float32)\n",
    "        self.omega *= np.sqrt(2 * self.gamma)\n",
    "        self.b = self.rs.uniform(0, 2*np.pi, size=self.c).astype(np.float32)\n",
    "        self.normalizer = np.sqrt(2./self.c)\n",
    "    \n",
    "    def __call__(self,x):\n",
    "        projection = np.dot(x, self.omega) + self.b\n",
    "        np.cos(projection, projection)\n",
    "        projection *= self.normalizer\n",
    "        return projection\n",
    "\n",
    "rff = RFF(xv)\n",
    "rffxv = rff(xv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import OneHotEncoder\n",
    "ohe = OneHotEncoder(sparse=False).fit(y.reshape(-1,1))\n",
    "y = ohe.transform(y.reshape(-1,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "rxx=np.dot(rffxv.T,rffxv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.linalg import cho_factor, cho_solve\n",
    "rcf = cho_factor(rxx+1e-12*np.eye(rxx.shape[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def agd(x,y,c=None):\n",
    "    nx, d = x.shape\n",
    "    ny, k = y.shape \n",
    "    assert nx == ny\n",
    "    n = nx\n",
    "    from scipy.linalg import cho_factor, cho_solve\n",
    "    if c is None:\n",
    "        xx=np.dot(x.T,x)\n",
    "        c = cho_factor(xx+1e-12*np.eye(xx.shape[0]))\n",
    "    li, nextli = (1,1)\n",
    "    u = np.zeros((d,k), dtype=np.float32)\n",
    "    w = np.zeros((d,k), dtype=np.float32)\n",
    "    t = np.zeros((d,k), dtype=np.float32)\n",
    "    p = np.zeros((n,k), dtype=np.float32)\n",
    "    z = np.zeros(n, dtype=np.float32)\n",
    "    g = np.zeros((d,k), dtype=np.float32)\n",
    "    for i in range(100):\n",
    "        np.dot(x, u, out=p)\n",
    "        np.max(p, axis=1, out=z)\n",
    "        p -= z.reshape(-1,1)\n",
    "        np.exp(p, out=p)\n",
    "        np.sum(p, axis=1, out=z)\n",
    "        p /= z.reshape(-1,1)\n",
    "        p -= y\n",
    "        np.dot(x.T, p, out=g)\n",
    "        gnorm = np.sum(g*g)/(n*d*k)  \n",
    "        print(i, gnorm)\n",
    "        if gnorm < 1e-4:\n",
    "            break\n",
    "        gi = (1-li) / nextli\n",
    "        t[:] = gi * w[:]                    \n",
    "        np.subtract(u, cho_solve(c,g), out=w)\n",
    "        np.add((1-gi) * w, t, out=u)\n",
    "        li=nextli\n",
    "        nextli=(1+np.sqrt(1+4*li*li))/2;\n",
    "    return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 0.8397011968\n",
      "1 0.6272526336\n",
      "2 0.3982569472\n",
      "3 0.2001848832\n",
      "4 0.0937900416\n",
      "5 0.0461673056\n",
      "6 0.0245917536\n",
      "7 0.0141343504\n",
      "8 0.0086822848\n",
      "9 0.0056444348\n",
      "10 0.0038504268\n",
      "11 0.0027363636\n",
      "12 0.0020139444\n",
      "13 0.0015277179\n",
      "14 0.0011897801\n",
      "15 0.0009482662\n",
      "16 0.0007713892\n",
      "17 0.0006389983\n",
      "18 0.0005379338\n",
      "19 0.0004593725\n",
      "20 0.000397266875\n",
      "21 0.000347388725\n",
      "22 0.0003067337\n",
      "23 0.0002731358\n",
      "24 0.000245013425\n",
      "25 0.00022119805\n",
      "26 0.0002008180375\n",
      "27 0.0001832147125\n",
      "28 0.0001678860875\n",
      "29 0.0001544444625\n",
      "30 0.000142587025\n",
      "31 0.00013207325\n",
      "32 0.0001227101875\n",
      "33 0.00011434135\n",
      "34 0.0001068359625\n",
      "35 0.000100084975\n",
      "36 9.399680625e-05\n"
     ]
    }
   ],
   "source": [
    "w = agd(rffxv,y,c=rcf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def agdilr(x,y,c=None):\n",
    "    nx, d = x.shape\n",
    "    ny, k = y.shape \n",
    "    assert nx == ny\n",
    "    n = nx\n",
    "    from scipy.linalg import cho_factor, cho_solve\n",
    "    from scipy.special import expit\n",
    "    if c is None:\n",
    "        xx=np.dot(x.T,x)\n",
    "        c = cho_factor(xx+1e-12*np.eye(xx.shape[0]))\n",
    "    li, nextli = (1,1)\n",
    "    u = np.zeros((d,k), dtype=np.float32)\n",
    "    w = np.zeros((d,k), dtype=np.float32)\n",
    "    t = np.zeros((d,k), dtype=np.float32)\n",
    "    p = np.zeros((n,k), dtype=np.float32)\n",
    "    z = np.zeros(n, dtype=np.float32)\n",
    "    g = np.zeros((d,k), dtype=np.float32)\n",
    "    for i in range(100):\n",
    "        np.dot(x, u, out=p)\n",
    "        expit(p, p)\n",
    "        p -= y\n",
    "        np.dot(x.T, p, out=g)\n",
    "        gnorm = np.sum(g*g)/(n*d*k)  \n",
    "        print(i, gnorm)\n",
    "        if gnorm < 1e-4:\n",
    "            break\n",
    "        gi = (1-li) / nextli\n",
    "        t[:] = gi * w[:]                    \n",
    "        np.subtract(u, cho_solve(c,g), out=w)\n",
    "        np.add((1-gi) * w, t, out=u)\n",
    "        li=nextli\n",
    "        nextli=(1+np.sqrt(1+4*li*li))/2;\n",
    "    return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 61.2215422976\n",
      "1 34.9249339392\n",
      "2 21.1837599744\n",
      "3 12.1602449408\n",
      "4 6.9052243968\n",
      "5 3.9821623296\n",
      "6 2.363305984\n",
      "7 1.4517590016\n",
      "8 0.9247705088\n",
      "9 0.6106792448\n",
      "10 0.4174430208\n",
      "11 0.294748416\n",
      "12 0.2144181632\n",
      "13 0.1602550784\n",
      "14 0.1227006336\n",
      "15 0.0959672576\n",
      "16 0.0764628672\n",
      "17 0.0619050432\n",
      "18 0.0508107584\n",
      "19 0.0421957376\n",
      "20 0.0353932768\n",
      "21 0.0299428608\n",
      "22 0.0255201408\n",
      "23 0.0218923456\n",
      "24 0.0188891968\n",
      "25 0.0163838352\n",
      "26 0.0142800096\n",
      "27 0.0125034736\n",
      "28 0.0109959824\n",
      "29 0.0097112656\n",
      "30 0.0086120728\n",
      "31 0.0076681112\n",
      "32 0.006854528\n",
      "33 0.006150772\n",
      "34 0.0055398092\n",
      "35 0.0050074096\n",
      "36 0.0045416876\n",
      "37 0.0041326784\n",
      "38 0.0037720372\n",
      "39 0.003452746\n",
      "40 0.0031689024\n",
      "41 0.0029155568\n",
      "42 0.0026885412\n",
      "43 0.0024843476\n",
      "44 0.0023000292\n",
      "45 0.0021330844\n",
      "46 0.0019814108\n",
      "47 0.0018432186\n",
      "48 0.0017169848\n",
      "49 0.0016014062\n",
      "50 0.0014953648\n",
      "51 0.0013978998\n",
      "52 0.0013081731\n",
      "53 0.0012254534\n",
      "54 0.0011490939\n",
      "55 0.001078521\n",
      "56 0.0010132271\n",
      "57 0.0009527599\n",
      "58 0.0008967056\n",
      "59 0.000844701\n",
      "60 0.0007964062\n",
      "61 0.00075151815\n",
      "62 0.0007097567\n",
      "63 0.0006708667\n",
      "64 0.0006346186\n",
      "65 0.00060080085\n",
      "66 0.00056921735\n",
      "67 0.0005396924\n",
      "68 0.00051206495\n",
      "69 0.000486187\n",
      "70 0.0004619219\n",
      "71 0.00043914755\n",
      "72 0.0004177529\n",
      "73 0.00039763195\n",
      "74 0.0003786922\n",
      "75 0.000360847875\n",
      "76 0.000344021375\n",
      "77 0.000328144025\n",
      "78 0.000313148\n",
      "79 0.00029897505\n",
      "80 0.000285570875\n",
      "81 0.0002728863\n",
      "82 0.0002608746\n",
      "83 0.00024949555\n",
      "84 0.000238710375\n",
      "85 0.000228482275\n",
      "86 0.00021877845\n",
      "87 0.00020956695\n",
      "88 0.000200821275\n",
      "89 0.00019251265\n",
      "90 0.00018461695\n",
      "91 0.0001771094625\n",
      "92 0.0001699697875\n",
      "93 0.00016317615\n",
      "94 0.000156710025\n",
      "95 0.000150551075\n",
      "96 0.0001446833\n",
      "97 0.0001390884375\n",
      "98 0.0001337527875\n",
      "99 0.000128662775\n"
     ]
    }
   ],
   "source": [
    "q = agdilr(rffxv,y,c=rcf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 1603.230629888\n",
      "1 1263.609053184\n",
      "2 950.191980544\n",
      "3 644.200726528\n",
      "4 419.374497792\n",
      "5 277.317189632\n",
      "6 189.474635776\n",
      "7 133.558771712\n",
      "8 96.842203136\n",
      "9 72.175239168\n",
      "10 55.305777152\n",
      "11 43.580964864\n",
      "12 35.2965632\n",
      "13 29.335881728\n",
      "14 24.955836416\n",
      "15 21.654812672\n",
      "16 19.091929088\n",
      "17 17.03514624\n",
      "18 15.328274432\n",
      "19 13.868142592\n",
      "20 12.588895232\n",
      "21 11.44962048\n",
      "22 10.425668608\n",
      "23 9.502051328\n",
      "24 8.668997632\n",
      "25 7.919024128\n",
      "26 7.245551616\n",
      "27 6.641988608\n",
      "28 6.10170368\n",
      "29 5.617979904\n",
      "30 5.18426368\n",
      "31 4.794210816\n",
      "32 4.44194048\n",
      "33 4.122111488\n",
      "34 3.830002432\n",
      "35 3.561542144\n",
      "36 3.313305088\n",
      "37 3.082521856\n",
      "38 2.867015424\n",
      "39 2.665138944\n",
      "40 2.4757248\n",
      "41 2.297963264\n",
      "42 2.131347584\n",
      "43 1.97556928\n",
      "44 1.830432128\n",
      "45 1.695756928\n",
      "46 1.571324544\n",
      "47 1.456835456\n",
      "48 1.351885568\n",
      "49 1.255970048\n",
      "50 1.168477696\n",
      "51 1.088727424\n",
      "52 1.01599488\n",
      "53 0.949546304\n",
      "54 0.888651776\n",
      "55 0.832633216\n",
      "56 0.780859264\n",
      "57 0.732776512\n",
      "58 0.687901312\n",
      "59 0.645833408\n",
      "60 0.606246528\n",
      "61 0.568886848\n",
      "62 0.533563616\n",
      "63 0.500131424\n",
      "64 0.468502656\n",
      "65 0.438607264\n",
      "66 0.410406304\n",
      "67 0.383880192\n",
      "68 0.359006848\n",
      "69 0.335779776\n",
      "70 0.314184224\n",
      "71 0.294202048\n",
      "72 0.275809344\n",
      "73 0.25897176\n",
      "74 0.243642432\n",
      "75 0.229759712\n",
      "76 0.217249904\n",
      "77 0.206027072\n",
      "78 0.19599168\n",
      "79 0.187040528\n",
      "80 0.17906064\n",
      "81 0.1719384\n",
      "82 0.165554352\n",
      "83 0.159799536\n",
      "84 0.154574176\n",
      "85 0.149784208\n",
      "86 0.145350864\n",
      "87 0.141208096\n",
      "88 0.137304576\n",
      "89 0.133595896\n",
      "90 0.130057112\n",
      "91 0.126667816\n",
      "92 0.12341856\n",
      "93 0.120308448\n",
      "94 0.117342792\n",
      "95 0.114526976\n",
      "96 0.111865616\n",
      "97 0.109367704\n",
      "98 0.107038248\n",
      "99 0.104874752\n"
     ]
    }
   ],
   "source": [
    "wlin = agd(xv,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savez('mnist_models.npz',v=vv,omega=rff.omega,b=rff.b,z=rff.normalizer,wlin=wlin,w=w,q=q)"
   ]
  },
  {
   "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}