{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fe5c93a5-a182-4b32-9831-5ca0c391def6",
   "metadata": {},
   "source": [
    "# 梯度累积\n",
    "\n",
    "[![在线运行](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0/resource/_static/logo_modelarts.png)](https://authoring-modelarts-cnnorth4.huaweicloud.com/console/lab?share-url-b64=aHR0cHM6Ly9vYnMuZHVhbHN0YWNrLmNuLW5vcnRoLTQubXlodWF3ZWljbG91ZC5jb20vbWluZHNwb3JlLXdlYnNpdGUvbm90ZWJvb2svcjIuMC90dXRvcmlhbHMvZXhwZXJ0cy96aF9jbi9vcHRpbWl6ZS9taW5kc3BvcmVfZ3JhZGllbnRfYWNjdW11bGF0aW9uLmlweW5i=&imageid=b8671c1e-c439-4ae2-b9c6-69b46db134ae) [![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0/resource/_static/logo_notebook.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r2.0/tutorials/experts/zh_cn/optimize/mindspore_gradient_accumulation.ipynb) [![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0/resource/_static/logo_download_code.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/r2.0/tutorials/experts/zh_cn/optimize/mindspore_gradient_accumulation.py) [![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0/resource/_static/logo_source.png)](https://gitee.com/mindspore/docs/blob/r2.0/tutorials/experts/source_zh_cn/optimize/gradient_accumulation.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cef28b35-ab1f-4051-a3a3-d8cf4af1ca45",
   "metadata": {},
   "source": [
    "## 概述\n",
    "\n",
    "本教程介绍梯度累积的训练算法，目的是为了解决由于内存不足，导致Batch size过大神经网络无法训练，或者网络模型过大无法加载的OOM（Out Of Memory）问题。\n",
    "\n",
    "## 梯度累积原理\n",
    "\n",
    "梯度累积是一种将训练神经网络的数据样本按Batch size拆分为几个小Batch的方式，然后按顺序进行计算。\n",
    "\n",
    "在进一步讨论梯度累积之前，我们来看看神经网络的计算过程。\n",
    "\n",
    "深度学习模型由许多相互连接的神经网络单元所组成，在所有神经网络层中，样本数据会不断向前传播。在通过所有层后，网络模型会输出样本的预测值，通过损失函数然后计算每个样本的损失值（误差）。神经网络通过反向传播，去计算损失值相对于模型参数的梯度。最后这些梯度信息用于对网络模型中的参数进行更新。\n",
    "\n",
    "优化器用于对网络模型权重参数更新的数学公式。以一个简单随机梯度下降(SGD)算法为例。\n",
    "\n",
    "假设Loss Function函数公式为：\n",
    "\n",
    "$$Loss(\\theta)=\\frac{1}{2}\\left(h(x^{k})-y^{k}\\right)^{2}$$\n",
    "\n",
    "在构建模型时，优化器用于计算最小化损失的算法。这里SGD算法利用Loss函数来更新权重参数公式为：\n",
    "\n",
    "$$\\theta_{i}=\\theta_{i-1}-lr * grad_{i}$$\n",
    "\n",
    "其中$\\theta$是网络模型中的可训练参数（权重或偏差），$lr$是学习率，$grad_{i}$是相对于网络模型参数的损失。\n",
    "\n",
    "梯度累积只计算神经网络模型，并不及时更新网络模型的参数，同时在计算的时候累积得到的梯度信息，最后统一使用累积的梯度来对参数进行更新。\n",
    "\n",
    "$$accumulated=\\sum_{i=0}^{N} grad_{i}$$\n",
    "\n",
    "在不更新模型变量的时候，实际上是把原来的数据Batch size分成几个小的Mini-Batch，每个step中使用的样本实际上是更小的数据集。\n",
    "\n",
    "在N个step内不更新变量，使所有Mini-Batch使用相同的模型变量来计算梯度，以确保计算出来得到相同的梯度和权重信息，算法上等价于使用原来没有切分的Batch size大小一样。即：\n",
    "\n",
    "$$\\theta_{i}=\\theta_{i-1}-lr * \\sum_{i=0}^{N} grad_{i}$$\n",
    "\n",
    "最终在上面步骤中累积梯度会产生与使用全局Batch size大小相同的梯度总和。\n",
    "\n",
    "![](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.0/tutorials/experts/source_zh_cn/optimize/images/GradientAccumulation1.png)\n",
    "\n",
    "当然在实际工程当中，关于调参和算法上有两点需要注意的：\n",
    "\n",
    "1. **学习率 learning rate**：一定条件下，Batch size越大训练效果越好，梯度累积则模拟了Batch size增大的效果，如果accumulation steps为4，则Batch size增大了4倍，根据经验，使用梯度累积的时候需要把学习率适当放大。\n",
    "\n",
    "2. **归一化 Batch Norm**：accumulation steps为4时进行Batch size模拟放大的效果，与真实Batch size相比，数据的分布其实并不完全相同，4倍Batch size的Batch Norm计算出来的均值和方差与实际数据均值和方差不太相同，因此有些实现中会使用Group Norm来代替Batch Norm。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0fcd2e0-8534-4a8b-9ea6-424c6ee5b741",
   "metadata": {},
   "source": [
    "## 梯度累积实现\n",
    "\n",
    "基于MindSpore的[函数式自动微分](https://www.mindspore.cn/tutorials/zh-CN/r2.0/beginner/autograd.html)机制，正向和反向执行完成后，函数将返回与训练参数相对应的梯度。因此我们需要设计一个梯度累积类Accumulator，对每一个Step产生的梯度值进行存储。下面是Accumulator的实现样例，我们需要维护两份与模型可训练参数的Shape相同的内部属性，分别为inner_grads和zeros。其中inner_grads用于存储累积的梯度值，zeros用于参数优化更新后的清零。同时，Accumulator内部维护了一个counter变量，在每一次正反向执行完成后，counter自增，通过对counter取模的方式来判断是否达到累积步数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "02712382-422d-42a9-9862-3028c272d7f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore as ms\n",
    "from mindspore import Tensor, Parameter, ops\n",
    "\n",
    "@ms.jit_class\n",
    "class Accumulator():\n",
    "    def __init__(self, optimizer, accumulate_step, clip_norm=1.0):\n",
    "        self.optimizer = optimizer\n",
    "        self.clip_norm = clip_norm\n",
    "        self.inner_grads = optimizer.parameters.clone(prefix=\"accumulate_\", init='zeros')\n",
    "        self.zeros = optimizer.parameters.clone(prefix=\"zeros_\", init='zeros')\n",
    "        self.counter = Parameter(Tensor(1, ms.int32), 'counter_')\n",
    "        assert accumulate_step > 0\n",
    "        self.accumulate_step = accumulate_step\n",
    "        self.map = ops.HyperMap()\n",
    "\n",
    "    def __call__(self, grads):\n",
    "        # 将单步获得的梯度累加至Accumulator的inner_grads\n",
    "        self.map(ops.partial(ops.assign_add), self.inner_grads, grads)\n",
    "        if self.counter % self.accumulate_step == 0:\n",
    "            # 如果达到累积步数，进行参数优化更新\n",
    "            self.optimizer(self.inner_grads)\n",
    "            # 完成参数优化更新后，清零inner_grads\n",
    "            self.map(ops.partial(ops.assign), self.inner_grads, self.zeros)\n",
    "        # 计算步数加一\n",
    "        ops.assign_add(self.counter, Tensor(1, ms.int32))\n",
    "\n",
    "        return True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6382173-b71e-4185-803f-53cf17cf076d",
   "metadata": {},
   "source": [
    "> `ms.jit_class`为MindSpore即时编译修饰器，可以将普通的Python类作为可编译计算图使用。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cf8b985-04b4-4271-b6bf-1e063fbbf583",
   "metadata": {},
   "source": [
    "接下来，我们使用[快速入门](https://www.mindspore.cn/tutorials/zh-CN/r2.0/beginner/quick_start.html)中手写数字识别模型验证梯度累加的效果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "73e7648e-c2aa-413f-9f88-e94e36a5e9d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/MNIST_Data.zip (10.3 MB)\n",
      "\n",
      "file_sizes: 100%|██████████████████████████| 10.8M/10.8M [00:06<00:00, 1.67MB/s]\n",
      "Extracting zip file...\n",
      "Successfully downloaded / unzipped to ./\n"
     ]
    }
   ],
   "source": [
    "from mindspore import nn\n",
    "from mindspore import value_and_grad\n",
    "from mindspore.dataset import vision, transforms\n",
    "from mindspore.dataset import MnistDataset\n",
    "\n",
    "from download import download\n",
    "\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",
    "\n",
    "def datapipe(path, batch_size):\n",
    "    image_transforms = [\n",
    "        vision.Rescale(1.0 / 255.0, 0),\n",
    "        vision.Normalize(mean=(0.1307,), std=(0.3081,)),\n",
    "        vision.HWC2CHW()\n",
    "    ]\n",
    "    label_transform = transforms.TypeCast(ms.int32)\n",
    "\n",
    "    dataset = MnistDataset(path)\n",
    "    dataset = dataset.map(image_transforms, 'image')\n",
    "    dataset = dataset.map(label_transform, 'label')\n",
    "    dataset = dataset.batch(batch_size)\n",
    "    return dataset\n",
    "\n",
    "class Network(nn.Cell):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.dense_relu_sequential = nn.SequentialCell(\n",
    "            nn.Dense(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Dense(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Dense(512, 10)\n",
    "        )\n",
    "\n",
    "    def construct(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.dense_relu_sequential(x)\n",
    "        return logits\n",
    "\n",
    "model = Network()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d97f0eb8-713d-4a52-b125-e267163dede4",
   "metadata": {},
   "source": [
    "假设我们在使用[快速入门](https://www.mindspore.cn/tutorials/zh-CN/r2.0/beginner/quick_start.html)中配置的`batch_size=64`会导致显存不足，此时我们设置累加步数为2，通过执行两次`batch_size=32`进行梯度累加。\n",
    "\n",
    "首先，我们使用Accumulator，传入实例化的optimizer，并配置累加步数。然后定义正向计算函数`forward_fn`，此时，由于梯度累加的需要，loss值需要进行相应的缩放。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1cfbe3a1-e275-4362-90cd-7f848ceece23",
   "metadata": {},
   "outputs": [],
   "source": [
    "accumulate_step = 2\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = nn.SGD(model.trainable_params(), 1e-2)\n",
    "accumulator = Accumulator(optimizer, accumulate_step)\n",
    "\n",
    "def forward_fn(data, label):\n",
    "    logits = model(data)\n",
    "    loss = loss_fn(logits, label)\n",
    "    # loss除以累加步数accumulate_step\n",
    "    return loss / accumulate_step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45cd732b-15ea-4a4f-9420-ac5389b78895",
   "metadata": {},
   "source": [
    "接下来继续使用`value_and_grad`函数进行函数变换，并构造单步训练函数`train_step`。此时我们使用实例化好的accumulator进行梯度累加，由于optimizer作为accumulator的内部属性，不需要单独执行。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "79f0d4f4-ed5d-4491-9c55-327817049243",
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_fn = value_and_grad(forward_fn, None, model.trainable_params())\n",
    "\n",
    "@ms.jit\n",
    "def train_step(data, label):\n",
    "    loss, grads = grad_fn(data, label)\n",
    "    accumulator(grads)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5fb7d77-7d64-4a4a-b35c-c385ee501a41",
   "metadata": {},
   "source": [
    "接下来，我们定义训练和评估逻辑，进行训练验证。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "86103b30-18ed-4831-8bea-3b9c9bcb51cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_loop(model, dataset, loss_fn, optimizer):\n",
    "    size = dataset.get_dataset_size()\n",
    "    model.set_train()\n",
    "    for batch, (data, label) in enumerate(dataset.create_tuple_iterator()):\n",
    "        loss = train_step(data, label)\n",
    "\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.asnumpy(), batch\n",
    "            print(f\"loss: {loss:>7f}  [{current:>3d}/{size:>3d}]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5a6993ea-c4fd-4d1a-bbb2-7bc0567162df",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_loop(model, dataset, loss_fn):\n",
    "    num_batches = dataset.get_dataset_size()\n",
    "    model.set_train(False)\n",
    "    total, test_loss, correct = 0, 0, 0\n",
    "    for data, label in dataset.create_tuple_iterator():\n",
    "        pred = model(data)\n",
    "        total += len(data)\n",
    "        test_loss += loss_fn(pred, label).asnumpy()\n",
    "        correct += (pred.argmax(1) == label).asnumpy().sum()\n",
    "    test_loss /= num_batches\n",
    "    correct /= total\n",
    "    print(f\"Test: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42ccd6d0-217f-45e2-985a-6dbc92bd3950",
   "metadata": {},
   "source": [
    "接下来同样进行3个epoch的训练，注意此时根据我们的假设，数据集需要设置`batch_size=32`，每两步进行累加。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cffeea83-ad68-4785-a8d5-4a5a7700f871",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = datapipe('MNIST_Data/train', 32)\n",
    "test_dataset = datapipe('MNIST_Data/test', 32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77494827-aba1-4800-bfd0-18e22a7a9f89",
   "metadata": {},
   "source": [
    "开始训练验证，此时由于batch_size调小需要训练的step数增加至2倍。最终Accuracy验证结果与[快速入门](https://www.mindspore.cn/tutorials/zh-CN/r2.0/beginner/quick_start.html)结果一致，均为92.0\\%左右。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "faa96153-c7f2-439d-a384-776e867d87e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 1.150851  [  0/1875]\n",
      "loss: 1.149633  [100/1875]\n",
      "loss: 1.145340  [200/1875]\n",
      "loss: 1.140591  [300/1875]\n",
      "loss: 1.134244  [400/1875]\n",
      "loss: 1.125991  [500/1875]\n",
      "loss: 1.100611  [600/1875]\n",
      "loss: 1.051961  [700/1875]\n",
      "loss: 0.925877  [800/1875]\n",
      "loss: 0.879966  [900/1875]\n",
      "loss: 0.750192  [1000/1875]\n",
      "loss: 0.617844  [1100/1875]\n",
      "loss: 0.470084  [1200/1875]\n",
      "loss: 0.560856  [1300/1875]\n",
      "loss: 0.359766  [1400/1875]\n",
      "loss: 0.502521  [1500/1875]\n",
      "loss: 0.299145  [1600/1875]\n",
      "loss: 0.383266  [1700/1875]\n",
      "loss: 0.239381  [1800/1875]\n",
      "Test: \n",
      " Accuracy: 84.8%, Avg loss: 0.528309 \n",
      "\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 0.390662  [  0/1875]\n",
      "loss: 0.250778  [100/1875]\n",
      "loss: 0.570571  [200/1875]\n",
      "loss: 0.196102  [300/1875]\n",
      "loss: 0.297634  [400/1875]\n",
      "loss: 0.192528  [500/1875]\n",
      "loss: 0.231240  [600/1875]\n",
      "loss: 0.144425  [700/1875]\n",
      "loss: 0.113696  [800/1875]\n",
      "loss: 0.233481  [900/1875]\n",
      "loss: 0.212078  [1000/1875]\n",
      "loss: 0.144562  [1100/1875]\n",
      "loss: 0.220822  [1200/1875]\n",
      "loss: 0.197890  [1300/1875]\n",
      "loss: 0.283782  [1400/1875]\n",
      "loss: 0.219684  [1500/1875]\n",
      "loss: 0.155505  [1600/1875]\n",
      "loss: 0.255665  [1700/1875]\n",
      "loss: 0.155548  [1800/1875]\n",
      "Test: \n",
      " Accuracy: 90.1%, Avg loss: 0.340294 \n",
      "\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 0.176077  [  0/1875]\n",
      "loss: 0.204260  [100/1875]\n",
      "loss: 0.339903  [200/1875]\n",
      "loss: 0.221457  [300/1875]\n",
      "loss: 0.244668  [400/1875]\n",
      "loss: 0.089163  [500/1875]\n",
      "loss: 0.159595  [600/1875]\n",
      "loss: 0.211632  [700/1875]\n",
      "loss: 0.096592  [800/1875]\n",
      "loss: 0.081018  [900/1875]\n",
      "loss: 0.190852  [1000/1875]\n",
      "loss: 0.139729  [1100/1875]\n",
      "loss: 0.049344  [1200/1875]\n",
      "loss: 0.122041  [1300/1875]\n",
      "loss: 0.198622  [1400/1875]\n",
      "loss: 0.133956  [1500/1875]\n",
      "loss: 0.144801  [1600/1875]\n",
      "loss: 0.076985  [1700/1875]\n",
      "loss: 0.103241  [1800/1875]\n",
      "Test: \n",
      " Accuracy: 92.0%, Avg loss: 0.281193 \n",
      "\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "epochs = 3\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train_loop(model, train_dataset, loss_fn, optimizer)\n",
    "    test_loop(model, test_dataset, loss_fn)\n",
    "print(\"Done!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MindSpore",
   "language": "python",
   "name": "mindspore"
  },
  "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": 5
}