{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train your first model\n",
    "\n",
    "This is the second of our [beginner tutorial series](https://github.com/awslabs/djl/tree/master/jupyter/tutorial) that will take you through creating, training, and running inference on a neural network. In this tutorial, you will learn how to train an image classification model that can recognize handwritten digits.\n",
    "\n",
    "## Preparation\n",
    "\n",
    "This tutorial requires the installation of the Java Jupyter Kernel. To install the kernel, see the [Jupyter README](https://github.com/awslabs/djl/blob/master/jupyter/README.md)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "// Add the snapshot repository to get the DJL snapshot artifacts\n",
    "// %mavenRepo snapshots https://oss.sonatype.org/content/repositories/snapshots/\n",
    "\n",
    "// Add the maven dependencies\n",
    "%maven ai.djl:api:0.5.0\n",
    "%maven ai.djl:basicdataset:0.5.0\n",
    "%maven ai.djl:model-zoo:0.5.0\n",
    "%maven ai.djl.mxnet:mxnet-engine:0.5.0\n",
    "%maven org.slf4j:slf4j-api:1.7.26\n",
    "%maven org.slf4j:slf4j-simple:1.7.26\n",
    "%maven net.java.dev.jna:jna:5.3.0\n",
    "        \n",
    "// See https://github.com/awslabs/djl/blob/master/mxnet/mxnet-engine/README.md\n",
    "// for more MXNet library selection options\n",
    "%maven ai.djl.mxnet:mxnet-native-auto:1.7.0-a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import java.nio.file.*;\n",
    "\n",
    "import ai.djl.*;\n",
    "import ai.djl.basicdataset.*;\n",
    "import ai.djl.ndarray.types.*;\n",
    "import ai.djl.training.*;\n",
    "import ai.djl.training.dataset.*;\n",
    "import ai.djl.training.initializer.*;\n",
    "import ai.djl.training.loss.*;\n",
    "import ai.djl.training.listener.*;\n",
    "import ai.djl.training.evaluator.*;\n",
    "import ai.djl.training.optimizer.*;\n",
    "import ai.djl.training.util.*;\n",
    "import ai.djl.basicmodelzoo.cv.classification.*;\n",
    "import ai.djl.basicmodelzoo.basic.*;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Prepare MNIST dataset for training\n",
    "\n",
    "When training a deep learning network, it is important to first understand the dataset.\n",
    "\n",
    "## Dataset\n",
    "\n",
    "A [Dataset](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/dataset/Dataset.html) is a collection of sample input/output pairs for the function represented by your neural network. Each single input/output is represented by a [Record](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/dataset/Record.html). Each record could have multiple arrays of inputs or outputs such as an image question and answer dataset where the input is both an image and a question about the image while the output is the answer to the question.\n",
    "\n",
    "Because data learning is highly parallelizable, training is often done not with a single record at a time but a [Batch](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/dataset/Batch.html) of records at a time. This can lead to significant performance gains, especially when working with images.\n",
    "\n",
    "### MNIST\n",
    "\n",
    "The dataset we will be using is [MNIST](https://en.wikipedia.org/wiki/MNIST_database), a database of handwritten digits. Each image contains a black and white digit from 0-9 in a 28x28 image. It is commonly used when getting started with deep learning because it is small and fast to train.\n",
    "\n",
    "![Mnist Image](https://upload.wikimedia.org/wikipedia/commons/2/27/MnistExamples.png)\n",
    "\n",
    "Once you understand your dataset, you should create an implementation of the [Dataset class](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/dataset/Dataset.html). In this case, we provide the MNIST dataset built-in to make it easy for you to use it.\n",
    "\n",
    "## Sampler\n",
    "\n",
    "Then, we must decide the parameters for loading data from the dataset. The only parameter we need for MNIST is the choice of [Sampler](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/dataset/Sampler.html). The sampler decides which and how many element from datasets are part of each batch when iterating through it. We will have it randomly shuffle the elements for the batch and use a batchSize of 32. The batchSize is usually the largest power of 2 that fits within memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int batchSize = 32;\n",
    "Mnist mnist = Mnist.builder().setSampling(batchSize, true).build();\n",
    "mnist.prepare(new ProgressBar());"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 2: Create your Model\n",
    "\n",
    "A [Model](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/Model.html) contains a neural network [Block](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/nn/Block.html) along with additional artifacts used for the training process. It possesses additional information about the inputs, outputs, shapes, and data types you will use. Generally, you will use Model once you have fully completed your Block.\n",
    "\n",
    "In this tutoral, we will use the built-in Multilayer Perceptron Block from the Model Zoo. To learn more, see the previous tutorial: [Create Your First Network](create_your_first_network.ipynb).\n",
    "\n",
    "Because images in the MNIST dataset are 28x28 grayscale images, we will create an MLP block with 28 x 28 input. The output will be 10 because there are 10 possible classes (0 to 9) each image could be. For the hidden layers, we have chosen `new int[] {128, 64}` by experimenting with different values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Model model = Model.newInstance();\n",
    "model.setBlock(new Mlp(28 * 28, 10, new int[] {128, 64}));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Create a Trainer\n",
    "\n",
    "Now, you can create a [`Trainer`](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/Trainer.html) to train your model. The trainer is the main class to orchestrate the training process. Usually, they will be opened using a try-with-resources and closed after training is over.\n",
    "\n",
    "The trainer takes an existing model and attempts to optimize the parameters inside the model's Block to best match the dataset. Most optimization is based upon [Stochastic Gradient Descent](https://en.wikipedia.org/wiki/Stochastic_gradient_descent) (SGD).\n",
    "\n",
    "## Step 3.1: Setup your training configurations\n",
    "\n",
    "Before you create your trainer, we we will need a [training configuration](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/DefaultTrainingConfig.html) that describes how to train your model.\n",
    "\n",
    "The following are a few common items you may need to configure your training:\n",
    "* **REQUIRED** [`Loss`](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/loss/Loss.html) function: A loss function is used to measure how well our model matches the dataset. Because the lower value of the function is better, it's called the \"loss\" function. The Loss is the only required argument to the model\n",
    "* [`Evaluator`](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/evaluator/Evaluator.html) function: An evaluator function is also used to measure how well our model matches the dataset. Unlike the loss, they are only there for people to look at and are not used for optimizing the model. Since many losses are not as intuitive, adding other evaluators such as Accuracy can help to understand how your model is doing. If you know of any useful evaluators, we recommend adding them.\n",
    "* batch size: To take the advantage of the natural parallelism, you usually train models with batches of input data items rather than a single item at a time. This should match the batch size provided in the model\",\n",
    "* [`Device`](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/Device.html): The device is what hardware should be used to train your model on. Typically, this is either CPU or GPU. DJL can automatically detect whether a GPU is available. If GPUs are available, it will run on a single GPU by default. If you need to train with multiple GPUs, you need to set devices as : `config.setDevices(Devices.getDevices(maxNumberOfGPUs))`.\n",
    "* [`Initializer`](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/initializer/Initializer.html): An `Initializer` is used to set the initial values of the model's parameters before training. This can usually be left as the default initializer.\n",
    "* [`Optimizer`](https://javadoc.io/static/ai.djl/api/0.5.0/index.html?ai/djl/training/optimizer/Optimizer.html): The optimizer is the code that updates the model parameters to minimize the loss function. There are a variety of optimizers, most of which offer improvements upon the basic SGD. When just starting, you can use the default optimizer. Later on, Customizing the optimizer can result in faster training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DefaultTrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())\n",
    "    //softmaxCrossEntropyLoss is a standard loss for classification problems\n",
    "    .addEvaluator(new Accuracy()) // Use accuracy so we humans can understand how accurate the model is\n",
    "    .addTrainingListeners(TrainingListener.Defaults.logging());\n",
    "\n",
    "// Now that we have our training configuration, we should create a new trainer for our model\n",
    "Trainer trainer = model.newTrainer(config);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 5: Initialize Training\n",
    "\n",
    "Before training your model, you have to initialize all of the parameters with default values. You can use the trainer for this initialization by passing in the input shape.\n",
    "\n",
    "* The first axis of the input shape is the batch size. This won't impact the parameter initialization, so you can use 1 here.\n",
    "* The second axis of the input shape of the MLP - the number of pixels in the input image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.initialize(new Shape(1, 28 * 28));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 6: Train your model\n",
    "\n",
    "Now, we can train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "// Deep learning is typically trained in epochs where each epoch trains the model on each item in the dataset once.\n",
    "int epoch = 2;\n",
    "\n",
    "for (int i = 0; i < epoch; ++i) {\n",
    "    int index = 0;\n",
    "    \n",
    "    // We iterate through the dataset once during this epoch\n",
    "    for (Batch batch : trainer.iterateDataset(mnist)) {\n",
    "        \n",
    "        // During trainBatch, we update the loss and evaluators with the results for the training batch.\n",
    "        trainer.trainBatch(batch);\n",
    "        \n",
    "        // Now, we update the model parameters based on the results of the latest trainBatch\n",
    "        trainer.step();\n",
    "        \n",
    "        // We must make sure to close the batch to ensure all the memory associated with the batch is cleared quickly.\n",
    "        // If the memory isn't closed after each batch, you will very quickly run out of memory on your GPU\n",
    "        batch.close();\n",
    "    }\n",
    "    // reset training and validation evaluators at end of epoch\n",
    "    trainer.endEpoch();\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 7: Save your model\n",
    "\n",
    "Once your model is trained, you should save it so that it can be reloaded later. You can also add metadata to it such as training accuracy, number of epochs trained, etc that can be used when loading the model or when examining it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Path modelDir = Paths.get(\"build/mlp\");\n",
    "Files.createDirectories(modelDir);\n",
    "\n",
    "model.setProperty(\"Epoch\", String.valueOf(epoch));\n",
    "\n",
    "model.save(modelDir, \"mlp\");\n",
    "\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summary\n",
    "\n",
    "Now, you've successfully trained a model that can recognize handwritten digits. You'll learn how to apply this model in the next chapter: [Run image classification with your model](image_classification_with_your_model.ipynb).\n",
    "\n",
    "You can find the complete source code for this tutorial in the [examples project](https://github.com/awslabs/djl/blob/master/examples/src/main/java/ai/djl/examples/training/TrainMnist.java)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java",
   "language": "java",
   "name": "java"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "Java",
   "pygments_lexer": "java",
   "version": "12.0.2+10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}