jupyter/train_your_first_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train your first model\n",
    "\n",
    "\n",
    "In this tutorial, you 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 Kernel. To install the Java Kernel, see the [README](https:/awslabs/djl/blob/master/jupyter/README.md)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%maven ai.djl:api:0.2.0\n",
    "%maven ai.djl:basicdataset:0.2.0\n",
    "%maven ai.djl:model-zoo:0.2.0\n",
    "%maven ai.djl:repository:0.2.0\n",
    "%maven ai.djl.mxnet:mxnet-engine:0.2.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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Include MXNet engine dependency\n",
    "\n",
    "This tutorial uses MXNet engine as its backend. MXNet has different [build flavor](https://mxnet.apache.org/get_started?version=v1.5.1&platform=linux&language=python&environ=pip&processor=cpu) and it is platform specific.\n",
    "Please read [here](https:/awslabs/djl/blob/master/examples/README.md#engine-selection) for how to select MXNet engine flavor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "String classifier = System.getProperty(\"os.name\").startsWith(\"Mac\") ? \"osx-x86_64\" : \"linux-x86_64\";\n",
    "\n",
    "%maven ai.djl.mxnet:mxnet-native-mkl:jar:${classifier}:1.6.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.metrics.*;\n",
    "import ai.djl.training.optimizer.*;\n",
    "import ai.djl.training.optimizer.learningrate.*;\n",
    "import ai.djl.training.util.*;\n",
    "import ai.djl.zoo.cv.classification.*;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Create your neural network\n",
    "\n",
    "In this tutoral, use the built-in MLP block from Model Zoo. To learn more about MLP block, see [Create Your First Network](create_your_first_network.ipynb)\n",
    "\n",
    "Images in MNIST dataset are 28x28 grayscale images, so create an MLP block with 28 x 28 input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Model model = Model.newInstance();\n",
    "model.setBlock(new Mlp(28, 28));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 2: Setup your training configurations\n",
    "\n",
    "\n",
    "The following are a few common items you need to configure your training:\n",
    "* batch size: To take the advantage of GPU, you usually train models in batches. Pick a batch size based on your model.\n",
    "* [`Initializer`](https://javadoc.djl.ai/api/0.2.0/index.html?ai/djl/training/initializer/Initializer.html): An `Initializer` initializes model parameters.\n",
    "* [`Loss`](https://javadoc.djl.ai/api/0.2.0/index.html?ai/djl/training/loss/Loss.html) function: A loss function is used to measure how good (or bad) our models are. Because the lower value of the function is better, it's called the \"loss\" function.\n",
    "* [`Optimizer`](https://javadoc.djl.ai/api/0.2.0/index.html?ai/djl/training/optimizer/Optimizer.html): Optimization algorithms allow us to continue updating model parameters and to minimize the value of the loss function.\n",
    "* `Device`: 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int batchSize = 32;\n",
    "\n",
    "Initializer initializer = new XavierInitializer();\n",
    "Loss loss = Loss.softmaxCrossEntropyLoss();\n",
    "Optimizer optimizer = Optimizer.adam()\n",
    "        .setRescaleGrad(1.0f / batchSize)\n",
    "        .build();\n",
    "Accuracy accuracy = new Accuracy();\n",
    "\n",
    "TrainingConfig config = new DefaultTrainingConfig(initializer, loss)\n",
    "        .setOptimizer(optimizer)\n",
    "        .addTrainingMetric(accuracy)\n",
    "        .setBatchSize(batchSize);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Prepare MNIST dataset for training\n",
    "\n",
    "The [MNIST](https://en.wikipedia.org/wiki/MNIST_database) dataset is a database of handwritten digits that is commonly used for training image classification models. \n",
    "\n",
    "We provide the MNIST dataset as a built-in dataset, which makes it easy for you to use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Mnist mnist = Mnist.builder(model.getNDManager()).setSampling(batchSize, true).build();\n",
    "mnist.prepare(new ProgressBar());"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 4: Create a Trainer\n",
    "\n",
    "Now, you can create a `Trainer` with your training configuration.\n",
    "You need to initialize the model parameters in the trainer with proper input shape:\n",
    "* The first axis of the input is batch, which won't impact the parameter, so you can use 1 here.\n",
    "* The second axis of the MNIST image, which is 28 * 28.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Trainer trainer = model.newTrainer(config);\n",
    "trainer.initialize(new Shape(1, 28 * 28));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 5: Train your model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ProgressBar progressBar = new ProgressBar(\"Training\", (int)(mnist.size() / batchSize));\n",
    "int epoch = 2;\n",
    "float trainingAccuracy = 0f;\n",
    "for (int i = 0; i < epoch; ++i) {\n",
    "    int index = 0;\n",
    "    for (Batch batch : trainer.iterateDataset(mnist)) {\n",
    "        trainer.trainBatch(batch);\n",
    "        trainer.step();\n",
    "        batch.close();\n",
    "\n",
    "        trainingAccuracy = accuracy.getValue();\n",
    "        progressBar.update(index++, String.format(\"Epoch: %d, Accuracy: %.3f\", i, trainingAccuracy));\n",
    "    }\n",
    "    // reset training and validation metric at end of epoch\n",
    "    trainer.resetTrainingMetrics();\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 6: Save your model\n",
    "\n",
    "While saving your model, you can add metadata to it such as: training accuracy, number epoch, etc."
   ]
  },
  {
   "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",
    "model.setProperty(\"Accuracy\", String.valueOf(trainingAccuracy));\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 in [examples project](https:/awslabs/djl/blob/master/examples/docs/train_your_first_model.md)."
   ]
  }
 ],
 "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": "11.0.5+10-LTS"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}