{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "[![Download Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0.0-alpha/resource/_static/logo_notebook_en.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r2.0.0-alpha/golden_stick/en/deployment/mindspore_convert.ipynb) [![Download Sample Code](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0.0-alpha/resource/_static/logo_download_code_en.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r2.0.0-alpha/golden_stick/en/deployment/mindspore_convert.py) [![View source files in Gitee](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0.0-alpha/resource/_static/logo_source_en.png)](https://gitee.com/mindspore/docs/blob/r2.0.0-alpha/docs/golden_stick/docs/source_en/deployment/convert.ipynb)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# MindSpore Golden Stick Network Conversion\n",
    "\n",
    "There are three ways to convert the model and export MindIR:\n",
    "\n",
    "1. Export MindIR after training;\n",
    "2. Export MindIR from ckpt;\n",
    "3. Configure the algorithm before training to automatically export MindIR."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Necessary prerequisites\n",
    "\n",
    "Firstly download dataset and create Lenet, and for demonstration convenience, we implemented one simplest MindSpore Golden Stick algorithm, called FooAlgo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from download import download\n",
    "\n",
    "import mindspore\n",
    "from mindspore import nn, Model, Tensor, export, context\n",
    "from mindspore.train import Accuracy\n",
    "from mindspore.train.callback import ModelCheckpoint\n",
    "import mindspore.dataset as ds\n",
    "import mindspore.dataset.vision as vision\n",
    "import mindspore.dataset.transforms as transforms\n",
    "from mindspore.dataset.vision import Inter\n",
    "from mindspore import dtype as mstype\n",
    "from mindspore.common.initializer import Normal\n",
    "from mindspore_gs import CompAlgo\n",
    "\n",
    "# Download data from open datasets\n",
    "url = \"https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/\" \\\n",
    "      \"notebook/datasets/MNIST_Data.zip\"\n",
    "path = download(url, \"./\", kind=\"zip\", replace=True)\n",
    "\n",
    "def create_dataset(data_path, batch_size=32, num_parallel_workers=1):\n",
    "    \"\"\"\n",
    "    create dataset for train or test\n",
    "    \"\"\"\n",
    "    # define dataset\n",
    "    mnist_ds = ds.MnistDataset(data_path)\n",
    "\n",
    "    resize_height, resize_width = 32, 32\n",
    "    rescale = 1.0 / 255.0\n",
    "    rescale_nml = 1 / 0.3081\n",
    "    shift_nml = -1 * 0.1307 / 0.3081\n",
    "\n",
    "    # define map operations\n",
    "    resize_op = vision.Resize((resize_height, resize_width), interpolation=Inter.LINEAR)  # Bilinear mode\n",
    "    rescale_nml_op = vision.Rescale(rescale_nml * rescale, shift_nml)\n",
    "    hwc2chw_op = vision.HWC2CHW()\n",
    "    type_cast_op = transforms.TypeCast(mstype.int32)\n",
    "\n",
    "    # apply map operations on images\n",
    "    mnist_ds = mnist_ds.map(operations=type_cast_op, input_columns=\"label\", num_parallel_workers=num_parallel_workers)\n",
    "    mnist_ds = mnist_ds.map(operations=resize_op, input_columns=\"image\", num_parallel_workers=num_parallel_workers)\n",
    "    mnist_ds = mnist_ds.map(operations=rescale_nml_op, input_columns=\"image\", num_parallel_workers=num_parallel_workers)\n",
    "    mnist_ds = mnist_ds.map(operations=hwc2chw_op, input_columns=\"image\", num_parallel_workers=num_parallel_workers)\n",
    "\n",
    "    # apply DatasetOps\n",
    "    mnist_ds = mnist_ds.shuffle(buffer_size=1024)\n",
    "    mnist_ds = mnist_ds.batch(batch_size, drop_remainder=True)\n",
    "\n",
    "    return mnist_ds\n",
    "\n",
    "train_dataset = create_dataset(\"MNIST_Data/train\", 32, 1)\n",
    "print(\"train dataset output shape: \", train_dataset.output_shapes())\n",
    "\n",
    "# initial network\n",
    "class LeNet5(nn.Cell):\n",
    "    def __init__(self, num_class=10, num_channel=1, include_top=True):\n",
    "        super(LeNet5, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(num_channel, 6, 5, pad_mode='valid')\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')\n",
    "        self.relu = nn.ReLU()\n",
    "        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        self.include_top = include_top\n",
    "        if self.include_top:\n",
    "            self.flatten = nn.Flatten()\n",
    "            self.fc1 = nn.Dense(16 * 5 * 5, 120, weight_init=Normal(0.02))\n",
    "            self.fc2 = nn.Dense(120, 84, weight_init=Normal(0.02))\n",
    "            self.fc3 = nn.Dense(84, num_class, weight_init=Normal(0.02))\n",
    "\n",
    "\n",
    "    def construct(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.max_pool2d(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.max_pool2d(x)\n",
    "        if not self.include_top:\n",
    "            return x\n",
    "        x = self.flatten(x)\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# set graph mode\n",
    "context.set_context(mode=context.GRAPH_MODE)\n",
    "\n",
    "# for demonstration convenience, we implemented one simplest MindSpore Golden Stick algorithm, called FooAlgo\n",
    "class FooAlgo(CompAlgo):\n",
    "    def apply(self, network: nn.Cell) -> nn.Cell:\n",
    "        return network\n",
    "\n",
    "print(\"init ok.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "```python\n",
    "train dataset output shape: [[32, 1, 32, 32], [32]]\n",
    "init ok.\n",
    "```"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export MindIR after training\n",
    "\n",
    "MindSpore Golden Stick algorithms provide a `convert` interface to convert network, and then you can use `mindspore.export` to export MindIR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "## 1) Create network and dataset.\n",
    "network = LeNet5(10)\n",
    "train_dataset = create_dataset(\"MNIST_Data/train\", 32, 1)\n",
    "## 2) Create an algorithm instance.\n",
    "algo = FooAlgo()\n",
    "## 3) Apply MindSpore Golden Stick algorithm to origin network.\n",
    "network_opt = algo.apply(network)\n",
    "## 4) Set up Model.\n",
    "net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction=\"mean\")\n",
    "net_opt = nn.Momentum(network_opt.trainable_params(), 0.01, 0.9)\n",
    "model = Model(network_opt, net_loss, net_opt, metrics={\"Accuracy\": Accuracy()})\n",
    "cbs = [ModelCheckpoint(prefix='network', directory='ckpt/')]\n",
    "## 5) Config callback in model.train, start training.\n",
    "cbs.extend(algo.callbacks())\n",
    "model.train(1, train_dataset, callbacks=cbs)\n",
    "## 6） Convert network.\n",
    "net_deploy = algo.convert(network_opt)\n",
    "## 7) Export MindIR\n",
    "inputs = Tensor(np.ones([32, 1, 32, 32]).astype(np.float32))  # user define\n",
    "export(net_deploy, inputs, file_name=\"net_1.mindir\", file_format=\"MINDIR\")\n",
    "## 8) Test MindIR\n",
    "file_path = \"./net_1.mindir\"\n",
    "file_path = os.path.realpath(file_path)\n",
    "if not os.path.exists(file_path):\n",
    "    print(\"Export MindIR failed!!!\")\n",
    "else:\n",
    "    print(\"Export MindIR success! MindIR path is: \", file_path)\n",
    "test_inputs = Tensor(np.ones([32, 1, 32, 32]).astype(np.float32))\n",
    "graph = mindspore.load(file_path)\n",
    "net = nn.GraphCell(graph)\n",
    "output = net(test_inputs)\n",
    "print(\"Test output MindIR success, result shape is: \", output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "```python\n",
    "Export MindIR success! MindIR path is: /home/workspace/golden_stick/net_1.mindir\n",
    "Test output MindIR success, result shape is: (32, 10)\n",
    "```"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Export from ckpt\n",
    "\n",
    "Using the ckpt file after training, call `convert` and `mindspore.export` interfaces to export MindIR.\n",
    "\n",
    "> Please run the sample code in the previous section first, this section requires the ckpt file generated by the training process in the previous section."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "## 1) Create network and dataset.\n",
    "network = LeNet5(10)\n",
    "train_dataset = create_dataset(\"MNIST_Data/train\", 32, 1)\n",
    "## 2) Create an algorithm instance.\n",
    "algo = FooAlgo()\n",
    "## 3) Apply MindSpore Golden Stick algorithm to origin network.\n",
    "network_opt = algo.apply(network)\n",
    "## 4） Convert network.\n",
    "net_deploy = algo.convert(network_opt, ckpt_path=\"ckpt/network-1_1875.ckpt\")  # ckpt from previous section\n",
    "## 5) Export MindIR\n",
    "inputs = Tensor(np.ones([32, 1, 32, 32]).astype(np.float32))  # user define\n",
    "export(net_deploy, inputs, file_name=\"net_2.mindir\", file_format=\"MINDIR\")\n",
    "## 6) Test MindIR\n",
    "file_path = \"./net_2.mindir\"\n",
    "file_path = os.path.realpath(file_path)\n",
    "if not os.path.exists(file_path):\n",
    "    print(\"Export MindIR failed!!!\")\n",
    "else:\n",
    "    print(\"Export MindIR success! MindIR path is: \", file_path)\n",
    "test_inputs = Tensor(np.ones([32, 1, 32, 32]).astype(np.float32))\n",
    "graph = mindspore.load(file_path)\n",
    "net = nn.GraphCell(graph)\n",
    "output = net(test_inputs)\n",
    "print(\"Test output MindIR success, result shape is: \", output.shape)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "```python\n",
    "Export MindIR success! MindIR path is: /home/workspace/golden_stick/net_2.mindir\n",
    "Test output MindIR success, result shape is: (32, 10)\n",
    "```"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Configure the algorithm to automatically export MindIR\n",
    "\n",
    "Configure the algorithm `set_save_mindir` interface before training to automatically export MindIR after training.\n",
    "\n",
    "> a) When using MindIR generated in this way for inference, the input shape of the MindIR must be consistent with the shape of the dataset at the time of training.\n",
    ">\n",
    "> b) There are two necessary operations to configure the algorithm to automatically export MindIR, `set_save_mindir(True)` and add the algorithm callback function `callbacks=algo.callbacks()` when configuring the callback function in `model.train` . MindIR output path `save_mindir_path` is saved by default as `./network.mindir` if not configured."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "## 1) Create network and dataset.\n",
    "network = LeNet5(10)\n",
    "train_dataset = create_dataset(\"MNIST_Data/train\", 32, 1)\n",
    "## 2) Create an algorithm instance.\n",
    "algo = FooAlgo()\n",
    "## 3) Enable automatically export MindIR after training.\n",
    "algo.set_save_mindir(save_mindir=True)\n",
    "## 4) Set MindIR output path, the default value for the path is 'network.mindir'.\n",
    "algo.set_save_mindir_path(save_mindir_path=\"net_3.mindir\")\n",
    "## 5) Apply MindSpore Golden Stick algorithm to origin network.\n",
    "network_opt = algo.apply(network)\n",
    "## 6) Set up Model.\n",
    "net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction=\"mean\")\n",
    "net_opt = nn.Momentum(network_opt.trainable_params(), 0.01, 0.9)\n",
    "model = Model(network_opt, net_loss, net_opt, metrics={\"Accuracy\": Accuracy()})\n",
    "## 7) Config callback in model.train, start training, then MindIR will be exported.\n",
    "model.train(1, train_dataset, callbacks=algo.callbacks())\n",
    "## 8) Test MindIR\n",
    "file_path = \"./net_3.mindir\"\n",
    "file_path = os.path.realpath(file_path)\n",
    "if not os.path.exists(file_path):\n",
    "    print(\"Export MindIR failed!!!\")\n",
    "else:\n",
    "    print(\"Export MindIR success! MindIR path is: \", file_path)\n",
    "test_inputs = Tensor(np.ones([32, 1, 32, 32]).astype(np.float32))\n",
    "graph = mindspore.load(file_path)\n",
    "net = nn.GraphCell(graph)\n",
    "output = net(test_inputs)\n",
    "print(\"Test output MindIR success, result shape is: \", output.shape)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "```python\n",
    "Export MindIR success! MindIR path is: /home/workspace/golden_stick/net_3.mindir\n",
    "Test output MindIR success, result shape is: (32, 10)\n",
    "```"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.5 ('mindspore')",
   "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.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "8c9da313289c39257cb28b126d2dadd33153d4da4d524f730c81a4aaccbd2ca7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}