ResNet With TensorFlow 2.0.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_layer(nf, kernel=(3, 3), strides=(1, 1), padding:str='valid', activation:str='relu', kernel_init='he_normal', input_layer:bool=False, input_shape=None, bn:bool = True, bn_mom=0.9):\n",
    "    layers = []\n",
    "    if input_layer:\n",
    "        layers.append(tf.keras.layers.Conv2D(nf, kernel, strides, padding, kernel_initializer=kernel_init, input_shape=input_shape))\n",
    "    else:\n",
    "        layers.append(tf.keras.layers.Conv2D(nf, kernel, strides, padding, kernel_initializer=kernel_init))\n",
    "    if bn:\n",
    "        layers.append(tf.keras.layers.BatchNormalization(momentum=bn_mom))\n",
    "    layers.append(tf.keras.layers.Activation(activation=activation))\n",
    "    return tf.keras.Sequential(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResBlock(tf.keras.Model):\n",
    "    \"\"\"Residual Block Consisting of 2 (conv + batch_norm + relu) layers\n",
    "    nf: number of filters\n",
    "    shortcut: there is a an identity  \n",
    "    \"\"\"\n",
    "    def __init__(self, nf, bottleneck:bool=False):\n",
    "        super(ResBlock, self).__init__()\n",
    "        self.nf, self.bottleneck = nf, bottleneck\n",
    "        if self.bottleneck:\n",
    "            self.strides = (2, 2)\n",
    "            self.bottleneck_block = conv_layer(self.nf, kernel=(1, 1), strides=self.strides, padding='same') \n",
    "        else:\n",
    "            self.strides = (1, 1)\n",
    "            \n",
    "        self.conv1 = conv_layer(self.nf, kernel=(3, 3), strides=self.strides, padding='same')\n",
    "        self.conv2 = conv_layer(self.nf, kernel=(3, 3), padding='same')\n",
    "    def call(self, x, training=False):\n",
    "        skip = x\n",
    "        fx = x\n",
    "        if self.bottleneck:\n",
    "            skip = self.bottleneck_block(x)\n",
    "        fx = self.conv1(fx)\n",
    "        fx = self.conv2(fx)\n",
    "        print(skip.shape, fx.shape)\n",
    "        fx = tf.add(skip, fx) # may be function to minimize memory consumption.\n",
    "        fx = tf.nn.relu(fx)\n",
    "        return fx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNet34(tf.keras.Model):\n",
    "    def __init__(self, input_shape, include_top=True, n_classes=1000):\n",
    "        super(ResNet34, self).__init__()\n",
    "        self.ishape, self.include_top, self.n_classes = input_shape, include_top, n_classes\n",
    "        self.conv1 = conv_layer(64, kernel=(7, 7), padding='same', strides=2, input_layer=True, input_shape=self.ishape)\n",
    "        self.maxpool = tf.keras.layers.MaxPool2D((3, 3), 2, padding='same')\n",
    "        self.res_blocks = tf.keras.Sequential()\n",
    "        for nf, num_blocks, downscale in zip([64, 128, 256, 512],\n",
    "                                             [3, 4, 6, 3], \n",
    "                                             [False, True, True, True]):\n",
    "            for i in range(num_blocks):\n",
    "                if i == 0 and downscale:\n",
    "                    self.res_blocks.add(ResBlock(nf, bottleneck=True))\n",
    "                else:\n",
    "                    self.res_blocks.add(ResBlock(nf))\n",
    "        if include_top:\n",
    "            self.avg = tf.keras.layers.GlobalAveragePooling2D()\n",
    "            self.flatten = tf.keras.layers.Flatten()\n",
    "            self.out = tf.keras.layers.Dense(self.n_classes, activation='softmax')\n",
    "            \n",
    "    def call(self, x):\n",
    "        print(x.shape)\n",
    "        x = self.conv1(x)\n",
    "        print(x.shape, x.shape)\n",
    "        x = self.maxpool(x)\n",
    "        print(x.shape, x.shape)\n",
    "        x = self.res_blocks(x)\n",
    "        if self.include_top:\n",
    "            x = self.avg(x)\n",
    "            print(x.shape, x.shape)\n",
    "            x = self.flatten(x)\n",
    "            x = self.out(x)\n",
    "        print(x.shape)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 224, 224, 3)\n",
      "(1, 112, 112, 64) (1, 112, 112, 64)\n",
      "(1, 56, 56, 64) (1, 56, 56, 64)\n",
      "(1, 56, 56, 64) (1, 56, 56, 64)\n",
      "(1, 56, 56, 64) (1, 56, 56, 64)\n",
      "(1, 56, 56, 64) (1, 56, 56, 64)\n",
      "(1, 28, 28, 128) (1, 28, 28, 128)\n",
      "(1, 28, 28, 128) (1, 28, 28, 128)\n",
      "(1, 28, 28, 128) (1, 28, 28, 128)\n",
      "(1, 28, 28, 128) (1, 28, 28, 128)\n",
      "(1, 14, 14, 256) (1, 14, 14, 256)\n",
      "(1, 14, 14, 256) (1, 14, 14, 256)\n",
      "(1, 14, 14, 256) (1, 14, 14, 256)\n",
      "(1, 14, 14, 256) (1, 14, 14, 256)\n",
      "(1, 14, 14, 256) (1, 14, 14, 256)\n",
      "(1, 14, 14, 256) (1, 14, 14, 256)\n",
      "(1, 7, 7, 512) (1, 7, 7, 512)\n",
      "(1, 7, 7, 512) (1, 7, 7, 512)\n",
      "(1, 7, 7, 512) (1, 7, 7, 512)\n",
      "(1, 512) (1, 512)\n",
      "(1, 1000)\n",
      "Model: \"res_net34_23\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "sequential_534 (Sequential)  (None, 112, 112, 64)      9728      \n",
      "_________________________________________________________________\n",
      "max_pooling2d_21 (MaxPooling multiple                  0         \n",
      "_________________________________________________________________\n",
      "sequential_535 (Sequential)  multiple                  21300480  \n",
      "_________________________________________________________________\n",
      "global_average_pooling2d_19  multiple                  0         \n",
      "_________________________________________________________________\n",
      "flatten_8 (Flatten)          multiple                  0         \n",
      "_________________________________________________________________\n",
      "dense_19 (Dense)             multiple                  513000    \n",
      "=================================================================\n",
      "Total params: 21,823,208\n",
      "Trainable params: 21,806,184\n",
      "Non-trainable params: 17,024\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "res34 = ResNet34((224, 224, 3))\n",
    "res34.build((1, 224, 224, 3))\n",
    "res34.summary()"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}