Added LeNet model implemented with PyTorch modules

.github/workflows/learn-github-actions.yml

@@ -1,5 +1,5 @@
 name: learn-github-actions
-run-name: ${{ github.actor }} is learning GitHub Actions
+run-name: ${{ github.actor }} has made changes
 on: [push]
 jobs:
   check-bats-version:

notebooks/functional.ipynb

@@ -0,0 +1,105 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 79,
+   "metadata": {
+    "collapsed": true,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from jax import numpy as jnp\n",
+    "from jax import jit, vmap"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "outputs": [],
+   "source": [
+    "@jit\n",
+    "def sigmoid(x):\n",
+    "    return 1 / (1 + jnp.exp(-1 * x))"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 75,
+   "outputs": [],
+   "source": [
+    "@jit\n",
+    "def relu(x):\n",
+    "    return x * (x > 0)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 98,
+   "outputs": [],
+   "source": [
+    "@jit\n",
+    "@vmap\n",
+    "def softmax(x):\n",
+    "    \"\"\"\n",
+    "    >>> jnp.sum(softmax(jnp.array([[1, 2, 4], [1, 2, 3], [1, 2, 3]])), axis=1)\n",
+    "    DeviceArray([1., 1., 1.], dtype=float32)\n",
+    "    \"\"\"\n",
+    "    return jnp.exp(x) / jnp.sum(jnp.exp(x))"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

notebooks/model.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 3,
    "metadata": {
     "collapsed": true,
     "pycharm": {
@@ -14,15 +14,82 @@
     "import random\n",
     "\n",
     "import numpy as np\n",
-    "from jax import numpy as jnp"
+    "from src.functional import sigmoid\n",
+    "\n",
+    "from jax import grad\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 106,
+   "outputs": [],
+   "source": [
+    "class DenseLayer:\n",
+    "    def __init__(self, total_nodes, input_size, activation=sigmoid):\n",
+    "        self.total_nodes = total_nodes\n",
+    "        self.input_size = input_size\n",
+    "        self.weights = np.random.rand(total_nodes, self.input_size + 1)\n",
+    "        self.activation = activation\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x_biased = np.concatenate((x, np.ones((len(x), 1))), axis=1)\n",
+    "        y = self.weights @ x_biased.T\n",
+    "        return self.activation(y)\n",
+    "\n",
+    "    def backprop(self, gradient):\n",
+    "        "
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 107,
+   "outputs": [],
+   "source": [
+    "l = DenseLayer(5, 5)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 108,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "Array([[0.7291562 ],\n       [0.84321564],\n       [0.8657799 ],\n       [0.8525716 ],\n       [0.89164424]], dtype=float32)"
+     },
+     "execution_count": 108,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "l.forward(np.random.rand(1, 5))"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
    "outputs": [],
    "source": [
-    "class"
+    "import torch\n",
+    "from torch import nn"
    ],
    "metadata": {
     "collapsed": false,
@@ -30,6 +97,72 @@
      "name": "#%%\n"
     }
    }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "outputs": [],
+   "source": [
+    "class CNN(nn.Module):\n",
+    "    def __init__(self, input_channels, num_classes):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.feature_layers = [input_channels, 6, 16, 120]\n",
+    "        self.kernels = [5, 5, 5]\n",
+    "        self.pools = [2, 2]\n",
+    "        self.feature_activations = [nn.Tanh for _ in range(len(self.channels) - 1)]\n",
+    "\n",
+    "        self.classifier_layers = [120, num_classes]\n",
+    "        self.classifier_activations = [nn.Tanh for _ in range(len(self.classifier_layers))]\n",
+    "\n",
+    "        feature_layers = []\n",
+    "        for idx, layer in enumerate(list(zip(self.feature_layers[:-1], self.feature_layers[1:]))):\n",
+    "            feature_layers.append(\n",
+    "                nn.Conv2d(in_channels=layer[0], out_channels=layer[1], kernel_size=self.kernels[idx])\n",
+    "            )\n",
+    "            feature_layers.append(self.feature_activations[idx])\n",
+    "\n",
+    "            if idx != len(self.feature_activations):\n",
+    "                feature_layers.append(nn.MaxPool2d(kernel_size=self.pools[2]))\n",
+    "\n",
+    "\n",
+    "        classifier_layers = []\n",
+    "        for idx, layer in enumerate(list(zip(self.classifier_layers[:-1], self.classifier_layers[1:]))):\n",
+    "            classifier_layers.append(\n",
+    "                nn.Linear(in_features=layer[0], out_features=layer[1])\n",
+    "            )\n",
+    "\n",
+    "            if idx != len(self.classifier_activations):\n",
+    "                classifier_layers.append(self.classifier_activations[idx])\n",
+    "\n",
+    "\n",
+    "        self.feature_extractor = nn.Sequential(*feature_layers)\n",
+    "        self.classifier = nn.Sequential(*classifier_layers)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x = self.feature_extractor(x)\n",
+    "        y = self.classifier_layers(torch.flatten(x, 1))\n",
+    "        p = nn.functional.softmax(y, dim=1)\n",
+    "        return y, p"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
   }
  ],
  "metadata": {

src/models.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+# coding: utf-8
+import torch
+from torch import nn
+
+
+class CNN(nn.Module):
+    def __init__(self, input_channels, num_classes):
+        super().__init__()
+
+        self.feature_layers = [input_channels, 6, 16]
+        self.kernels = [5, 5]
+        self.pools = [2, 2]
+        self.feature_activations = [nn.ReLU for _ in range(len(self.feature_layers) - 1)]
+
+        self.classifier_layers = [400, 120, 84, num_classes]
+        self.classifier_activations = [nn.ReLU for _ in range(len(self.classifier_layers) - 1)]
+
+        feature_layers = []
+        for idx, layer in enumerate(list(zip(self.feature_layers[:-1], self.feature_layers[1:]))):
+            feature_layers.append(
+                nn.Conv2d(in_channels=layer[0], out_channels=layer[1],
+                          kernel_size=self.kernels[idx], padding=2 if idx == 0 else 0)
+            )
+            feature_layers.append(self.feature_activations[idx]())
+            feature_layers.append(nn.MaxPool2d(kernel_size=self.pools[idx]))
+
+        classifier_layers = []
+        for idx, layer in enumerate(list(zip(self.classifier_layers[:-1], self.classifier_layers[1:]))):
+            classifier_layers.append(
+                nn.Linear(in_features=layer[0], out_features=layer[1])
+            )
+            classifier_layers.append(self.classifier_activations[idx]())
+
+        self.feature_extractor = nn.Sequential(*feature_layers)
+        self.classifier = nn.Sequential(*classifier_layers)
+
+    def forward(self, x):
+        x = self.feature_extractor(x)
+        y = self.classifier(torch.flatten(x))
+        return y