KerasPruningInCallback.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "sublime-adelaide",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "copyrighted-length",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "general-commonwealth",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.compat.v1.disable_eager_execution()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "proper-eugene",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('./keras-surgeon/src/')\n",
    "# sys.path.append('./tf-keras-surgeon/src/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "encouraging-passing",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kerassurgeon.operations import delete_channels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "equipped-trademark",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "from tensorflow.keras.callbacks import Callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cathedral-birth",
   "metadata": {},
   "outputs": [],
   "source": [
    "def copy_model_inplace(model1, model2):\n",
    "    model2.compile(\n",
    "        optimizer=model1.optimizer,\n",
    "        loss=model1.loss,\n",
    "        metrics=model1.metrics)    \n",
    "    \n",
    "    for k in dir(model1):\n",
    "        if k.startswith('__'):\n",
    "            continue\n",
    "\n",
    "        #if k in ['_compiled_trainable_state', '_layer_call_argspecs']:\n",
    "        #    continue\n",
    "\n",
    "        if callable(getattr(model, k)):\n",
    "            continue\n",
    "\n",
    "        if hasattr(model.__class__, k) and isinstance(getattr(model.__class__, k), property):\n",
    "            continue\n",
    "\n",
    "        try:\n",
    "            setattr(model1, k, getattr(model2, k))\n",
    "        except Exception as e:\n",
    "            pass    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "american-colon",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "grand-luther",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/ruslan/anaconda3/lib/python3.8/site-packages/keras/layers/normalization/batch_normalization.py:532: _colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/ruslan/anaconda3/lib/python3.8/site-packages/keras/layers/normalization/batch_normalization.py:532: _colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n"
     ]
    }
   ],
   "source": [
    "model = MobileNetV2(input_shape=(224, 224, 3), weights=None, include_top=True, classes=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "promising-warning",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(optimizer=Adam(learning_rate=1e-3),\n",
    "              loss='binary_crossentropy',\n",
    "              metrics=['binary_accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "necessary-niagara",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PruningCallback(Callback):\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        pruned_model = delete_channels(self.model, self.model.layers[1], [0, 1])\n",
    "        copy_model_inplace(self.model, pruned_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "mounted-lloyd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([3, 3, 3, 32])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.layers[1].weights[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "abstract-cooperative",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 32 samples\n",
      "Epoch 1/4\n",
      "32/32 [==============================] - ETA: 0s - loss: 0.6004 - binary_accuracy: 0.2500Deleting 2/32 channels from layer: Conv1\n",
      "expanded_conv_depthwise (3, 3) (112, 112, 32)\n",
      "(32, 112, 112)\n",
      "[0, 1]\n",
      "Config {'name': 'expanded_conv_depthwise', 'trainable': True, 'dtype': 'float32', 'kernel_size': (3, 3), 'strides': (1, 1), 'padding': 'same', 'data_format': 'channels_last', 'dilation_rate': (1, 1), 'groups': 1, 'activation': 'linear', 'use_bias': False, 'bias_initializer': {'class_name': 'Zeros', 'config': {}}, 'bias_regularizer': None, 'activity_regularizer': None, 'bias_constraint': None, 'depth_multiplier': 1, 'depthwise_initializer': {'class_name': 'GlorotUniform', 'config': {'seed': None}}, 'depthwise_regularizer': None, 'depthwise_constraint': None}\n",
      "Outputmask Shape (112, 112, 32)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ruslan/anaconda3/lib/python3.8/site-packages/keras/engine/functional.py:1410: CustomMaskWarning: Custom mask layers require a config and must override get_config. When loading, the custom mask layer must be passed to the custom_objects argument.\n",
      "  layer_config = serialize_layer_fn(layer)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32/32 [==============================] - 14s 442ms/sample - loss: 0.6004 - binary_accuracy: 0.2500\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-7-f3456320983e>:14: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
      "  if callable(getattr(model, k)):\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/4\n",
      "32/32 [==============================] - ETA: 0s - loss: 0.2590 - binary_accuracy: 0.2500Deleting 2/30 channels from layer: Conv1\n",
      "expanded_conv_depthwise (3, 3) (112, 112, 30)\n",
      "(30, 112, 112)\n",
      "[0, 1]\n",
      "Config {'name': 'expanded_conv_depthwise', 'trainable': True, 'dtype': 'float32', 'kernel_size': (3, 3), 'strides': (1, 1), 'padding': 'same', 'data_format': 'channels_last', 'dilation_rate': (1, 1), 'groups': 1, 'activation': 'linear', 'use_bias': False, 'bias_initializer': {'class_name': 'Zeros', 'config': {}}, 'bias_regularizer': None, 'activity_regularizer': None, 'bias_constraint': None, 'depth_multiplier': 1, 'depthwise_initializer': {'class_name': 'GlorotUniform', 'config': {'seed': None}}, 'depthwise_regularizer': None, 'depthwise_constraint': None}\n",
      "Outputmask Shape (112, 112, 30)\n",
      "32/32 [==============================] - 12s 377ms/sample - loss: 0.2590 - binary_accuracy: 0.2500\n",
      "Epoch 3/4\n",
      "32/32 [==============================] - ETA: 0s - loss: 0.0173 - binary_accuracy: 0.2500Deleting 2/28 channels from layer: Conv1\n",
      "expanded_conv_depthwise (3, 3) (112, 112, 28)\n",
      "(28, 112, 112)\n",
      "[0, 1]\n",
      "Config {'name': 'expanded_conv_depthwise', 'trainable': True, 'dtype': 'float32', 'kernel_size': (3, 3), 'strides': (1, 1), 'padding': 'same', 'data_format': 'channels_last', 'dilation_rate': (1, 1), 'groups': 1, 'activation': 'linear', 'use_bias': False, 'bias_initializer': {'class_name': 'Zeros', 'config': {}}, 'bias_regularizer': None, 'activity_regularizer': None, 'bias_constraint': None, 'depth_multiplier': 1, 'depthwise_initializer': {'class_name': 'GlorotUniform', 'config': {'seed': None}}, 'depthwise_regularizer': None, 'depthwise_constraint': None}\n",
      "Outputmask Shape (112, 112, 28)\n",
      "32/32 [==============================] - 15s 457ms/sample - loss: 0.0173 - binary_accuracy: 0.2500\n",
      "Epoch 4/4\n",
      "32/32 [==============================] - ETA: 0s - loss: 0.0032 - binary_accuracy: 0.2500Deleting 2/26 channels from layer: Conv1\n",
      "expanded_conv_depthwise (3, 3) (112, 112, 26)\n",
      "(26, 112, 112)\n",
      "[0, 1]\n",
      "Config {'name': 'expanded_conv_depthwise', 'trainable': True, 'dtype': 'float32', 'kernel_size': (3, 3), 'strides': (1, 1), 'padding': 'same', 'data_format': 'channels_last', 'dilation_rate': (1, 1), 'groups': 1, 'activation': 'linear', 'use_bias': False, 'bias_initializer': {'class_name': 'Zeros', 'config': {}}, 'bias_regularizer': None, 'activity_regularizer': None, 'bias_constraint': None, 'depth_multiplier': 1, 'depthwise_initializer': {'class_name': 'GlorotUniform', 'config': {'seed': None}}, 'depthwise_regularizer': None, 'depthwise_constraint': None}\n",
      "Outputmask Shape (112, 112, 26)\n",
      "32/32 [==============================] - 18s 550ms/sample - loss: 0.0032 - binary_accuracy: 0.2500\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f3ef08fa1f0>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(x=np.random.random((32, 224, 224, 3)), y=np.random.randint(0, 2, size=(32)), epochs=4, callbacks=[PruningCallback()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "final-quantity",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([3, 3, 3, 24])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.layers[1].weights[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "latter-signal",
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}