{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7b1f6350",
   "metadata": {},
   "source": [
    "<!--- Licensed to the Apache Software Foundation (ASF) under one -->\n",
    "<!--- or more contributor license agreements.  See the NOTICE file -->\n",
    "<!--- distributed with this work for additional information -->\n",
    "<!--- regarding copyright ownership.  The ASF licenses this file -->\n",
    "<!--- to you under the Apache License, Version 2.0 (the -->\n",
    "<!--- \"License\"); you may not use this file except in compliance -->\n",
    "<!--- with the License.  You may obtain a copy of the License at -->\n",
    "\n",
    "<!---   http://www.apache.org/licenses/LICENSE-2.0 -->\n",
    "\n",
    "<!--- Unless required by applicable law or agreed to in writing, -->\n",
    "<!--- software distributed under the License is distributed on an -->\n",
    "<!--- \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->\n",
    "<!--- KIND, either express or implied.  See the License for the -->\n",
    "<!--- specific language governing permissions and limitations -->\n",
    "<!--- under the License. -->\n",
    "\n",
    "# MXNet Gluon Fit API\n",
    "\n",
    "In this tutorial, you will learn how to use the [Gluon Fit API](https://cwiki.apache.org/confluence/display/MXNET/Gluon+Fit+API+-+Tech+Design) which is the easiest way to train deep learning models using the [Gluon API](../index.rst) in Apache MXNet.\n",
    "\n",
    "With the Fit API, you can train a deep learning model with a minimal amount of code. Just specify the network, loss function and the data you want to train on. You don't need to worry about the boiler plate code to loop through the dataset in batches (often called as 'training loop'). Advanced users can train with bespoke training loops, and many of these use cases will be covered by the Fit API.\n",
    "\n",
    "To demonstrate the Fit API, you will train an image classification model using the [ResNet-18](https://arxiv.org/abs/1512.03385) neural network architecture. The model will be trained using the [Fashion-MNIST dataset](https://github.com/zalandoresearch/fashion-mnist).\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "To complete this tutorial, you will need:\n",
    "\n",
    "- [MXNet](https://mxnet.apache.org/get_started) (The version of MXNet will be >= 1.5.0, you can use `pip install mxnet` to get 1.5.0 release pip package or build from source with master, refer to [MXNet installation](https://mxnet.apache.org/get_started?version=master&platform=linux&language=python&environ=pip&processor=cpu)\n",
    "- [Jupyter Notebook](https://jupyter.org/index.html) (For interactively running the provided .ipynb file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6abb813",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mxnet as mx\n",
    "from mxnet import gluon\n",
    "from mxnet.gluon.model_zoo import vision\n",
    "from mxnet.gluon.contrib.estimator import estimator\n",
    "from mxnet.gluon.contrib.estimator.event_handler import TrainBegin, TrainEnd, EpochEnd, CheckpointHandler\n",
    "\n",
    "gpu_count = mx.device.num_gpus()\n",
    "device = [mx.gpu(i) for i in range(gpu_count)] if gpu_count > 0 else mx.cpu()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26f19e1",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "[Fashion-MNIST](https://github.com/zalandoresearch/fashion-mnist) dataset consists of fashion items divided into ten categories: t-shirt/top, trouser, pullover, dress, coat, sandal, shirt, sneaker, bag and ankle boot.\n",
    "\n",
    "- It has 60,000 grayscale images of size 28 * 28 for training.\n",
    "- It has 10,000 grayscale images of size 28 * 28 for testing/validation.\n",
    "\n",
    "We will use the ```gluon.data.vision``` package to directly import the Fashion-MNIST dataset and perform pre-processing on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a39bb33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the training data\n",
    "fashion_mnist_train = gluon.data.vision.FashionMNIST(train=True)\n",
    "\n",
    "# Get the validation data\n",
    "fashion_mnist_val = gluon.data.vision.FashionMNIST(train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e46c7d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms = [gluon.data.vision.transforms.Resize(224), # We pick 224 as the model we use takes an input of size 224.\n",
    "                gluon.data.vision.transforms.ToTensor()]\n",
    "\n",
    "# Now we will stack all these together.\n",
    "transforms = gluon.data.vision.transforms.Compose(transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ecca89",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply the transformations\n",
    "fashion_mnist_train = fashion_mnist_train.transform_first(transforms)\n",
    "fashion_mnist_val = fashion_mnist_val.transform_first(transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78ff279e",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256 # Batch size of the images\n",
    "num_workers = 4 # The number of parallel workers for loading the data using Data Loaders.\n",
    "\n",
    "train_data_loader = gluon.data.DataLoader(fashion_mnist_train, batch_size=batch_size,\n",
    "                                          shuffle=True, num_workers=num_workers)\n",
    "val_data_loader = gluon.data.DataLoader(fashion_mnist_val, batch_size=batch_size,\n",
    "                                        shuffle=False, num_workers=num_workers)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e202285e",
   "metadata": {},
   "source": [
    "## Model and Optimizers\n",
    "\n",
    "Let's load the resnet-18 model architecture from [Gluon Model Zoo](../../../../api/gluon/model_zoo/index.rst) and initialize its parameters. The Gluon Model Zoo contains a repository of pre-trained models as well the model architecture definitions. We are using the model architecture from the model zoo in order to train it from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eec06d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_18_v1 = vision.resnet18_v1(pretrained=False, classes = 10)\n",
    "resnet_18_v1.initialize(init = mx.init.Xavier(), device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea75a08f",
   "metadata": {},
   "source": [
    "We will be using `SoftmaxCrossEntropyLoss` as the loss function since this is a multi-class classification problem. We will be using `sgd` (Stochastic Gradient Descent) as the optimizer.\n",
    "You can experiment with a [different loss](../../../../api/gluon/loss/index.rst) or [optimizer](../../../../api/optimizer/index.rst) as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a4a3fd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50b6595e",
   "metadata": {},
   "source": [
    "Let's define the trainer object for training the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b2b3f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.04 # You can experiment with your own learning rate here\n",
    "num_epochs = 2 # You can run training for more epochs\n",
    "trainer = gluon.Trainer(resnet_18_v1.collect_params(),\n",
    "                        'sgd', {'learning_rate': learning_rate})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "008c6ae5",
   "metadata": {},
   "source": [
    "## Train using Fit API\n",
    "\n",
    "As stated earlier, the Fit API greatly simplifies the boiler plate code and complexity for training using MXNet Gluon.\n",
    "\n",
    "In the basic usage example, with just 2 lines of code, we will set up our model for training.\n",
    "\n",
    "### Basic Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c062cd6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_acc = mx.gluon.metric.Accuracy() # Metric to monitor\n",
    "\n",
    "# Define the estimator, by passing to it the model, loss function, metrics, trainer object and device\n",
    "est = estimator.Estimator(net=resnet_18_v1,\n",
    "                          loss=loss_fn,\n",
    "                          train_metrics=train_acc,\n",
    "                          trainer=trainer,\n",
    "                          device=device)\n",
    "\n",
    "# ignore warnings for nightly test on CI only\n",
    "import warnings\n",
    "with warnings.catch_warnings():\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "    # Magic line\n",
    "    est.fit(train_data=train_data_loader,\n",
    "        epochs=num_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc187984",
   "metadata": {},
   "source": [
    "```text\n",
    "    Training begin: using optimizer SGD with current learning rate 0.0400 <!--notebook-skip-line-->\n",
    "    Train for 2 epochs. <!--notebook-skip-line-->\n",
    "\n",
    "    [Epoch 0] finished in 25.110s: train_accuracy : 0.7877 train_softmaxcrossentropyloss0 : 0.5905 <!--notebook-skip-line-->\n",
    "\n",
    "    [Epoch 1] finished in 23.595s: train_accuracy : 0.8823 train_softmaxcrossentropyloss0 : 0.3197 <!--notebook-skip-line-->\n",
    "    Train finished using total 48s at epoch 1. train_accuracy : 0.8823 train_softmaxcrossentropyloss0 : 0.3197 <!--notebook-skip-line-->\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f51a3e5f",
   "metadata": {},
   "source": [
    "### Advanced Usage\n",
    "\n",
    "The Fit API is also customizable with several `Event Handlers` which give a fine grained control over the steps in training and exposes callback methods that provide control over the stages involved in training. Available callback methods are: `train_begin`, `train_end`, `batch_begin`, `batch_end`, `epoch_begin` and `epoch_end`.\n",
    "\n",
    "You can use built-in event handlers such as `LoggingHandler`, `CheckpointHandler` or `EarlyStoppingHandler` to log and save the model at certain time-steps during training. You can also stop the training when the model's performance plateaus.\n",
    "There are also some default utility handlers that will be added to your estimator by default. For example, `StoppingHandler` is used to control when the training ends, based on number of epochs or number of batches trained.\n",
    "`MetricHandler` is used to calculate training metrics at end of each batch and epoch.\n",
    "`ValidationHandler` is used to validate your model on test data at each epoch's end and then calculate validation metrics.\n",
    "You can create these utility handlers with different configurations and pass to estimator. This will override the default handler configuration.\n",
    "You can create a custom handler by inheriting one or multiple\n",
    "[base event handlers](https://github.com/apache/incubator-mxnet/blob/master/python/mxnet/gluon/contrib/estimator/event_handler.py#L32)\n",
    " including: `TrainBegin`, `TrainEnd`, `EpochBegin`, `EpochEnd`, `BatchBegin`, `BatchEnd`.\n",
    "\n",
    "\n",
    "### Custom Event Handler\n",
    "\n",
    "Here we will showcase an example custom event handler that inherits features from a few base handler classes.\n",
    "Our custom event handler is a simple one: record the loss values at the end of every epoch in our training phase.\n",
    "\n",
    "Note: For each of the method, the `Estimator` object is passed along, so you can access training metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e99e915a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LossRecordHandler(TrainBegin, TrainEnd, EpochEnd):\n",
    "    def __init__(self):\n",
    "        super(LossRecordHandler, self).__init__()\n",
    "        self.loss_history = {}\n",
    "\n",
    "    def train_begin(self, estimator, *args, **kwargs):\n",
    "        print(\"Training begin\")\n",
    "\n",
    "    def train_end(self, estimator, *args, **kwargs):\n",
    "        # Print all the losses at the end of training\n",
    "        print(\"Training ended\")\n",
    "        for loss_name in self.loss_history:\n",
    "            for i, loss_val in enumerate(self.loss_history[loss_name]):\n",
    "                print(\"Epoch: {}, Loss name: {}, Loss value: {}\".format(i, loss_name, loss_val))\n",
    "\n",
    "    def epoch_end(self, estimator, *args, **kwargs):\n",
    "        for metric in estimator.train_metrics:\n",
    "            # look for train Loss in training metrics\n",
    "            # we wrapped loss value as a metric to record it\n",
    "            if isinstance(metric, mx.gluon.metric.Loss):\n",
    "                loss_name, loss_val = metric.get()\n",
    "                # append loss value for this epoch\n",
    "                self.loss_history.setdefault(loss_name, []).append(loss_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "216fea5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's reset the model, trainer and accuracy objects from above\n",
    "\n",
    "resnet_18_v1.initialize(force_reinit=True, init = mx.init.Xavier(), device=device)\n",
    "trainer = gluon.Trainer(resnet_18_v1.collect_params(),\n",
    "                        'sgd', {'learning_rate': learning_rate})\n",
    "train_acc = mx.gluon.metric.Accuracy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15113413",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the estimator, by passing to it the model, loss function, metrics, trainer object and device\n",
    "est = estimator.Estimator(net=resnet_18_v1,\n",
    "                          loss=loss_fn,\n",
    "                          train_metrics=train_acc,\n",
    "                          trainer=trainer,\n",
    "                          device=device)\n",
    "\n",
    "# Define the handlers, let's say in built Checkpointhandler\n",
    "checkpoint_handler = CheckpointHandler(model_dir='./',\n",
    "                                       model_prefix='my_model',\n",
    "                                       monitor=train_acc,  # Monitors a metric\n",
    "                                       save_best=True)  # Save the best model in terms of\n",
    "# Let's instantiate another handler which we defined above\n",
    "loss_record_handler = LossRecordHandler()\n",
    "# ignore warnings for nightly test on CI only\n",
    "import warnings\n",
    "with warnings.catch_warnings():\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "    # Magic line\n",
    "    est.fit(train_data=train_data_loader,\n",
    "            val_data=val_data_loader,\n",
    "            epochs=num_epochs,\n",
    "            event_handlers=[checkpoint_handler, loss_record_handler]) # Add the event handlers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c52cfa75",
   "metadata": {},
   "source": [
    "```text\n",
    "    Training begin: using optimizer SGD with current learning rate 0.0400 <!--notebook-skip-line-->\n",
    "    Train for 2 epochs. <!--notebook-skip-line-->\n",
    "\n",
    "    [Epoch 0] finished in 25.236s: train_accuracy : 0.7917 train_softmaxcrossentropyloss0 : 0.5741 val_accuracy : 0.6612 val_softmaxcrossentropyloss0 : 0.8627 <!--notebook-skip-line-->\n",
    "\n",
    "    [Epoch 1] finished in 24.892s: train_accuracy : 0.8826 train_softmaxcrossentropyloss0 : 0.3229 val_accuracy : 0.8474 val_softmaxcrossentropyloss0 : 0.4262 <!--notebook-skip-line-->\n",
    "\n",
    "    Train finished using total 50s at epoch 1. train_accuracy : 0.8826 train_softmaxcrossentropyloss0 : 0.3229 val_accuracy : 0.8474 val_softmaxcrossentropyloss0 : 0.4262 <!--notebook-skip-line-->\n",
    "\n",
    "    Training begin <!--notebook-skip-line-->\n",
    "    Epoch 1, loss 0.5741 <!--notebook-skip-line-->\n",
    "    Epoch 2, loss 0.3229 <!--notebook-skip-line-->\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c57b9451",
   "metadata": {},
   "source": [
    "You can load the saved model, by using the `load_parameters` API in Gluon. For more details refer to the [Loading model parameters from file tutorial](../blocks/save_load_params.ipynb#Loading-model-parameters-from-file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9bd0dde",
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet_18_v1 = vision.resnet18_v1(pretrained=False, classes=10)\n",
    "resnet_18_v1.load_parameters('./my_model-best.params', device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a061c2bf",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "- For more hands on learning about deep learning, check out [Dive into Deep Learning](https://d2l.ai)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}