12_KerasCustomLayersExample.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Keras Custom Layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example code for the lecture series \"Machine Learning for Physicists\" by Florian Marquardt\n",
    "\n",
    "Lecture 12\n",
    "\n",
    "See https://machine-learning-for-physicists.org and the current course website linked there!\n",
    "\n",
    "This notebook is distributed under the Attribution-ShareAlike 4.0 International (CC BY-SA 4.0) license:\n",
    "\n",
    "https://creativecommons.org/licenses/by-sa/4.0/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to:\n",
    "- construct a keras custom layer (for a simple convolutional layer that respects periodic boundary conditions in 1D)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from tensorflow.keras import Sequential # Sequential is the neural-network class\n",
    "from tensorflow.keras.layers import Lambda # Dense is the standard network layer\n",
    "\n",
    "# array math:\n",
    "import numpy as np\n",
    "\n",
    "# plotting:\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "matplotlib.rcParams['figure.dpi']=300 # highres display\n",
    "\n",
    "# for updating display \n",
    "# (very simple animation)\n",
    "from IPython.display import clear_output\n",
    "from time import sleep\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PeriodicConvolution(keras.layers.Layer):\n",
    "    # init gets called when setting up the network\n",
    "    def __init__(self,kernel_size=3,**kwargs):\n",
    "        self.kernel_size=kernel_size\n",
    "        super(PeriodicConvolution, self).__init__(**kwargs)\n",
    "\n",
    "    # build gets called when the network is first evaluated\n",
    "    # that means, the size of the input is now known\n",
    "    def build(self, input_shape):\n",
    "        self.w = self.add_weight(\n",
    "            shape=[self.kernel_size],\n",
    "            initializer=\"random_normal\",\n",
    "            trainable=True,\n",
    "        )\n",
    "        \n",
    "    # call gets called whenever the network is evaluated\n",
    "    # (actually it may get called only once, to build up the\n",
    "    # symbolic graph, but that is a detail)\n",
    "    def call(self, inputs):\n",
    "        j0=int((self.kernel_size-1)/2)\n",
    "        # do convolution using tf.roll, which respects periodic\n",
    "        # boundary conditions!\n",
    "        # Note: unlike numpy, we cannot simply assign a zeros array first,\n",
    "        # because we are not allowed to assign afterwards to the Tensor,\n",
    "        # so we just initialize by treating j==0 separately\n",
    "        for j in range(self.kernel_size):\n",
    "            if j==0:\n",
    "                z=self.w[j]*tf.roll(inputs,shift=j-j0,axis=1)\n",
    "            else:\n",
    "                z+=self.w[j]*tf.roll(inputs,shift=j-j0,axis=1)\n",
    "        return z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "Net=Sequential()\n",
    "Net.add(PeriodicConvolution(kernel_size=3))\n",
    "\n",
    "Net.compile(loss='mean_square_error', optimizer='adam')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_in=np.array([[0.,0.,3.,0.,0.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.         -0.20112531  0.03615593 -0.02203112  0.        ]]\n"
     ]
    }
   ],
   "source": [
    "y_out=Net.predict_on_batch(y_in)\n",
    "print(y_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Variable 'sequential/periodic_convolution/Variable:0' shape=(3,) dtype=float32, numpy=array([-0.06704177,  0.01205198, -0.00734371], dtype=float32)>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Net.layers[0].w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Variable 'UnreadVariable' shape=(3,) dtype=float32, numpy=array([-1.,  0.,  1.], dtype=float32)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Net.layers[0].w.assign(np.array([-1,0,1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0. -3.  0.  3.  0.]]\n"
     ]
    }
   ],
   "source": [
    "y_out=Net.predict_on_batch(y_in)\n",
    "print(y_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 3.  0.  0. -3.  0.]]\n"
     ]
    }
   ],
   "source": [
    "y_out=Net.predict_on_batch(np.array([[0.,0.,0.,0.,3.]]))\n",
    "print(y_out)"
   ]
  },
  {
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}