reddit12k_jmlr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch    \n",
    "torch.manual_seed(123)\n",
    "import random\n",
    "random.seed(123)\n",
    "\n",
    "import core.config as config\n",
    "\n",
    "import torch.nn as nn\n",
    "import os\n",
    "import shutil\n",
    "import itertools\n",
    "\n",
    "from chofer_tda_datasets import Reddit12kJmlr\n",
    "from chofer_tda_datasets.transforms import Hdf5GroupToDict\n",
    "from core.utils import *\n",
    "from sklearn.model_selection import ShuffleSplit\n",
    "from torch.utils.data import Subset\n",
    "\n",
    "from torchph.nn.slayer import SLayerRationalHat, LinearRationalStretchedBirthLifeTimeCoordinateTransform, prepare_batch\n",
    "\n",
    "\n",
    "from collections import Counter, defaultdict\n",
    "\n",
    "%matplotlib notebook\n",
    "\n",
    "\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = str(3)\n",
    "\n",
    "class RedditCollate:   \n",
    "    def __init__(self, dataset, cuda=True):\n",
    "        self.cuda = cuda\n",
    "        self.dataset = dataset\n",
    "        \n",
    "    def __call__(self, sample_target_iter):\n",
    "        x, y = dict_sample_target_iter_concat(sample_target_iter)\n",
    "        \n",
    "        for k in self.dataset.keys_not_essential:\n",
    "            batch_view = x[k]\n",
    "            x[k] = prepare_batch(batch_view, 2)\n",
    "            \n",
    "        for k in self.dataset.keys_essential:\n",
    "            batch_view = x[k]\n",
    "            x[k] = prepare_batch(batch_view, 1)            \n",
    "\n",
    "        y = torch.LongTensor(y)    \n",
    "\n",
    "        if self.cuda:\n",
    "            # Shifting the necessary parts of the prepared batch to the cuda\n",
    "            x = {k: collection_cascade(v,\n",
    "                                       lambda x: isinstance(x, tuple),\n",
    "                                       lambda x: (x[0].cuda(), x[1].cuda(), x[2], x[3]))\n",
    "                 for k, v in x.items()}\n",
    "\n",
    "            y = y.cuda()\n",
    "\n",
    "        return x, y\n",
    "    \n",
    "\n",
    "class train_env:\n",
    "    n_epochs = 300\n",
    "    lr_initial = 0.5\n",
    "    lr_epoch_step = 20\n",
    "    batch_size = 100\n",
    "    train_size = 0.9\n",
    "    nu = 0.01\n",
    "    momentum = 0.9\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = Reddit12kJmlr(data_root_folder_path=config.paths.data_root_dir)\n",
    "dataset.keys_essential = ('dim_0_ess', 'dim_1_ess')\n",
    "dataset.keys_not_essential = ('dim_0',)\n",
    "dataset.keys_of_interrest = dataset.keys_essential + dataset.keys_not_essential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unsqueeze_essentials(x):\n",
    "    for k in dataset.keys_essential:\n",
    "        if len(x[k]) > 0:\n",
    "            x[k] = x[k].unsqueeze(1)\n",
    "            \n",
    "    return x\n",
    "\n",
    "def coordinate_transform(x):\n",
    "    t = LinearRationalStretchedBirthLifeTimeCoordinateTransform(nu=train_env.nu)\n",
    "    for k in dataset.keys_not_essential:\n",
    "        x[k] = t(x[k])   \n",
    "        \n",
    "    return x\n",
    "\n",
    "dataset.data_transforms = \\\n",
    "[\n",
    "    Hdf5GroupToDict(),\n",
    "    numpy_to_torch_cascade,   \n",
    "    coordinate_transform, \n",
    "    unsqueeze_essentials\n",
    "]\n",
    "dataset.target_transforms = [lambda x: int(x) - 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reddit_collate = RedditCollate(dataset, cuda=False)         "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def LinearCell(n_in, n_out):\n",
    "    m = nn.Sequential(nn.Linear(n_in, n_out), \n",
    "                      nn.BatchNorm1d(n_out), \n",
    "                      nn.ReLU(), \n",
    "                     )\n",
    "    m.out_features = m[0].out_features\n",
    "    return m\n",
    "\n",
    "\n",
    "def Slayer(n_elements, point_dim):\n",
    "    return SLayerRationalHat(n_elements, point_dimension=2, radius_init=500)   \n",
    "\n",
    "\n",
    "class Reddit5KModel(nn.Module):    \n",
    "    def __init__(self):\n",
    "        super().__init__()  \n",
    "        \n",
    "        dim_0_n_elements = 150\n",
    "        dim_0_ess_n_elements = 5\n",
    "        dim_1_ess_n_elements = 50\n",
    "        \n",
    "        self.dim_0 = nn.Sequential(\n",
    "            Slayer(dim_0_n_elements, 2),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "        self.dim_0_ess = nn.Sequential(            \n",
    "            Slayer(dim_0_ess_n_elements, 1),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "        self.dim_1_ess = nn.Sequential(            \n",
    "            Slayer(dim_1_ess_n_elements, 1),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "        \n",
    "\n",
    "        cls_in_size = dim_0_n_elements + dim_0_ess_n_elements + dim_1_ess_n_elements\n",
    "        n_1 = 2*cls_in_size#int(cls_in_size/2)\n",
    "#         cls_in_size = self.dim_0_linear.out_features + self.dim_0_ess_linear.out_features + self.dim_1_ess_linear.out_features\n",
    "        self.cls = nn.Sequential(          \n",
    "                               nn.Dropout(0.2),\n",
    "                               LinearCell(cls_in_size, n_1),\n",
    "                               nn.Dropout(0.2),\n",
    "                               LinearCell(n_1, int(n_1/4)),\n",
    "                               LinearCell(int(n_1/4), int(n_1/16)),\n",
    "                               nn.Linear(int(n_1/16), 11))\n",
    "                         \n",
    "    def forward(self, x):\n",
    "        x_dim_0 = self.dim_0(x['dim_0'])        \n",
    "        \n",
    "        x_dim_0_ess = self.dim_0_ess(x['dim_0_ess'])   \n",
    "        \n",
    "        x_dim_1_ess = self.dim_1_ess(x['dim_1_ess'])       \n",
    "        \n",
    "        x = torch.cat([x_dim_0, x_dim_0_ess, x_dim_1_ess], dim=1)\n",
    "        x = self.cls(x)\n",
    "            \n",
    "        return x\n",
    "        \n",
    "    def centers_init(self, sample_target_iter):\n",
    "        x, _ = dict_sample_target_iter_concat(sample_target_iter)\n",
    "        \n",
    "        def get_centers(samples, n_centers):    \n",
    "            s = torch.cat(samples, dim=0)\n",
    "            \n",
    "            s = list({tuple(row) for row in s})\n",
    "            \n",
    "            if len(s) < n_centers:\n",
    "                while len(s) < n_centers:\n",
    "                    s += s \n",
    "                    \n",
    "            s = np.array(s)\n",
    "            \n",
    "            if s.ndim == 1:\n",
    "                s = s.reshape(-1, 1)\n",
    "            \n",
    "            print('Running k-means++')\n",
    "            kmeans = sklearn.cluster.KMeans(n_clusters=n_centers, \n",
    "                                            init='k-means++', \n",
    "                                            random_state=123, \n",
    "                                            n_jobs=10, \n",
    "                                            n_init=1)                           \n",
    "            kmeans.fit(s)\n",
    "            centers = kmeans.cluster_centers_\n",
    "            centers = torch.from_numpy(centers).float()    \n",
    "            return centers\n",
    "        \n",
    "        for key, slayer in zip(['dim_0', 'dim_0_ess', 'dim_1_ess'], [self.dim_0[0], self.dim_0_ess[0], self.dim_1_ess[0]]):\n",
    "            samples = x[key]       \n",
    "            \n",
    "            centers = get_centers(samples, slayer.centers.size(0))\n",
    "            slayer.centers.data = centers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def experiment(train_slayers):      \n",
    "    \n",
    "    stats_of_runs = []\n",
    "    \n",
    "    splitter = ShuffleSplit(n_splits=10, \n",
    "                            train_size=train_env.train_size, \n",
    "                            test_size=1-train_env.train_size, \n",
    "                            random_state=123)\n",
    "    train_test_splits = list(splitter.split(X=dataset.targets, y=dataset.targets))\n",
    "\n",
    "    train_test_splits = [(train_i.tolist(), test_i.tolist()) for train_i, test_i in train_test_splits]\n",
    "    \n",
    "    for run_i, (train_i, test_i) in enumerate(train_test_splits):\n",
    "        print('')\n",
    "        print('Run', run_i)\n",
    "        \n",
    "        model = Reddit5KModel()\n",
    "        print('Model built!')\n",
    "        \n",
    "        model.centers_init([dataset[i] for i in train_i])\n",
    "        model.cuda()\n",
    "\n",
    "        stats = defaultdict(list)\n",
    "        stats_of_runs.append(stats)  \n",
    "        \n",
    "        stats['model'] = model        \n",
    "        \n",
    "        opt=torch.optim.SGD(model.parameters() if train_slayers else model.cls.parameters(), \n",
    "                            lr=train_env.lr_initial, \n",
    "                            momentum=train_env.momentum)\n",
    "\n",
    "        for i_epoch in range(1, train_env.n_epochs+1):      \n",
    "\n",
    "            model.train()  \n",
    "            \n",
    "            dl_train = DataLoader(Subset(dataset, train_i),\n",
    "                              batch_size=train_env.batch_size, \n",
    "                              collate_fn=reddit_collate,\n",
    "                              shuffle=True,                               \n",
    "                              num_workers=5,\n",
    "                              drop_last=True)\n",
    "\n",
    "            dl_test = DataLoader(Subset(dataset, test_i),\n",
    "                                 batch_size=train_env.batch_size, \n",
    "                                 collate_fn=reddit_collate, \n",
    "                                 num_workers=5)\n",
    "\n",
    "            epoch_loss = 0\n",
    "\n",
    "            if i_epoch % train_env.lr_epoch_step == 0:\n",
    "                adapt_lr(opt, lambda lr: lr*0.5)\n",
    "\n",
    "            for i_batch, (x, y) in enumerate(dl_train, 1):   \n",
    "                \n",
    "                x = {k: collection_cascade(v,\n",
    "                                       lambda x: isinstance(x, tuple),\n",
    "                                       lambda x: (x[0].cuda(), x[1].cuda(), x[2], x[3]))\n",
    "                     for k, v in x.items()}\n",
    "\n",
    "                y = y.cuda()\n",
    "\n",
    "                y = torch.autograd.Variable(y)\n",
    "\n",
    "                def closure():\n",
    "                    opt.zero_grad()\n",
    "                    y_hat = model(x)            \n",
    "                    loss = nn.functional.cross_entropy(y_hat, y)   \n",
    "                    loss.backward()\n",
    "                    return loss\n",
    "\n",
    "                loss = opt.step(closure)\n",
    "\n",
    "                epoch_loss += float(loss)\n",
    "                stats['loss_by_batch'].append(float(loss))\n",
    "                stats['centers'].append(model.dim_0[0].centers.data.cpu().numpy())\n",
    "                stats['radius'].append(model.dim_0[0].radius.data.cpu().numpy())\n",
    "\n",
    "                print(\"Epoch {}/{}, Batch {}/{}\".format(i_epoch, train_env.n_epochs, i_batch, len(dl_train)), end=\"     \\r\")\n",
    "\n",
    "            stats['train_loss_by_epoch'].append(epoch_loss/len(dl_train))\n",
    "\n",
    "#             print(\"\\n\\r testing...\")\n",
    "            model.eval()    \n",
    "            true_samples = 0\n",
    "            seen_samples = 0\n",
    "            epoch_test_loss = 0\n",
    "            for i_batch, (x, y) in enumerate(dl_test):\n",
    "                \n",
    "                x = {k: collection_cascade(v,\n",
    "                                       lambda x: isinstance(x, tuple),\n",
    "                                       lambda x: (x[0].cuda(), x[1].cuda(), x[2], x[3]))\n",
    "                 for k, v in x.items()}\n",
    "\n",
    "                y = y.cuda()\n",
    "\n",
    "                y_hat = model(x)\n",
    "                epoch_test_loss += float(nn.functional.cross_entropy(y_hat, torch.autograd.Variable(y.cuda())).data)\n",
    "\n",
    "                y_hat = y_hat.max(dim=1)[1].data.long()\n",
    "\n",
    "                true_samples += (y_hat == y).sum()\n",
    "                seen_samples += y.size(0)  \n",
    "\n",
    "            test_acc = true_samples.item()/seen_samples\n",
    "            stats['test_accuracy'].append(test_acc)\n",
    "            stats['test_loss_by_epoch'].append(epoch_test_loss/len(dl_test))\n",
    "    return stats_of_runs\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "res_trained_slayer = experiment(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accs = [np.mean(s['test_accuracy'][-10:]) for s in res_trained_slayer]\n",
    "print(accs)\n",
    "print(np.mean(accs))\n",
    "print(np.std(accs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_rigid_slayer = experiment(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "accs = [np.mean(s['test_accuracy'][-10:]) for s in res_rigid_slayer]\n",
    "print(accs)\n",
    "print(np.mean(accs))\n",
    "print(np.std(accs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "stats = res_trained_slayer[-1]\n",
    "\n",
    "c_start = stats['centers'][3]\n",
    "c_end = stats['centers'][-1]\n",
    "\n",
    "plt.plot(c_start[:,0], c_start[:, 1], 'bo', label='center initialization')\n",
    "plt.plot(c_end[:,0], c_end[:, 1], 'ro', label='center learned')\n",
    "\n",
    "all_centers = numpy.stack(stats['centers'], axis=0)\n",
    "for i in range(all_centers.shape[1]):\n",
    "    points = all_centers[:,i, :]\n",
    "    plt.plot(points[:, 0], points[:, 1], '-k')\n",
    "    \n",
    "plt.legend()\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(stats['train_loss_by_epoch'], label='train_loss')\n",
    "plt.plot(stats['test_loss_by_epoch'], label='test_loss')\n",
    "plt.plot(stats['test_accuracy'], label='test_acccuracy')\n",
    "plt.legend()\n",
    "\n",
    "plt.figure()\n",
    "radius = stats['radius']\n",
    "radius = np.stack(radius, axis=1)\n",
    "for curve in radius:\n",
    "    plt.plot(curve)\n",
    "plt.show()"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}