{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. 第二章 - 如何训练一个神经网络（简单版）\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "目录：\n",
    "\n",
    "第一部分 神经网络\n",
    "* 2.0 神经网络简介  \n",
    "\n",
    "第二部分 线性回归\n",
    "* 2.1 加载数据\n",
    "* 2.2 定义模型\n",
    "* 2.3 定义损失函数\n",
    "* 2.4 选择优化器\n",
    "* 2.5 训练模型并验证\n",
    "* 2.6 附：可视化验证"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.0 神经网络简介"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以监督学习为例，我们训练一个神经网络，目标是：给定输入数据及其正确的标签，将输入数据传入我们的神经网络中，通过神经网络的复杂计算，得到对应的输出值，使这个输出值尽可能地与正确标签相符。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "为了实现上述目标，我们以这样一种简单的方式理解神经网络的训练流程：  \n",
    "* 首先，我们拿到含有正确标签的数据（获取数据集）；  \n",
    "* 然后，我们定义一系列含有待定参数的计算（定义神经网络，其参数就是我们的学习目标）；  \n",
    "* 接着，我们将数据按照第二步定义好的规则，计算得到对应的输出值（前向传播）；  \n",
    "* 随后，我们根据一定规则，比较输出值与正确标签的差异（计算损失函数）；  \n",
    "* 最后，我们根据输出值与正确标签的差异大小，反观我们当前计算中所使用的各个参数，对参数进行优化更新（反向传播以及参数优化）；  \n",
    "以此反复，优化我们的计算，最终寻找到该计算（神经网络）的最佳参数。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对应地，在 Jittor (计图) 中，我们可以按如下顺序训练一个神经网络：\n",
    "1. 加载数据\n",
    "2. 定义模型\n",
    "3. 定义损失函数\n",
    "4. 选择优化器\n",
    "5. 训练模型（并验证）\n",
    "\n",
    "现在，我们按照上述流程，实现一个简单的神经网络吧！（线性回归问题）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;2m[i 0202 22:58:27.866990 76 compiler.py:847] Jittor(1.2.2.27) src: /home/llt/.local/lib/python3.7/site-packages/jittor\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:27.869186 76 compiler.py:848] g++ at /usr/bin/g++\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:27.870261 76 compiler.py:849] cache_path: /home/llt/.cache/jittor/default/g++\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:27.886880 76 __init__.py:257] Found /usr/local/cuda/bin/nvcc(10.2.89) at /usr/local/cuda/bin/nvcc.\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:27.959711 76 __init__.py:257] Found gdb(8.1.0) at /usr/bin/gdb.\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:27.976226 76 __init__.py:257] Found addr2line(2.30) at /usr/bin/addr2line.\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:28.026845 76 compiler.py:889] pybind_include: -I/usr/include/python3.7m -I/usr/local/lib/python3.7/dist-packages/pybind11/include\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:28.051528 76 compiler.py:891] extension_suffix: .cpython-37m-x86_64-linux-gnu.so\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:28.245442 76 __init__.py:169] Total mem: 62.78GB, using 16 procs for compiling.\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:30.601735 76 jit_compiler.cc:21] Load cc_path: /usr/bin/g++\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:30.801683 76 init.cc:54] Found cuda archs: [75,]\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:30.994661 76 __init__.py:257] Found mpicc(2.1.1) at /usr/bin/mpicc.\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:31.058753 76 compiler.py:654] handle pyjt_include/home/llt/.local/lib/python3.7/site-packages/jittor/extern/mpi/inc/mpi_warper.h\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:31.108971 76 compile_extern.py:287] Downloading nccl...\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:31.187842 76 compile_extern.py:16] found /usr/local/cuda/include/cublas.h\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:31.189916 76 compile_extern.py:16] found /usr/lib/x86_64-linux-gnu/libcublas.so\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:31.780064 76 compile_extern.py:16] found /usr/local/cuda/include/cudnn.h\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:31.781034 76 compile_extern.py:16] found /usr/local/cuda/lib64/libcudnn.so\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:31.809573 76 compiler.py:654] handle pyjt_include/home/llt/.local/lib/python3.7/site-packages/jittor/extern/cuda/cudnn/inc/cudnn_warper.h\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:33.621312 76 compile_extern.py:16] found /usr/local/cuda/include/curand.h\u001b[m\n",
      "\u001b[38;5;2m[i 0202 22:58:33.622719 76 compile_extern.py:16] found /usr/local/cuda/lib64/libcurand.so\u001b[m\n"
     ]
    }
   ],
   "source": [
    "# 加载计图\n",
    "import jittor as jt\n",
    "\n",
    "# 开启 GPU 加速\n",
    "# jt.flags.use_cuda = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 任务：线性回归问题\n",
    "\n",
    "任务描述如下：  \n",
    "* 已知 x 和 y 具有一定的线性关系。\n",
    "* 给定 x 的值，用模型预测 y 的值。 \n",
    "\n",
    "解决步骤如下：\n",
    "* 首先，我们会随机生成一些具有线性关系的数据点 $(x, y)$，当作我们的数据集；  \n",
    "* 然后，定义我们的计算模型为 $y=a + b \\cdot x $；  \n",
    "* 接着，我们用计图的内置函数，选择我们的损失函数和参数优化器；  \n",
    "* 最后，我们将完成训练模型的主代码块以及验证部分的代码。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 加载数据"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先，我们要准备好实验的数据集。 \n",
    "\n",
    "在这个线性回归问题中，我们会随机生成 100 个数据点 $(x, y)$。其中，x 和 y 潜在的线性关系为 $y=a + b \\cdot x $ （我们会为 $y$ 设置一定的噪音 ）。这里，$a$ 和 $b$ 为我们模型将要学习的参数。**我们先手动设置 $a = 1$，$b = 2$ 来生成数据集。然后，我们用这样一个数据集训练我们的模型，看看模型是否有能力学习到这两个参数值。**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# 设定种子，保持结果的可复制性。\n",
    "np.random.seed(2021)  \n",
    "\n",
    "# 初始化 100 个数据点。其中，x 为输入数据，y 为正确的标签值（通过 x 预测的值）。\n",
    "x = np.random.rand(100).reshape(100,1)\n",
    "y = 1 + 2 * x + 0.1 * np.random.randn(100,1)             # 线性关系设置为 a = 1, b = 2,并利用正态分布，给 y 的值设定噪音\n",
    "\n",
    "# 将我们的数据点，切分为训练集和验证集（先随机切分索引，再根据索引切分数据集）\n",
    "index = np.arange(100)                                   # 生成 100 个索引值\n",
    "np.random.shuffle(index)                                 # 将索引随机排序\n",
    "train_index = index[:80]                                 # 训练集数据对应的索引\n",
    "val_index = index[-20:]                                  # 验证集数据对应的索引\n",
    "\n",
    "# 根据索引，切分训练集和验证集\n",
    "x_train, y_train = x[train_index], y[train_index]\n",
    "x_val, y_val = x[val_index], y[val_index]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "将数据集准备好后，我们通过 Matplotlib 查看一下数据点的分布"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x7f31b3c92d50>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12804 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, axs = plt.subplots(nrows = 1, ncols = 2)\n",
    "axs[0].scatter(x_train, y_train)                         # 展示训练集上的数据点\n",
    "axs[1].scatter(x_val, y_val)                             # 展示验证集上的数据点"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，我们将数据点加载到 Jittor 中，转化为 Jittor 可操作的 Var 类型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 将 NumPy 数组转化成 Var\n",
    "x_train_var = jt.array(x_train)\n",
    "y_train_var = jt.array(y_train)\n",
    "x_val_var = jt.array(x_val)\n",
    "y_val_var = jt.array(y_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 定义模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "模型的定义：我们定义模型需要继承 Jittor 的 Module 类，并实现  \\_\\_init__  函数和 execute 函数。\n",
    "* \\_\\_init__ 函数： 用于定义模型由哪些参数或操作组成； \n",
    "* execute 函数： 定义了模型执行的顺序和模型的返回值。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "模型 1\n",
    "\"\"\"\n",
    "from jittor import Module\n",
    "\n",
    "class FirstModel(Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        # 随机初始化参数 a 和 b\n",
    "        self.a = jt.rand(1)\n",
    "        self.b = jt.rand(1) \n",
    "\n",
    "    def execute(self, x):\n",
    "        # 模型通过输入的 x 值，进行与参数 a 和参数 b 的计算，得到预测的 y 值，并返回计算结果\n",
    "        y_pred = self.a + self.b * x\n",
    "        return y_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，用我们定义好的模型类，创建一个模型实例。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FirstModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在，我们来瞧一眼这个实例模型，它初始化时随机的参数是多少："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'a': jt.Var([0.2822725], dtype=float32), 'b': jt.Var([0.15348685], dtype=float32)}\n"
     ]
    }
   ],
   "source": [
    "print(model.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 定义损失函数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们从 Jittor 的函数库里选择 MSE （均方误差）作为衡量 “模型输出值” 与 “正确标签” 差异大小的标准，即损失函数。  \n",
    "\n",
    "（提示： Jittor 内置的损失函数和优化器都在 nn 类中。您只需导入 nn 类，即可轻松地使用这些函数。）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入 nn 类\n",
    "from jittor import nn\n",
    "\n",
    "# 设置损失函数\n",
    "loss_function = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 选择优化器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们选择 Jittor 内置的 SGD （Stochastic Gradient Descent，随机梯度下降） 作为模型参数的优化器，并设置学习率 $learning\\_rate = 0.1$。  \n",
    "\n",
    "注意：在创建优化器实例的时候，我们需要将模型参数传入，代表我们的优化器将对这些参数进行优化更新。\n",
    "\n",
    "（提示：模型的参数可通过 model.parameters() 获取。）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置学习率\n",
    "learning_rate = 0.1\n",
    "\n",
    "# 传入模型参数，创建优化器实例\n",
    "optimizer = nn.SGD(model.parameters(), learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.5 训练模型并验证"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先，我们完成模型训练的代码块："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, x_train, y_train, loss_function, optimizer):\n",
    "    model.train()                                 # 开启训练模式\n",
    "    y_pred = model(x_train)                       # 将输入值 x_train 传入模型，计算得到输出值（即预测值） y_pred\n",
    "    loss = loss_function(y_train, y_pred)         # 通过损失函数，计算真实值 y_train 和预测值 y_pred 的差异大小\n",
    "    optimizer.step(loss)                          # 优化器根据计算出来的损失函数值对模型参数进行优化、更新\n",
    "    return loss                                  # 返回本次训练的 Loss 值，以便记录"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "随后，我们完成模型验证的代码块："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def val(model, x_val, y_val, loss_function):\n",
    "    model.eval()                                  # 开启验证模式，不更新模型参数\n",
    "    y_pred = model(x_val)                         # 将输入值 x_val 传入模型，计算得到输出值（即预测值） y_pred\n",
    "    loss = loss_function(y_val, y_pred)           # 通过损失函数，计算真实值 y_val 和预测值 y_pred 的差异大小     \n",
    "    return loss                                  # 返回本次验证的 Loss 值，以便记录"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，我们完成模型训练并验证的主代码块：  \n",
    "在执行前后，我们会分别打印出模型的参数，看我们的训练是否将模型参数训练成我们预期的 $a = 1$，$b = 2$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before training: \n",
      " {'a': jt.Var([0.2822725], dtype=float32), 'b': jt.Var([0.15348685], dtype=float32)}\n",
      "After training: \n",
      " {'a': jt.Var([0.9964309], dtype=float64), 'b': jt.Var([2.01100074], dtype=float64)}\n"
     ]
    }
   ],
   "source": [
    "# 打印训练前的模型参数\n",
    "print(\"Before training: \\n\", model.state_dict())\n",
    "\n",
    "# 设置迭代次数（在这个案例中，一个纪元（Epoch）即是一个迭代（Iteration））\n",
    "epochs = 500\n",
    "\n",
    "# 初始化空列表，分别用于记录训练集和验证集上的 Loss 值\n",
    "train_loss_list = list()\n",
    "val_loss_list = list()\n",
    "\n",
    "# 循环迭代训练\n",
    "for epoch in range(epochs):\n",
    "    # 在训练集上进行训练，将更新模型参数。\n",
    "    train_loss = train(model, x_train_var, y_train_var, loss_function, optimizer)\n",
    "    train_loss_list.append(train_loss)\n",
    "    # 在验证集上进行验证，模型参数不做更新。\n",
    "    val_loss = val(model, x_val_var, y_val_var, loss_function)\n",
    "    val_loss_list.append(val_loss)\n",
    "    \n",
    "# 打印训练结束后的模型参数\n",
    "print(\"After training: \\n\", model.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "从上述打印结果，可以看到，经过我们的训练，模型参数由最初的随机值朝着正确的方向发生了改变。  \n",
    "在训练结束后，我们的模型参数已明显的接近于 $a = 1$，$b = 2$。  \n",
    "这说明，我们的模型训练是成功的！"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.6 附：可视化验证"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，让我们利用可视化工具，验证一下实验结果吧。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Loss 在训练集和验证集上不同的下降趋势：**  \n",
    "* Loss 值越大，代表通过模型计算出的预测值 y_pred 和真实值 y 的差距越大；  \n",
    "* Loss 值越小，说明 y_pred 和 y 越来越接近，代表模型预测得越来越准确。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13584 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 作图前 50 个 Epoch 中 Loss 的变化趋势\n",
    "plt.plot(train_loss_list[:50],'r', label=\"Training Loss\")\n",
    "plt.plot(val_loss_list[:50],'g', label=\"Validation Loss\")\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.ylabel(\"Loss\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**比较验证集上 “原始数据点” 和 “模型预测直线” ：**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 14804 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 利用训练好的模型，对验证集上的 x 进行计算，预测出 y_pred\n",
    "y_pred = model(x_val_var)  \n",
    "\n",
    "# Matplotlib 作图。注意，需将计图的 Var 类型转化为 NumPy 数组\n",
    "plt.scatter(x_val_var.numpy(), y_val_var.numpy(), label=\"Validation Data\")      # 原始数据点\n",
    "plt.plot(x_val_var.numpy(), y_pred.numpy(), 'r', label=\"Model Prediction\")      # 模型预测结果\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 📣\n",
    "恭喜您！已成功完成了线性回归的任务！🎉🎉🎉\n",
    "\n",
    "您可能觉得这个模型还太过简单、过于理论，无法应用到实际的神经网络训练中。  \n",
    "那么，请您继续最终章的挑战。  \n",
    "在终章中，我们会以上述模型为雏形，建立一个健全的神经网络，解决一个实际的分类问题。"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
