Tsetlin_Lottery_Machine.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Tsetlin Lottery Machine",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/littell/scratchwork/blob/master/Tsetlin_Lottery_Machine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7tiCgeg0CQ7m",
        "colab_type": "text"
      },
      "source": [
        "## Installing PyTorch 1.1.0\n",
        "\n",
        "While we don't need autograd during the reinforcement learning phase, we can make use of the GPU acceleration, and having the whole thing be a PyTorch module means we can train it as an ordinary neural network afterward."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-IwgbV3X8ynf",
        "colab_type": "code",
        "outputId": "04a70ce6-d441-4cf2-c0e1-0afb4738ba90",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 340
        }
      },
      "source": [
        "from os.path import exists\n",
        "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
        "\n",
        "platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
        "cuda_output = !ldconfig -p|grep cudart.so|sed -e 's/.*\\.\\([0-9]*\\)\\.\\([0-9]*\\)$/cu\\10/'    \n",
        "accelerator = cuda_output[0] if exists('/dev/nvidia0') else 'cpu'\n",
        "print(f\"Device = {accelerator}\")\n",
        "version='1.1.0'\n",
        "torch_url=f\"http://download.pytorch.org/whl/{accelerator}/torch-{version}-{platform}-linux_x86_64.whl\"\n",
        "!pip install -U {torch_url} torchvision"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Device = cu100\n",
            "Collecting torch==1.1.0 from http://download.pytorch.org/whl/cu100/torch-1.1.0-cp36-cp36m-linux_x86_64.whl\n",
            "\u001b[?25l  Downloading http://download.pytorch.org/whl/cu100/torch-1.1.0-cp36-cp36m-linux_x86_64.whl (770.7MB)\n",
            "\u001b[K     |████████████████████████████████| 770.7MB 2.1MB/s \n",
            "\u001b[?25hRequirement already up-to-date: torchvision in /usr/local/lib/python3.6/dist-packages (0.3.0)\n",
            "Requirement already satisfied, skipping upgrade: numpy in /usr/local/lib/python3.6/dist-packages (from torch==1.1.0) (1.16.4)\n",
            "Requirement already satisfied, skipping upgrade: six in /usr/local/lib/python3.6/dist-packages (from torchvision) (1.12.0)\n",
            "Requirement already satisfied, skipping upgrade: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision) (4.3.0)\n",
            "Requirement already satisfied, skipping upgrade: olefile in /usr/local/lib/python3.6/dist-packages (from pillow>=4.1.1->torchvision) (0.46)\n",
            "Installing collected packages: torch\n",
            "  Found existing installation: torch 1.1.0\n",
            "    Uninstalling torch-1.1.0:\n",
            "      Successfully uninstalled torch-1.1.0\n",
            "Successfully installed torch-1.1.0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "torch"
                ]
              }
            }
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PnrVBXnrV_KD",
        "colab_type": "text"
      },
      "source": [
        "## Import libraries into Python, then build the dataset loaders"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x8HNRNICWDO6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.utils.data as data\n",
        "from torch.autograd import Variable\n",
        "from torchvision import datasets, transforms"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sSzJS8pCWHdl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class DatasetLoader(object):\n",
        "  \n",
        "    def __init__(self, dataset='MNIST', dataset_dir='./data', batch_size=100):\n",
        "        dataset_ = {\n",
        "            'MNIST': datasets.MNIST,\n",
        "            'CIFAR10': datasets.CIFAR10\n",
        "        }[dataset]\n",
        "        \n",
        "        self.size = { \n",
        "            'MNIST': 28*28, \n",
        "            'CIFAR10': 3*32*32 \n",
        "        }[dataset]\n",
        "\n",
        "        transform = {\n",
        "            'MNIST': transforms.ToTensor(),\n",
        "            'CIFAR10': transforms.Compose([\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
        "                ])\n",
        "        }[dataset]\n",
        "\n",
        "        train_dataset = dataset_(root=dataset_dir,\n",
        "                                 train=True,\n",
        "                                 transform=transform,\n",
        "                                 download=True)\n",
        "\n",
        "        self.train = data.DataLoader(dataset=train_dataset,\n",
        "                                      batch_size=batch_size,\n",
        "                                      shuffle=True)\n",
        "\n",
        "        test_dataset = dataset_(root=dataset_dir,\n",
        "                                 train=False,\n",
        "                                 transform=transform,\n",
        "                                 download=True)\n",
        "\n",
        "        self.test = data.DataLoader(dataset=test_dataset,\n",
        "                                      batch_size=batch_size,\n",
        "                                      shuffle=False)\n",
        "        "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hjSG2QQ2WeiK",
        "colab_type": "text"
      },
      "source": [
        "## Make the module that we're going to train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5k6PJ3dYmqLJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "class TsetlinSelector(nn.Module):\n",
        "  \n",
        "    def __init__(self, size, max=20):\n",
        "        super(TsetlinSelector, self).__init__()\n",
        "        self.size = size\n",
        "        self.max = max\n",
        "        \n",
        "        dtype = torch.FloatTensor\n",
        "        \n",
        "        # initialize a vector of the appropriate size, drawn from a Bernoulli distribution\n",
        "        self.automata = nn.Parameter((torch.randn(size) > 0.5).type(dtype), requires_grad=False)\n",
        "      \n",
        "    def forward(self, x):\n",
        "        return x * (self.automata > 0.0).float()\n",
        "      \n",
        "    def reward(self, p):\n",
        "        ''' Reward p% of automata '''\n",
        "        automata_on = (self.automata > 0.0).float()                  # 1.0 for automata that are \"on\", 0.0 otherwise\n",
        "        reward_mask = (torch.rand_like(self.automata) < p).float()   # Only reward p% of the automata\n",
        "        self.automata += (automata_on * 2 - 1) * reward_mask         # Add 1.0 for automata that are \"on\", -1.0 otherwise\n",
        "        self.automata.clamp_(-self.max, self.max+1)\n",
        "        \n",
        "    def penalize(self, p):\n",
        "        ''' Penalize p% of automata '''\n",
        "        automata_on = (self.automata > 0.0).float()                  # 1.0 for automata that are \"on\", 0.0 otherwise\n",
        "        reward_mask = (torch.rand_like(self.automata) < p).float()   # Only penalize p% of the automata\n",
        "        self.automata -= (automata_on * 2 - 1) * reward_mask         # Add 1.0 for automata that are \"on\", -1.0 otherwise\n",
        "        \n",
        "    def __repr__(self):\n",
        "        return f\"Automata={self.automata.data}, size={self.automata.shape}, type={type(self.automata)}\"\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j24JkaeNWk4n",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch.nn.functional as F\n",
        "\n",
        "      \n",
        "class Net(nn.Module):\n",
        "  \n",
        "    def __init__(self, image_size=28*28, h1_size = 500, h2_size = 100, output_size=10, dropout_rate=0.1):\n",
        "        super(Net, self).__init__()\n",
        "        self.image_size, self.output_size = image_size, output_size\n",
        "        self.h1_size, self.h2_size = h1_size, h2_size\n",
        "        self.dropout_rate = dropout_rate\n",
        "        \n",
        "        \"\"\"3-Layer Fully-connected NN\"\"\"\n",
        "        #self.net = nn.Sequential(\n",
        "            #nn.Linear(image_dim, 500),\n",
        "            #dropout(0.2, 0.19, 500, dropout_method),\n",
        "            #nn.ReLU(),\n",
        "            #nn.Linear(500, 100),\n",
        "            #dropout(0.5, 0.4, 100, dropout_method),\n",
        "            #nn.ReLU(),\n",
        "            #nn.Linear(100, 10)\n",
        "        #)\n",
        "        ''' The above was the original network, a standard architecture for testing\n",
        "        regularization methods.  I decomposed it below in order to have more direct\n",
        "        access to the parameters '''\n",
        "        \n",
        "        dtype = torch.FloatTensor\n",
        "        \n",
        "        self.m1 = TsetlinSelector(h1_size)\n",
        "        self.m2 = TsetlinSelector(h2_size)\n",
        "        \n",
        "        self.W1 = nn.Parameter(torch.randn(h1_size, image_size).type(dtype), requires_grad=True)\n",
        "        #torch.nn.init.xavier_uniform_(self.W1.data)\n",
        "        self.b1 = nn.Parameter(torch.randn(h1_size).type(dtype) * 0.001, requires_grad=True)\n",
        "        self.W2 = nn.Parameter(torch.randn(h2_size, h1_size).type(dtype), requires_grad=True)\n",
        "        #torch.nn.init.xavier_uniform_(self.W2.data)\n",
        "        self.b2 = nn.Parameter(torch.randn(h2_size).type(dtype) * 0.001, requires_grad=True)\n",
        "        self.W3 = nn.Parameter(torch.randn(output_size, h2_size).type(dtype), requires_grad=True)\n",
        "        #torch.nn.init.xavier_uniform_(self.W3.data)\n",
        "        self.b3 = nn.Parameter(torch.randn(output_size).type(dtype) * 0.001, requires_grad=True)\n",
        "        \n",
        "        self.dropout1 = nn.Dropout(dropout_rate)\n",
        "        self.dropout2 = nn.Dropout(dropout_rate)\n",
        "        \n",
        "    \n",
        "    def forward(self, x):\n",
        "        x = F.linear(x, self.W1, self.b1)\n",
        "        #x = self.dropout1(x)\n",
        "        x = self.m1(x)\n",
        "        x = nn.ReLU()(x)\n",
        "        x = F.linear(x, self.W2, self.b2)\n",
        "        #x = self.dropout2(x)\n",
        "        x = self.m2(x)\n",
        "        x = nn.ReLU()(x)\n",
        "        x = F.linear(x, self.W3, self.b3)\n",
        "        return x\n",
        "      \n",
        "    def reward(self, p):\n",
        "        self.m1.reward(p)\n",
        "        self.m2.reward(p)\n",
        "        \n",
        "    def penalize(self, p):\n",
        "        self.m1.penalize(p)\n",
        "        self.m2.penalize(p)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BFLeTo7mW7w2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Solver(object):\n",
        "  \n",
        "    def __init__(self, net, dataset, dropout_rate=0.1, n_epochs=100, lr=0.001):\n",
        "        self.n_epochs = n_epochs\n",
        "        self.dataset = dataset\n",
        "        self.net = net.cuda()\n",
        "        self.loss_fn = nn.CrossEntropyLoss().cuda()\n",
        "        self.optimizer = optim.Adam(self.net.parameters(), lr=lr)\n",
        "             \n",
        "    def step(self, logits, labels):\n",
        "        loss = self.loss_fn(logits, labels)\n",
        "        self.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        self.optimizer.step()\n",
        "        return float(loss.data)\n",
        "            \n",
        "    def train(self):\n",
        "        self.net.train()\n",
        "        \n",
        "        for epoch_i in (range(self.n_epochs)):\n",
        "            epoch_loss = 0\n",
        "            for images, labels in self.dataset.train:\n",
        "                images = Variable(images).view(-1, self.net.image_size).cuda()\n",
        "                labels = Variable(labels).cuda()\n",
        "\n",
        "                logits = self.net(images)\n",
        "                \n",
        "                epoch_loss += self.step(logits, labels)\n",
        "\n",
        "            epoch_loss /= len(self.dataset.train.dataset)\n",
        "            print('Epoch %s | loss: %s' % (epoch_i, epoch_loss))\n",
        "            \n",
        "    def evaluate(self):\n",
        "        \n",
        "        total = 0.\n",
        "        correct = 0.\n",
        "        self.net.eval()\n",
        "        for images, labels in self.dataset.test:\n",
        "            images = Variable(images).view(-1, self.net.image_size).cuda()\n",
        "\n",
        "            logits = self.net(images)\n",
        "\n",
        "            _, predicted = torch.max(logits.data, 1)\n",
        "            total += labels.size(0) * 1.\n",
        "            correct += (predicted.cpu().data == labels).sum().item()\n",
        "        print('=========================================')\n",
        "        print('Test accuracy: %s' % float(100.0 * correct / total))\n",
        "                "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1NIXkAi11mYi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ReinforcementSolver(Solver):\n",
        "  \n",
        "\n",
        "    def step(self, logits, labels):\n",
        "        _, predicted = torch.max(logits.data, 1)\n",
        "        incorrect = (predicted != labels.data).sum().item()\n",
        "        percent_incorrect = incorrect / len(predicted)\n",
        "        self.net.reward(0.9 * (1. - percent_incorrect))\n",
        "        self.net.penalize(0.1 * percent_incorrect)\n",
        "        return incorrect\n",
        "          "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_kpvVKPoJiwv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import random\n",
        "\n",
        "class MixedSolver(Solver):\n",
        "  \n",
        "    def adam_step(self, logits, labels):\n",
        "        loss = self.loss_fn(logits, labels)\n",
        "        self.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        self.optimizer.step()\n",
        "        return float(loss.data)\n",
        "      \n",
        "    def reinforcement_step(self, logits, labels):\n",
        "        _, predicted = torch.max(logits.data, 1)\n",
        "        incorrect = (predicted != labels.data).sum().item()\n",
        "        percent_incorrect = incorrect / len(predicted)\n",
        "        self.net.reward(0.9 * (1. - percent_incorrect))\n",
        "        self.net.penalize(0.1 * percent_incorrect)\n",
        "        return incorrect\n",
        "      \n",
        "    def step(self, logits, labels):\n",
        "        if random.random() > 0.5:\n",
        "            return self.adam_step(logits, labels)\n",
        "        else:\n",
        "            return self.reinforcement_step(logits, labels)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NlDhYSIeXYTT",
        "colab_type": "code",
        "outputId": "f90427b4-ec14-4179-fcf9-515a9fa5dbdf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1751
        }
      },
      "source": [
        "dataset = DatasetLoader('MNIST', batch_size=10)\n",
        "net = Net(image_size=28*28, h1_size=1000, h2_size=600)\n",
        "standard_solver = MixedSolver(net, dataset, n_epochs=100)\n",
        "standard_solver.train()\n",
        "standard_solver.evaluate()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 0 | loss: 7.475450570703335\n",
            "Epoch 1 | loss: 1.7961884170142026\n",
            "Epoch 2 | loss: 1.1536926459478587\n",
            "Epoch 3 | loss: 0.7362711226151829\n",
            "Epoch 4 | loss: 0.5706952831190033\n",
            "Epoch 5 | loss: 0.43342907724315\n",
            "Epoch 6 | loss: 0.36825820870570725\n",
            "Epoch 7 | loss: 0.2909423584133289\n",
            "Epoch 8 | loss: 0.24558685042870715\n",
            "Epoch 9 | loss: 0.20688970821643768\n",
            "Epoch 10 | loss: 0.17706360710580613\n",
            "Epoch 11 | loss: 0.14983166409963516\n",
            "Epoch 12 | loss: 0.1673409612043659\n",
            "Epoch 13 | loss: 0.1259439014232203\n",
            "Epoch 14 | loss: 0.10925592018023629\n",
            "Epoch 15 | loss: 0.11155467434674211\n",
            "Epoch 16 | loss: 0.10661213779965728\n",
            "Epoch 17 | loss: 0.09809046063778523\n",
            "Epoch 18 | loss: 0.10653187368728216\n",
            "Epoch 19 | loss: 0.08198832884632681\n",
            "Epoch 20 | loss: 0.0781372562187759\n",
            "Epoch 21 | loss: 0.08132444439005743\n",
            "Epoch 22 | loss: 0.08408074555088649\n",
            "Epoch 23 | loss: 0.07266684997335617\n",
            "Epoch 24 | loss: 0.0778233510460457\n",
            "Epoch 25 | loss: 0.06719883419621422\n",
            "Epoch 26 | loss: 0.05784094525367488\n",
            "Epoch 27 | loss: 0.05149668456069271\n",
            "Epoch 28 | loss: 0.05642537143265011\n",
            "Epoch 29 | loss: 0.05698384797345959\n",
            "Epoch 30 | loss: 0.04576452587872967\n",
            "Epoch 31 | loss: 0.048732324567872\n",
            "Epoch 32 | loss: 0.04887359770238571\n",
            "Epoch 33 | loss: 0.04128638673850657\n",
            "Epoch 34 | loss: 0.04397777445393072\n",
            "Epoch 35 | loss: 0.039763067035575146\n",
            "Epoch 36 | loss: 0.053182431011521265\n",
            "Epoch 37 | loss: 0.050133876214164014\n",
            "Epoch 38 | loss: 0.03742264664195778\n",
            "Epoch 39 | loss: 0.0419024805241823\n",
            "Epoch 40 | loss: 0.051440166739236744\n",
            "Epoch 41 | loss: 0.03922014366661157\n",
            "Epoch 42 | loss: 0.04298262208157248\n",
            "Epoch 43 | loss: 0.03709732215643912\n",
            "Epoch 44 | loss: 0.03890704060611339\n",
            "Epoch 45 | loss: 0.03214491254063948\n",
            "Epoch 46 | loss: 0.033942272089361805\n",
            "Epoch 47 | loss: 0.02704593253788844\n",
            "Epoch 48 | loss: 0.03801963169374164\n",
            "Epoch 49 | loss: 0.03587232751046501\n",
            "Epoch 50 | loss: 0.029419641161101902\n",
            "Epoch 51 | loss: 0.038109307324740535\n",
            "Epoch 52 | loss: 0.0350092247891659\n",
            "Epoch 53 | loss: 0.031267369999287566\n",
            "Epoch 54 | loss: 0.030020243528833188\n",
            "Epoch 55 | loss: 0.03408800482993441\n",
            "Epoch 56 | loss: 0.02662511117729118\n",
            "Epoch 57 | loss: 0.027416378690402494\n",
            "Epoch 58 | loss: 0.03254682849923468\n",
            "Epoch 59 | loss: 0.031864596148997136\n",
            "Epoch 60 | loss: 0.03039858922184116\n",
            "Epoch 61 | loss: 0.02935209828975334\n",
            "Epoch 62 | loss: 0.028788587806460177\n",
            "Epoch 63 | loss: 0.025762254554928588\n",
            "Epoch 64 | loss: 0.03069135643450871\n",
            "Epoch 65 | loss: 0.028811525642042694\n",
            "Epoch 66 | loss: 0.030246406579401083\n",
            "Epoch 67 | loss: 0.028399388638349895\n",
            "Epoch 68 | loss: 0.025424569870978898\n",
            "Epoch 69 | loss: 0.029713701671844624\n",
            "Epoch 70 | loss: 0.027126225709812737\n",
            "Epoch 71 | loss: 0.021807332592209485\n",
            "Epoch 72 | loss: 0.03080468023791909\n",
            "Epoch 73 | loss: 0.032317379208801976\n",
            "Epoch 74 | loss: 0.021094295428082714\n",
            "Epoch 75 | loss: 0.024779102239956607\n",
            "Epoch 76 | loss: 0.025555728337774225\n",
            "Epoch 77 | loss: 0.022492151704706764\n",
            "Epoch 78 | loss: 0.02727081231173749\n",
            "Epoch 79 | loss: 0.02362733495235453\n",
            "Epoch 80 | loss: 0.023099412977101202\n",
            "Epoch 81 | loss: 0.024013499566142765\n",
            "Epoch 82 | loss: 0.02282557071904342\n",
            "Epoch 83 | loss: 0.019406987038651034\n",
            "Epoch 84 | loss: 0.02499323732214592\n",
            "Epoch 85 | loss: 0.019057962967907103\n",
            "Epoch 86 | loss: 0.02563307674800541\n",
            "Epoch 87 | loss: 0.02647619085977397\n",
            "Epoch 88 | loss: 0.030066159084807056\n",
            "Epoch 89 | loss: 0.017869264308013954\n",
            "Epoch 90 | loss: 0.01830750308662655\n",
            "Epoch 91 | loss: 0.021460259472186848\n",
            "Epoch 92 | loss: 0.02305941932668599\n",
            "Epoch 93 | loss: 0.020295950623254764\n",
            "Epoch 94 | loss: 0.0233335454537378\n",
            "Epoch 95 | loss: 0.021975688307869154\n",
            "Epoch 96 | loss: 0.01957560849294062\n",
            "Epoch 97 | loss: 0.02215587967793884\n",
            "Epoch 98 | loss: 0.018988811288597328\n",
            "Epoch 99 | loss: 0.01937820955460229\n",
            "=========================================\n",
            "Test accuracy: 97.58\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}