convert_2_tflite.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 转成 tflite 格式\n",
    "\n",
    "模型训练好了之后要实际应用。对于模型部署有很多成熟的方案，服务器端有 Nvidia 的 TensorRT， Intel 的 OpenVINO 等。移动端也有各种不同的选择，这里仅仅使用 TensorFlow官方的 TfLite做演示。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先还是模型的定义，因为我们要使用模型，那么首先要知道模型是什么样子的。这里直接将训练中使用的模型代码拿来即可。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModel(tf.keras.Model):\n",
    "    def __init__(self):\n",
    "        super(MyModel, self).__init__()\n",
    "        # 模型有两个主要部分，特征提取层和分类器\n",
    "\n",
    "        # 这里是特征提取层\n",
    "        self.feature = tf.keras.models.Sequential()\n",
    "        self.feature.add(self.conv(64))\n",
    "        self.feature.add(self.conv(64, add_pooling=True))\n",
    "\n",
    "        self.feature.add(self.conv(128))\n",
    "        self.feature.add(self.conv(128, add_pooling=True))\n",
    "\n",
    "        self.feature.add(self.conv(256))\n",
    "        self.feature.add(self.conv(256))\n",
    "        self.feature.add(self.conv(256, add_pooling=True))\n",
    "\n",
    "        self.feature.add(self.conv(512))\n",
    "        self.feature.add(self.conv(512))\n",
    "        self.feature.add(self.conv(512, add_pooling=True))\n",
    "\n",
    "        self.feature.add(self.conv(512))\n",
    "        self.feature.add(self.conv(512))\n",
    "        self.feature.add(self.conv(512, add_pooling=True))\n",
    "\n",
    "        self.feature.add(tf.keras.layers.GlobalAveragePooling2D())\n",
    "\n",
    "        self.feature.add(tf.keras.layers.Dense(4096, activation=\"relu\"))\n",
    "        self.feature.add(tf.keras.layers.BatchNormalization())\n",
    "        self.feature.add(tf.keras.layers.Dense(4096, activation=\"relu\"))\n",
    "        self.feature.add(tf.keras.layers.BatchNormalization())\n",
    "\n",
    "        self.feature.add(tf.keras.layers.Dropout(0.5))\n",
    "\n",
    "        # 这个简单的机构是分类器\n",
    "        self.pred = tf.keras.layers.Dense(100)\n",
    "\n",
    "    def conv(self, filters, add_pooling=False):\n",
    "        # 模型大量使用重复模块构建，\n",
    "        # 这里将重复模块提取出来，简化模型构建过程\n",
    "        model = tf.keras.models.Sequential(\n",
    "            [\n",
    "                tf.keras.layers.Conv2D(filters, 3, padding=\"same\", activation=\"relu\"),\n",
    "                tf.keras.layers.BatchNormalization(),\n",
    "            ]\n",
    "        )\n",
    "        if add_pooling:\n",
    "            model.add(\n",
    "                tf.keras.layers.MaxPool2D(\n",
    "                    pool_size=(2, 2), strides=None, padding=\"same\"\n",
    "                )\n",
    "            )\n",
    "        return model\n",
    "\n",
    "    def call(self, x):\n",
    "        # call 用来定义模型各个结构之间的运算关系\n",
    "\n",
    "        x = self.feature(x)\n",
    "        return tf.nn.softmax(self.pred(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "需要注意上面模型定义里面， call 的最后，输出 pred 之前，多加了一个 softmax 。这是因为原始模型输出的是未经 softmax 激活的 logits，虽然大小可以判断，但不具备可读性。在服务器端可以很方便的使用 python 或其他语言的 api 转换为 概率模式，但是这里在 tflite 中，我们需要提前做一步转换。\n",
    "\n",
    "不用担心这里的一个操作会导致什么错误，因为 softmax 本质只是一个操作，没有新引入的变量，所以模型结构虽然有了变化，但是之前训练用的 ckpt 这里还是可以成功加载。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "有了模型的定义之后，我们可以加载训练好的模型，跟模型训练的时候类似，我们可以直接加载模型训练中的 checkpoint。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model lodaded\n"
     ]
    }
   ],
   "source": [
    "model = MyModel()\n",
    "ckpt = tf.train.get_checkpoint_state(\".\")\n",
    "if ckpt:\n",
    "    model.load_weights(ckpt.model_checkpoint_path)\n",
    "    print(\"model lodaded\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "另外 tflite converter 转换模型的时候，需要提前知道模型的输入输出尺寸，而我们这里只是加载了模型，还没有做过运算，模型输入输出尺寸无从得知，因此下面要进行一步运算，使用一个随手生成的数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Layer my_model is casting an input tensor from dtype float64 to the layer's dtype of float32, which is new behavior in TensorFlow 2.  The layer has dtype float32 because it's dtype defaults to floatx.\n",
      "\n",
      "If you intended to run this layer in float32, you can safely ignore this warning. If in doubt, this warning is likely only an issue if you are porting a TensorFlow 1.X model to TensorFlow 2.\n",
      "\n",
      "To change all layers to have dtype float64 by default, call `tf.keras.backend.set_floatx('float64')`. To change just this layer, pass dtype='float64' to the layer constructor. If you are the author of this layer, you can disable autocasting by passing autocast=False to the base Layer constructor.\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 100), dtype=float32, numpy=\n",
       "array([[6.97592739e-03, 3.50881674e-05, 2.99538297e-05, 9.69473086e-03,\n",
       "        3.26752067e-02, 3.38254467e-05, 4.28270176e-03, 7.65197736e-04,\n",
       "        1.38145303e-02, 4.08006570e-04, 9.17119253e-03, 2.54529703e-04,\n",
       "        8.80055968e-03, 2.19247136e-02, 1.69239156e-02, 5.20564441e-04,\n",
       "        3.75563162e-04, 6.06878370e-04, 6.36276382e-04, 2.92343990e-04,\n",
       "        1.59879904e-02, 1.20332901e-04, 2.19296475e-04, 6.52071601e-03,\n",
       "        7.25438772e-03, 3.08928802e-03, 6.33171108e-03, 7.33097717e-02,\n",
       "        2.71364674e-03, 1.47544459e-04, 2.18871821e-04, 1.30667789e-02,\n",
       "        8.32068981e-05, 1.33725349e-02, 3.60220962e-04, 6.36282144e-04,\n",
       "        1.49801731e-01, 4.53493092e-03, 6.90849358e-03, 2.50804285e-03,\n",
       "        2.28558830e-03, 2.93506833e-04, 2.17173016e-04, 6.17122687e-02,\n",
       "        4.72656591e-03, 8.37453044e-05, 5.78573830e-02, 9.84658371e-04,\n",
       "        2.13809806e-04, 4.66867127e-02, 1.56987342e-03, 4.50606551e-03,\n",
       "        8.11364222e-03, 2.54643857e-01, 7.31819382e-05, 2.53929669e-04,\n",
       "        1.14092210e-04, 1.92468477e-04, 4.86115430e-04, 5.69899520e-03,\n",
       "        2.48575471e-02, 6.32752490e-04, 1.73449895e-04, 7.40825897e-03,\n",
       "        6.58004865e-05, 4.25508991e-02, 1.53475814e-03, 3.16437334e-04,\n",
       "        7.92641134e-04, 1.53127138e-03, 1.25322631e-03, 1.36587163e-03,\n",
       "        4.86249628e-04, 4.44090801e-06, 5.09359059e-04, 2.73428050e-05,\n",
       "        1.49042840e-04, 1.34321291e-03, 1.14154152e-03, 1.02548487e-03,\n",
       "        5.78144354e-05, 4.14523629e-06, 2.89787731e-05, 8.65990575e-03,\n",
       "        5.34841965e-05, 4.77750227e-03, 5.88633156e-05, 3.68489673e-05,\n",
       "        3.07894632e-04, 6.99279262e-05, 1.73384428e-03, 1.56005364e-04,\n",
       "        2.17610766e-04, 1.08611165e-03, 2.03584181e-03, 1.68785002e-04,\n",
       "        1.03705411e-03, 5.88253699e-03, 2.21012553e-04, 1.15140676e-04]],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(np.ones([1,128,128,3]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "根据官方文档，tflite converter可以接受 3 种不同形式的输入，这里使用其中一种。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/dl/miniconda3/lib/python3.7/site-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "INFO:tensorflow:Assets written to: vgg/assets\n"
     ]
    }
   ],
   "source": [
    "tf.saved_model.save(model, \"vgg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "之后直接调用 tflite converter 进行转换就可以了，当然转换过程中也可以做一些优化操作，包括量化等等，具体参考官方文档即可。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "converter = tf.lite.TFLiteConverter.from_saved_model(\"vgg\")\n",
    "tflite_model = converter.convert()\n",
    "\n",
    "with open(\"converted_model.tflite\", \"wb\") as of:\n",
    "    of.write(tflite_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "生成的模型就可以拿到官方的 tflite demo apk 里面使用了。\n",
    "\n",
    "> 不过还有些地方需要注意，官方的 apk 对于输入的图像做了 90 度旋转，但是旋转的方向错了，导致我们这里的模型需要把手机倒立才能正确识别。\n",
    "> 另外，这里的模型训练的时候，输入数据做了处理，所以 apk 里面 输入后"
   ]
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}