hw14_life_long_learning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "hw14_life_long_learning.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "LD5roJkIvoRj"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Iallen520/lhy_DL_Hw/blob/master/hw14_life_long_learning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LD5roJkIvoRj",
        "colab_type": "text"
      },
      "source": [
        "## LifeLong Machine Learning\n",
        "\n",
        "### 助教的投影片連結\n",
        "[投影片](https://docs.google.com/presentation/d/13JmcOZ9i_m5xJbRBKNMAKE1fIzGhyaeLck3frY0B2xY/edit?usp=sharing)\n",
        "\n",
        "### 定義\n",
        "老師的影片有詳細說明定義 這裡不細提 詳細可以參考 [lifelong learning](https://youtu.be/7qT5P9KJnWo) \n",
        "\n",
        "\n",
        "### 方法\n",
        "在2019年底，有人提出了一個大匯整將lifelong learning 的方法，從2016- 2019 年初 的模型做了歸類，大致上可以分成三種大方法\n",
        "* Replay-based methods\n",
        "* Regularization-based methods\n",
        "* Parameter isolation methods\n",
        "\n",
        "<img src=\"https://i.ibb.co/VDFJkWG/2019-12-29-17-25.png\" width=\"100%\">\n",
        "\n",
        "在這次的作業之中，我們要走過一次regularization-based methods 裡面的 prior-focused的兩種方法 分別是 EWC 和 MAS 這兩種方法\n",
        "\n",
        "圖片出處 [Continual Learning in Neural\n",
        "Networks](https://arxiv.org/pdf/1910.02718.pdf)\n",
        "\n",
        "若有任何問題，歡迎來信至助教信箱 ntu-ml-2020spring-ta@googlegroups.com\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "ILD0GKIgJPHb",
        "colab": {}
      },
      "source": [
        "%%capture\n",
        "# !pip3 install torch torchvision"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "On1VZz4HUIJw",
        "colab_type": "text"
      },
      "source": [
        "# Import Libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dLnpJTNtje_J",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%capture\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import torch.utils.data as data\n",
        "import torch.utils.data.sampler as sampler\n",
        "import torchvision\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "import numpy as np\n",
        "import os\n",
        "import random\n",
        "from copy import deepcopy\n",
        "import json\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KRGf4QAKFzz9",
        "colab_type": "text"
      },
      "source": [
        "# 模型"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zZvlf0Wv7YdU",
        "colab_type": "text"
      },
      "source": [
        " >因為本次作業強調的是lifelong learning 的訓練方法，並非疊模型，所以今天我們所舉的例子，都會使用同一個模型來做訓練只是應用上不同lifelong learning的訓練方法， 在這次的作業的例子內 我們使用的是 一個 六層的 fully-connected layer 的 模型 加上 relu的 activation function."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fkc3Z-4JYaPe",
        "colab_type": "text"
      },
      "source": [
        "## Basic Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g8aQRs7ss3nx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Model(nn.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "    super(Model, self).__init__()\n",
        "    self.fc1 = nn.Linear(3*32*32, 1024)\n",
        "    self.fc2 = nn.Linear(1024, 512)\n",
        "    self.fc3 = nn.Linear(512, 256)\n",
        "    self.fc4 = nn.Linear(256, 128)\n",
        "    self.fc5 = nn.Linear(128, 128)\n",
        "    self.fc6 = nn.Linear(128, 10)\n",
        "    self.relu = nn.ReLU()\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = x.view(-1, 3*32*32)\n",
        "    x = self.fc1(x)\n",
        "    x = self.relu(x)\n",
        "    x = self.fc2(x)\n",
        "    x = self.relu(x)\n",
        "    x = self.fc3(x)\n",
        "    x = self.relu(x)\n",
        "    x = self.fc4(x)\n",
        "    x = self.relu(x)\n",
        "    x = self.fc5(x)\n",
        "    x = self.relu(x)\n",
        "    x = self.fc6(x)\n",
        "    return x\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AfeD8zfAccZe",
        "colab_type": "text"
      },
      "source": [
        "以下我們將依序介紹這兩種方法 EWC 跟 MAS "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qwri0rVpjd6h",
        "colab_type": "text"
      },
      "source": [
        "## EWC"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e3NePROJyWyW",
        "colab_type": "text"
      },
      "source": [
        "### Elastic Weight Consolidation\n",
        "\n",
        "#### 概念\n",
        "老師在影片中已經把核心概念介紹給大家，那在這邊我想大家都非常了解了這個方法的概念，我們就直接進入主題\n",
        "\n",
        "今天我們的任務 是在學習連續的兩個 task task A 跟 task B:\n",
        "\n",
        "在 EWC 作法下 他的 loss function 會被定義如下\n",
        " $$\\mathcal{L}_B = \\mathcal{L}(\\theta) + \\sum_{i} \\frac{\\lambda}{2} F_i (\\theta_{i} - \\theta_{A,i}^{*})^2  $$\n",
        "\n",
        "先解釋這個 loss function 裡的變數，$\\mathcal{L}_B$ 是指 task B 的 loss, 會等於 正常的loss function $\\mathcal{L}(\\theta)$ (如果 是 classification 的問題,就是 cross entropy 的 loss function) 加上一個正則項 (regularization term) \n",
        "\n",
        "這個正則項的由兩個部份組成，第一個是 $F_i$ 也是這個方法的核心, 第二個部份是 $(\\theta_{i} - \\theta_{A,i}^{*})^2$  ,  $\\theta_{A,i}^{*}$ 代表的是 訓練完task A 存下來模型第 i 個參數的值, $\\theta_i$ 代表的是目前模型第i個參數的值，注意一點是模型的架構在這種 regularization based 的方法上，都是固定ㄉ，目前模型跟 task A 存下來的模型 架構都一樣只是值不一樣。底下我將說明這個 $F_i$ 是怎麼實做出來\n",
        "\n",
        "在老師的影片中，老師是以只有兩個參數的模型舉例子，那假設我今天模型就是一個 neural network(參數不只兩個) 該怎麼辦呢？   \n",
        "\n",
        "$F_i$ 對應到老師的影片敘述是指第i個參數的守衛，假設這個參數對 task A 很重要，那這個 $F_i$ 的值就會很大，這個參數盡量不能被更動...\n",
        "\n",
        "實際上這個參數的算法 即是 如下的式子\n",
        "\n",
        "$$ F = [ \\nabla \\log(p(y_n | x_n, \\theta_{A}^{*}) \\nabla \\log(p(y_n | x_n, \\theta_{A}^{*})^T ] $$ \n",
        "\n",
        "$F$ 之中 只以對角線的值去近似各個參數的 $F_i$ 值\n",
        "\n",
        "$p(y_n | x_n, \\theta_{A}^{*})$ 指的就是模型在給定之前 task 的 data $x_n$ 以及 給定 訓練完 task A 存下來的模型參數 $\\theta_A^*$ 得到 $y_n$($x_n$ 對應的 label ) 的 posterior probability.\n",
        "那統整一下作法就是 再對這個 $p(y_n | x_n, \\theta_{A}^{*})$ 取 log 再取 gradient 並且平方 ( parameter.grad )^2.\n",
        "\n",
        "每一個參數我都可以使用 pytorch 的 backward 之後再取 gradient 的性質算出各自的 $F_i$.\n",
        "\n",
        "有關這個 $F$ 其實博大精深，是來自於 fisher information matrix. 底下我放上有關這個lifelong learning 在 fisher information matrix 上是怎麼簡單的近似到這一項，簡單的推導來自 [Continual Learning in Neural\n",
        "Networks](https://arxiv.org/pdf/1910.02718.pdf) 第2.4.1 小節 與 2.4 節\n",
        "\n",
        "For You Information: [Elastic Weight Consolidation](https://arxiv.org/pdf/1612.00796.pdf)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K511GmRzyYWa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "class EWC(object):\n",
        "  \"\"\"\n",
        "    @article{kirkpatrick2017overcoming,\n",
        "        title={Overcoming catastrophic forgetting in neural networks},\n",
        "        author={Kirkpatrick, James and Pascanu, Razvan and Rabinowitz, Neil and Veness, Joel and Desjardins, Guillaume and Rusu, Andrei A and Milan, Kieran and Quan, John and Ramalho, Tiago and Grabska-Barwinska, Agnieszka and others},\n",
        "        journal={Proceedings of the national academy of sciences},\n",
        "        year={2017},\n",
        "        url={https://arxiv.org/abs/1612.00796}\n",
        "    }\n",
        "  \"\"\"\n",
        "  def __init__(self, model: nn.Module, dataloaders: list, device):\n",
        "    \n",
        "    self.model = model\n",
        "    self.dataloaders = dataloaders\n",
        "    self.device = device\n",
        "    \n",
        "    self.params = {n: p for n, p in self.model.named_parameters() if p.requires_grad} #抓出模型的所有參數\n",
        "    self._means = {} # 初始化 平均參數\n",
        "    self._precision_matrices = self._calculate_importance() # 產生 EWC 的 Fisher (F) 矩陣 \n",
        "    \n",
        "    for n, p in self.params.items():\n",
        "      self._means[n] = p.clone().detach() # 算出每個參數的平均 （用之前任務的資料去算平均）\n",
        "  \n",
        "  def _calculate_importance(self):\n",
        "    precision_matrices = {}\n",
        "    for n, p in self.params.items(): # 初始化 Fisher (F) 的矩陣（都補零）\n",
        "      precision_matrices[n] = p.clone().detach().fill_(0)\n",
        "\n",
        "    self.model.eval()\n",
        "    dataloader_num=len(self.dataloaders)\n",
        "    number_data = sum([len(loader) for loader in self.dataloaders])\n",
        "    for dataloader in self.dataloaders:\n",
        "      for data in dataloader:\n",
        "        self.model.zero_grad()\n",
        "        input = data[0].to(self.device)\n",
        "        output = self.model(input).view(1, -1)\n",
        "        label = output.max(1)[1].view(-1)\n",
        "        \n",
        "        ############################################################################\n",
        "        #####                      產生 EWC 的 Fisher(F) 矩陣                    #####\n",
        "        ############################################################################    \n",
        "        loss = F.nll_loss(F.log_softmax(output, dim=1), label)             \n",
        "        loss.backward()                                                    \n",
        "                                                                           \n",
        "        for n, p in self.model.named_parameters():                         \n",
        "            precision_matrices[n].data += p.grad.data ** 2 / number_data   \n",
        "                                                                   \n",
        "    precision_matrices = {n: p for n, p in precision_matrices.items()}\n",
        "    return precision_matrices\n",
        "\n",
        "  def penalty(self, model: nn.Module):\n",
        "    loss = 0\n",
        "    for n, p in model.named_parameters():\n",
        "      _loss = self._precision_matrices[n] * (p - self._means[n]) ** 2\n",
        "      loss += _loss.sum()\n",
        "    return loss\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jmsPw6avjl5B",
        "colab_type": "text"
      },
      "source": [
        "## MAS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C7fSYpALrAVw",
        "colab_type": "text"
      },
      "source": [
        "### Memory Aware Synapses\n",
        "概念:\n",
        "老師的影片中，將它歸類到和 EWC 一樣的方法，只是算這個 important weight 的方式不太一樣.底下我將說明這個方法該怎麼實做\n",
        "\n",
        "MAS:\n",
        "在 MAS 內，學習一個連續的 tasks, task A, 和 task B, 他的 loss function 定義如下:\n",
        "\n",
        "$$\\mathcal{L}_B = \\mathcal{L}(\\theta) + \\sum_{i} \\frac{\\lambda}{2} \\Omega_i (\\theta_{i} - \\theta_{A,i}^{*})^2$$\n",
        "\n",
        "和 ewc不同的是 式子中的 $F_i$ 被取代成 $\\Omega_i$ , $\\Omega_i$ 來自於以下的式子：\n",
        "\n",
        "$$\\Omega_i = || \\frac{\\partial \\ell_2^2(M(x_k; \\theta))}{\\partial \\theta_i} || $$ \n",
        "\n",
        "$x_k$ 是 來自於 前面 task 的 sample data。 式子上的作法就是對最後模型的 output vector (最後一層)做 l2 norm 後取平方 再對各自的weight微分(取gradient) 並且取 該 gradient 的絕對值，在該paper 中其實也可以對各個層的 output vector 做 l2 norm ( local 版本)，這邊只實做 global 的版本。\n",
        "\n",
        "\n",
        "For Your Information: \n",
        "[Memory Aware Synapses](https://arxiv.org/pdf/1711.09601.pdf)\n",
        " \n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "btFvFJMmqxE0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "class MAS(object):\n",
        "    \"\"\"\n",
        "    @article{aljundi2017memory,\n",
        "      title={Memory Aware Synapses: Learning what (not) to forget},\n",
        "      author={Aljundi, Rahaf and Babiloni, Francesca and Elhoseiny, Mohamed and Rohrbach, Marcus and Tuytelaars, Tinne},\n",
        "      booktitle={ECCV},\n",
        "      year={2018},\n",
        "      url={https://eccv2018.org/openaccess/content_ECCV_2018/papers/Rahaf_Aljundi_Memory_Aware_Synapses_ECCV_2018_paper.pdf}\n",
        "    }\n",
        "    \"\"\"\n",
        "    def __init__(self, model: nn.Module, dataloaders: list, device):\n",
        "        self.model = model \n",
        "        self.dataloaders = dataloaders\n",
        "        self.params = {n: p for n, p in self.model.named_parameters() if p.requires_grad} #抓出模型的所有參數\n",
        "        self._means = {} # 初始化 平均參數\n",
        "        self.device = device\n",
        "        self._precision_matrices = self.calculate_importance() # 產生 MAS 的 Omega(Ω) 矩陣\n",
        "    \n",
        "        for n, p in self.params.items():\n",
        "            self._means[n] = p.clone().detach()\n",
        "    \n",
        "    def calculate_importance(self):\n",
        "        print('Computing MAS')\n",
        "\n",
        "        precision_matrices = {}\n",
        "        for n, p in self.params.items():\n",
        "            precision_matrices[n] = p.clone().detach().fill_(0) # 初始化 Omega(Ω) 矩陣（都補零）\n",
        "\n",
        "        self.model.eval()\n",
        "        dataloader_num = len(self.dataloaders)\n",
        "        num_data = sum([len(loader) for loader in self.dataloaders])\n",
        "        for dataloader in self.dataloaders:\n",
        "            for data in dataloader:\n",
        "                self.model.zero_grad()\n",
        "                output = self.model(data[0].to(self.device))\n",
        "\n",
        "                #######################################################################################\n",
        "                #####  產生 MAS 的 Omega(Ω) 矩陣 ( 對 output 向量 算他的 l2 norm 的平方) 再取 gradient  #####\n",
        "                #######################################################################################\n",
        "                output.pow_(2)                                                   \n",
        "                loss = torch.sum(output,dim=1)                                   \n",
        "                loss = loss.mean()                                               \n",
        "                loss.backward()                                                  \n",
        "                                          \n",
        "                for n, p in self.model.named_parameters():                      \n",
        "                    precision_matrices[n].data += p.grad.abs() / num_data ## difference with EWC      \n",
        "\n",
        "        precision_matrices = {n: p for n, p in precision_matrices.items()}\n",
        "        return precision_matrices\n",
        "\n",
        "    def penalty(self, model: nn.Module):\n",
        "        loss = 0\n",
        "        for n, p in model.named_parameters():\n",
        "            _loss = self._precision_matrices[n] * (p - self._means[n]) ** 2\n",
        "            loss += _loss.sum()\n",
        "        return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NWQ_JBlbFnKV",
        "colab_type": "text"
      },
      "source": [
        "# 資料"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vdep_aMvUYqI",
        "colab_type": "text"
      },
      "source": [
        "## 資料預處理\n",
        "- 轉換 MNIST  ($1*28*28$) 到 ($3*32*32$)\n",
        "- 轉換 USPS   ($1*16*16$) 到 ($3*32*32$)\n",
        "- 正規化 圖片"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xVHrWsHfIPtY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Convert2RGB(object):\n",
        "  \n",
        "  def __init__(self, num_channel):\n",
        "    self.num_channel = num_channel\n",
        "\n",
        "  def __call__(self, img):                                                                                                                                                                                                                              \n",
        "    # If the channel of img is not equal to desired size,\n",
        "    # then expand the channel of img to desired size.\n",
        "    img_channel = img.size()[0]\n",
        "    img = torch.cat([img] * (self.num_channel - img_channel + 1), 0)\n",
        "    return img\n",
        "\n",
        "\n",
        "class Pad(object):\n",
        "\n",
        "  def __init__(self, size, fill=0, padding_mode='constant'):\n",
        "    self.size = size\n",
        "    self.fill = fill\n",
        "    self.padding_mode = padding_mode\n",
        "    \n",
        "  def __call__(self, img):\n",
        "    # If the H and W of img is not equal to desired size,\n",
        "    # then pad the channel of img to desired size.\n",
        "    img_size = img.size()[1]\n",
        "    assert ((self.size - img_size) % 2 == 0)\n",
        "    padding = (self.size - img_size) // 2\n",
        "    padding = (padding, padding, padding, padding)\n",
        "    return F.pad(img, padding, self.padding_mode, self.fill)\n",
        "\n",
        "def get_transform():\n",
        "  transform = transforms.Compose([transforms.ToTensor(),\n",
        "                                  Pad(32),\n",
        "                                  Convert2RGB(3),\n",
        "                                  transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])])\n",
        "  return transform\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PW4r9Jd-etyG",
        "colab_type": "text"
      },
      "source": [
        "## 準備 資料集\n",
        "- MNIST   : 一張圖片資料大小:  $28*28*1$, 灰階 , 10 個種類\n",
        "- SVHN    : 一張圖片資料大小:  $32*32*3$, RGB , 10 個種類\n",
        "- USPS    : 一張圖片資料大小:  $16*16*1$, 灰階 , 10 個種類"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HPIeRDtIox0M",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Data():\n",
        "\n",
        "  def __init__(self, path):\n",
        "\n",
        "    transform = get_transform()\n",
        "\n",
        "    self.MNIST_dataset = datasets.MNIST(root = os.path.join(path, \"MNIST\"),\n",
        "                                        transform=transform,\n",
        "                                        train = True,\n",
        "                                        download = True)\n",
        "\n",
        "    self.SVHN_dataset = datasets.SVHN(root = os.path.join(path, \"SVHN\"),\n",
        "                                      transform=transform,\n",
        "                                      split='train',\n",
        "                                      download = True)\n",
        "\n",
        "    self.USPS_dataset = datasets.USPS(root = os.path.join(path, \"USPS\"),\n",
        "                                            transform=transform,\n",
        "                                            train = True,\n",
        "                                            download = True)\n",
        "    \n",
        "  def get_datasets(self):\n",
        "      a = [(self.SVHN_dataset, \"SVHN\"),(self.MNIST_dataset, \"MNIST\"),(self.USPS_dataset, \"USPS\")]\n",
        "      return a\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eMtV82EPjsld",
        "colab_type": "text"
      },
      "source": [
        "## 建立 Dataloader\n",
        "- *.train_loader: 拿取訓練集並訓練 \\\\\n",
        "- *.val_loader: 拿取驗證集並驗測結果 \\\\"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "29-5g8ZHjs_3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Dataloader():\n",
        "\n",
        "  def __init__(self, dataset, batch_size, split_ratio=0.1):\n",
        "    self.dataset = dataset[0]\n",
        "    self.name = dataset[1]\n",
        "    train_sampler, val_sampler = self.split_dataset(split_ratio)\n",
        "\n",
        "    self.train_dataset_size = len(train_sampler)\n",
        "    self.val_dataset_size = len(val_sampler)\n",
        "\n",
        "    self.train_loader = data.DataLoader(self.dataset, batch_size = batch_size, sampler=train_sampler)\n",
        "    self.val_loader = data.DataLoader(self.dataset, batch_size = batch_size, sampler=val_sampler)\n",
        "    self.train_iter = self.infinite_iter()\n",
        "\n",
        "  def split_dataset(self, split_ratio):\n",
        "    data_size = len(self.dataset)\n",
        "    split = int(data_size * split_ratio)\n",
        "    indices = list(range(data_size))\n",
        "    np.random.shuffle(indices)\n",
        "    train_idx, valid_idx = indices[split:], indices[:split]\n",
        "\n",
        "    train_sampler = sampler.SubsetRandomSampler(train_idx)\n",
        "    val_sampler = sampler.SubsetRandomSampler(valid_idx)\n",
        "    return train_sampler, val_sampler\n",
        "    \n",
        "  def infinite_iter(self):\n",
        "    it = iter(self.train_loader)\n",
        "    while True:\n",
        "      try:\n",
        "        ret = next(it)\n",
        "        yield ret\n",
        "      except StopIteration:\n",
        "        it = iter(self.train_loader)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vzG5BWtHGA3p",
        "colab_type": "text"
      },
      "source": [
        "# 小工具"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vMBoCSH5MBLN",
        "colab_type": "text"
      },
      "source": [
        "## 儲存模型"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uCZuQrWiMGmH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def save_model(model, optimizer, store_model_path):\n",
        "  # save model and optimizer\n",
        "  torch.save(model.state_dict(), f'{store_model_path}.ckpt')\n",
        "  torch.save(optimizer.state_dict(), f'{store_model_path}.opt')\n",
        "  return\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nde98xvAMxAd",
        "colab_type": "text"
      },
      "source": [
        "##載入模型\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FGzZ2Yp2MxK-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def load_model(model, optimizer, load_model_path):\n",
        "  # load model and optimizer\n",
        "  print(f'Load model from {load_model_path}')\n",
        "  model.load_state_dict(torch.load(f'{load_model_path}.ckpt'))\n",
        "  optimizer.load_state_dict(torch.load(f'{load_model_path}.opt'))\n",
        "  return model, optimizer\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eoz6awEcOIAz",
        "colab_type": "text"
      },
      "source": [
        "## 建立模型 & 優化器"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TvWqv_JlOOix",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def build_model(data_path, batch_size, learning_rate): \n",
        "  # create model\n",
        "  model = Model().to(device)\n",
        "  optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
        "  data = Data(data_path)\n",
        "  datasets = data.get_datasets()\n",
        "  tasks = []\n",
        "  for dataset in datasets:\n",
        "    tasks.append(Dataloader(dataset, batch_size))\n",
        "\n",
        "  return model, optimizer, tasks\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74De0sS-O50R",
        "colab_type": "text"
      },
      "source": [
        "# 訓練"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_q9Co3vuGWfu",
        "colab_type": "text"
      },
      "source": [
        "## 正常訓練 ( baseline )"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TBnE9GbiO8Ob",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def normal_train(model, optimizer, task, total_epochs, summary_epochs):\n",
        "  model.train()\n",
        "  model.zero_grad()\n",
        "  ceriation = nn.CrossEntropyLoss()\n",
        "  losses = []\n",
        "  loss = 0.0\n",
        "  for epoch in range(summary_epochs):\n",
        "    imgs, labels = next(task.train_iter)\n",
        "    imgs, labels = imgs.to(device), labels.to(device)\n",
        "    outputs = model(imgs)\n",
        "    ce_loss = ceriation(outputs, labels)\n",
        "    \n",
        "    optimizer.zero_grad()\n",
        "    ce_loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    loss += ce_loss.item()\n",
        "    if (epoch + 1) % 50 == 0:\n",
        "      loss = loss / 50\n",
        "      print (\"\\r\", \"train task {} [{}] loss: {:.3f}      \".format(task.name, (total_epochs + epoch + 1), loss), end=\" \")\n",
        "      losses.append(loss)\n",
        "      loss = 0.0\n",
        "    \n",
        "  return model, optimizer, losses\n",
        "  "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m2FlojHR_4qb",
        "colab_type": "text"
      },
      "source": [
        "## EWC 訓練"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nLHALesw_61i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def ewc_train(model, optimizer, task, total_epochs, summary_epochs, ewc, lambda_ewc):\n",
        "  model.train()\n",
        "  model.zero_grad()\n",
        "  ceriation = nn.CrossEntropyLoss()\n",
        "  losses = []\n",
        "  loss = 0.0\n",
        "  for epoch in range(summary_epochs):\n",
        "    imgs, labels = next(task.train_iter)\n",
        "    imgs, labels = imgs.to(device), labels.to(device)\n",
        "    outputs = model(imgs)\n",
        "    ce_loss = ceriation(outputs, labels)\n",
        "    total_loss = ce_loss\n",
        "    ewc_loss = ewc.penalty(model)\n",
        "    total_loss += lambda_ewc * ewc_loss \n",
        "    \n",
        "    optimizer.zero_grad()\n",
        "    total_loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    loss += total_loss.item()\n",
        "    if (epoch + 1) % 50 == 0:\n",
        "      loss = loss / 50\n",
        "      print (\"\\r\", \"train task {} [{}] loss: {:.3f}      \".format(task.name, (total_epochs + epoch + 1), loss), end=\" \")\n",
        "      losses.append(loss)\n",
        "      loss = 0.0\n",
        "    \n",
        "  return model, optimizer, losses"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0B6to7GuqvPX",
        "colab_type": "text"
      },
      "source": [
        "## MAS 訓練"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fWhZz9uZquew",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def mas_train(model, optimizer, task, total_epochs, summary_epochs, mas_tasks, lambda_mas,alpha=0.8):\n",
        "  model.train()\n",
        "  model.zero_grad()\n",
        "  ceriation = nn.CrossEntropyLoss()\n",
        "  losses = []\n",
        "  loss = 0.0\n",
        "  for epoch in range(summary_epochs):\n",
        "    imgs, labels = next(task.train_iter)\n",
        "    imgs, labels = imgs.to(device), labels.to(device)\n",
        "    outputs = model(imgs)\n",
        "    ce_loss = ceriation(outputs, labels)\n",
        "    total_loss = ce_loss\n",
        "    mas_tasks.reverse()\n",
        "    if len(mas_tasks) > 1:\n",
        "        preprevious = 1 - alpha\n",
        "        scalars = [alpha,preprevious]\n",
        "        for mas,scalar in zip(mas_tasks[:2],scalars):\n",
        "            mas_loss = mas.penalty(model)\n",
        "            total_loss += lambda_mas * mas_loss * scalar\n",
        "    elif len(mas_tasks) == 1:\n",
        "        mas_loss = mas_tasks[0].penalty(model)\n",
        "        total_loss += lambda_mas * mas_loss\n",
        "    else:\n",
        "        pass\n",
        "    \n",
        "    optimizer.zero_grad()\n",
        "    total_loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    loss += total_loss.item()\n",
        "    if (epoch + 1) % 50 == 0:\n",
        "      loss = loss / 50\n",
        "      print (\"\\r\", \"train task {} [{}] loss: {:.3f}      \".format(task.name, (total_epochs + epoch + 1), loss), end=\" \")\n",
        "      losses.append(loss)\n",
        "      loss = 0.0\n",
        "    \n",
        "  return model, optimizer, losses"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6cuHVXxAfHrA",
        "colab_type": "text"
      },
      "source": [
        "## 驗證\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZBp-n3FrfOCe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def val(model, task):\n",
        "  model.eval()\n",
        "  correct_cnt = 0\n",
        "  for imgs, labels in task.val_loader:\n",
        "    imgs, labels = imgs.to(device), labels.to(device)\n",
        "    outputs = model(imgs)\n",
        "    _, pred_label = torch.max(outputs.data, 1)\n",
        "\n",
        "    correct_cnt += (pred_label == labels.data).sum().item()\n",
        "    \n",
        "  return correct_cnt / task.val_dataset_size\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dFEYmPAlx_SX",
        "colab_type": "text"
      },
      "source": [
        "## 主訓練程序\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cJ54vDP2yC2S",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train_process(model, optimizer, tasks, config):\n",
        "  task_loss, acc = {}, {}\n",
        "  for task_id, task in enumerate(tasks):\n",
        "    print ('\\n')\n",
        "    total_epochs = 0\n",
        "    task_loss[task.name] = []\n",
        "    acc[task.name] = []\n",
        "    if config.mode == 'basic' or task_id == 0:\n",
        "      while (total_epochs < config.num_epochs):\n",
        "        model, optimizer, losses = normal_train(model, optimizer, task, total_epochs, config.summary_epochs)\n",
        "        task_loss[task.name] +=  losses\n",
        "\n",
        "        for subtask in range(task_id + 1):\n",
        "          acc[tasks[subtask].name].append(val(model, tasks[subtask]))\n",
        "\n",
        "        total_epochs += config.summary_epochs\n",
        "        if total_epochs % config.store_epochs == 0 or total_epochs >= config.num_epochs:\n",
        "          save_model(model, optimizer, config.store_model_path)\n",
        "    \n",
        "    if config.mode == 'ewc' and task_id > 0:\n",
        "      old_dataloaders = []\n",
        "      for old_task in range(task_id): \n",
        "        old_dataloaders += [tasks[old_task].val_loader]\n",
        "      ewc = EWC(model, old_dataloaders, device)\n",
        "      while (total_epochs < config.num_epochs):\n",
        "        model, optimizer, losses = ewc_train(model, optimizer, task, total_epochs, config.summary_epochs, ewc, config.lifelong_coeff)\n",
        "        task_loss[task.name] +=  losses\n",
        "\n",
        "        for subtask in range(task_id + 1):\n",
        "          acc[tasks[subtask].name].append(val(model, tasks[subtask]))\n",
        "\n",
        "        total_epochs += config.summary_epochs\n",
        "        if total_epochs % config.store_epochs == 0 or total_epochs >= config.num_epochs:\n",
        "          save_model(model, optimizer, config.store_model_path)\n",
        "\n",
        "    if config.mode == 'mas' and task_id > 0:\n",
        "      old_dataloaders = []\n",
        "      mas_tasks = []\n",
        "      for old_task in range(task_id): \n",
        "        old_dataloaders += [tasks[old_task].val_loader]\n",
        "        mas = MAS(model, old_dataloaders, device)\n",
        "        mas_tasks += [mas]\n",
        "      while (total_epochs < config.num_epochs):\n",
        "        model, optimizer, losses = mas_train(model, optimizer, task, total_epochs, config.summary_epochs, mas_tasks, config.lifelong_coeff)\n",
        "        task_loss[task.name] +=  losses\n",
        "\n",
        "        for subtask in range(task_id + 1):\n",
        "          acc[tasks[subtask].name].append(val(model, tasks[subtask]))\n",
        "\n",
        "        total_epochs += config.summary_epochs\n",
        "        if total_epochs % config.store_epochs == 0 or total_epochs >= config.num_epochs:\n",
        "          save_model(model, optimizer, config.store_model_path)\n",
        "\n",
        "    if config.mode == 'scp' and task_id > 0:\n",
        "      pass\n",
        "      ########################################\n",
        "      ##       TODO 區塊 （ PART 2 ）         ##\n",
        "      ########################################\n",
        "      ##    PART 2  implementation 的部份    ##\n",
        "      ##   你也可以寫別的 regularization 方法  ##\n",
        "      ##    助教這裡有提供的是  scp    的 作法   ##\n",
        "      ##     Slicer Cramer Preservation     ##\n",
        "      ########################################\n",
        "      ########################################\n",
        "      ##       TODO 區塊 （ PART 2 ）         ##\n",
        "      ########################################\n",
        "  return task_loss, acc\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7PbfgB3n9eoT",
        "colab_type": "text"
      },
      "source": [
        "# 設定"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3kWSZ4w39gzj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        " class configurations(object):\n",
        "  def __init__(self):\n",
        "    self.batch_size = 256\n",
        "    self.num_epochs = 10000\n",
        "    self.store_epochs = 250\n",
        "    self.summary_epochs = 250\n",
        "    self.learning_rate = 0.0005\n",
        "    self.load_model = False\n",
        "    self.store_model_path = \"./model\"\n",
        "    self.load_model_path = \"./model\"\n",
        "    self.data_path = \"./data\"\n",
        "    self.mode = None\n",
        "    self.lifelong_coeff = 0.5\n",
        "\n",
        "###### 你也可以自己設定參數   ########\n",
        "###### 但上面的參數 是這次作業的預設直 #########\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w464f4KOLUh6",
        "colab_type": "text"
      },
      "source": [
        "#主程式區塊\n",
        "- 給 EWC, MAS 超參數 $\\lambda$ \n",
        "- 訓練"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AJwVkorvLaSh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "the order is svhn -> mnist -> usps\n",
        "===============================================\n",
        "\n",
        "\"\"\"\n",
        "# import tqdm\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    mode_list = ['mas','ewc','basic']\n",
        "\n",
        "    ## hint: 謹慎的去選擇 lambda 超參數 / ewc: 80~400, mas: 0.1 - 10\n",
        "    ############################################################################\n",
        "    #####                           TODO 區塊 （ PART 1 ）                   #####\n",
        "    ############################################################################ \n",
        "    coeff_list = [0, 0 ,0 ]  ## 你需要在這 微調 lambda 參數, mas, ewc, baseline=0##  \n",
        "    ############################################################################\n",
        "    #####                           TODO 區塊 （ PART 1 ）                   #####\n",
        "    ############################################################################ \n",
        "    \n",
        "    config = configurations()\n",
        "    count = 0\n",
        "    for mode in mode_list:\n",
        "        config.mode = mode\n",
        "        config.lifelong_coeff = coeff_list[count]\n",
        "        print(\"{} training\".format(config.mode))    \n",
        "        model, optimizer, tasks = build_model(config.load_model_path, config.batch_size, config.learning_rate)\n",
        "        print (\"Finish build model\")\n",
        "        if config.load_model:\n",
        "            model, optimizer = load_model(model, optimizer, config.load_model_path)\n",
        "        task_loss, acc = train_process(model, optimizer, tasks, config)\n",
        "        with open(f'./{config.mode}_acc.txt', 'w') as f:\n",
        "            json.dump(acc, f)\n",
        "        count += 1\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DSJX338dA2He",
        "colab_type": "text"
      },
      "source": [
        "# 畫出 Result 圖片"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4X_RvV4my5Jl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "\n",
        "def plot_result(mode_list, task1, task2, task3):\n",
        "  \n",
        "    #draw the lines\n",
        "    count = 0\n",
        "    for reg_name in mode_list:\n",
        "        label = reg_name\n",
        "        with open(f'./{reg_name}_acc.txt', 'r') as f:\n",
        "            acc = json.load(f)\n",
        "        if count == 0: \n",
        "            color= 'red'\n",
        "        elif count  == 1:\n",
        "            color= 'blue'\n",
        "        else:\n",
        "            color = 'purple'\n",
        "        ax1=plt.subplot(3, 1, 1)\n",
        "        plt.plot(range(len(acc[task1])),acc[task1],color,label=label)\n",
        "        ax1.set_ylabel(task1)\n",
        "        ax2=plt.subplot(3, 1, 2,sharex=ax1,sharey=ax1)\n",
        "        plt.plot(range(len(acc[task3]),len(acc[task1])),acc[task2],color,label=label)\n",
        "        ax2.set_ylabel(task2)\n",
        "        ax3=plt.subplot(3, 1, 3,sharex=ax1,sharey=ax1)\n",
        "        ax3.set_ylabel(task3)\n",
        "        plt.plot(range(len(acc[task2]),len(acc[task1])),acc[task3],color,label=label)\n",
        "        count += 1\n",
        "    plt.ylim((0.02,1.02))\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "    return \n",
        "\n",
        "mode_list = ['ewc','mas','basic']\n",
        "plot_result(mode_list,'SVHN','MNIST','USPS')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "43QGlXTxyzw_",
        "colab_type": "text"
      },
      "source": [
        "在今年 ICLR 2020 的 paper，有以這兩種方法做 baseline，並對這兩種方法各自做了一個 geometry view，也提出新的方法，有興趣的人可以參考\n",
        "\n",
        "paper link 如下 [SLICED CRAMER´ SYNAPTIC CONSOLIDATION FOR\n",
        "PRESERVING DEEPLY LEARNED REPRESENTATIONS](https://openreview.net/pdf?id=BJge3TNKwH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cbMUPaN_zAs7",
        "colab_type": "text"
      },
      "source": [
        "# 進階 \n",
        "請實做其他的 regularization 的方法，助教有提供的是 SCP 的作法，\n",
        "\n",
        "你也可以考慮實做出 SI, Rimennian Walk, IMM, 或是上面的方法, \n",
        "\n",
        "你可以參考助教上方的寫法，寫出雷同的 class 跟 training 來 train，\n",
        "\n",
        "記得畫出與上方雷同的 evaluation 圖表 (show result) example 需要比對的話 可以參考助教給的 slide。\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m3aOQ2XI-Prm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def sample_spherical(npoints, ndim=3):\n",
        "    vec = np.random.randn(ndim, npoints)\n",
        "    vec /= np.linalg.norm(vec, axis=0)\n",
        "    return vec"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bcjtln1T6U7T",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class SCP(object):\n",
        "    \"\"\"\n",
        "    OPEN REVIEW VERSION:\n",
        "    https://openreview.net/forum?id=BJge3TNKwH\n",
        "    \"\"\"\n",
        "    def __init__(self, model: nn.Module, dataloaders: list, L: int, device):\n",
        "        self.model = model \n",
        "        self.dataloaders = dataloaders\n",
        "        self.params = {n: p for n, p in self.model.named_parameters() if p.requires_grad}\n",
        "        self._means = {}\n",
        "        self.L= L\n",
        "        self.device = device\n",
        "        self._precision_matrices = self.calculate_importance()\n",
        "    \n",
        "        for n, p in self.params.items():\n",
        "            self._means[n] = p.clone().detach()\n",
        "    \n",
        "    def calculate_importance(self):\n",
        "        print('Computing SCP')\n",
        "\n",
        "        precision_matrices = {}\n",
        "        for n, p in self.params.items():\n",
        "            precision_matrices[n] = p.clone().detach().fill_(0)\n",
        "\n",
        "        self.model.eval()\n",
        "        dataloader_num = len(self.dataloaders)\n",
        "        num_data = sum([len(loader) for loader in self.dataloaders])\n",
        "        for dataloader in self.dataloaders:\n",
        "            for data in dataloader:\n",
        "                self.model.zero_grad()\n",
        "                output = self.model(data[0].to(self.device))\n",
        "\n",
        "                ####################################################################################\n",
        "                #####                            TODO 區塊 （ PART 2 ）                           #####\n",
        "                ####################################################################################\n",
        "                ##### 產生 SCP 的 Gamma(Γ) 矩陣（ 如同 MAS 的 Omega(Ω) 矩陣, EWC 的 Fisher(F) 矩陣 ）#####\n",
        "                ####################################################################################\n",
        "                #####        1.對所有資料的 Output vector 取 平均 得到 平均 vector φ(:,θ_A* )       #####\n",
        "                ####################################################################################\n",
        "\n",
        "                ####################################################################################\n",
        "                #####   2. 隨機 從 單位球殼 取樣 L 個 vector ξ #（ Hint: sample_spherical() ）      #####\n",
        "                ####################################################################################\n",
        "\n",
        "                ####################################################################################\n",
        "                #####   3.    每一個 vector ξ 和 vector φ( :,θ_A* )內積得到 scalar ρ               #####\n",
        "                #####           對 scalar ρ 取 backward ， 每個參數得到各自的 gradient ∇ρ           #####\n",
        "                #####       每個參數的 gradient ∇ρ 取平方 取 L 平均 得到 各個參數的 Γ scalar          #####  \n",
        "                #####              所有參數的  Γ scalar 組合而成其實就是 Γ 矩陣                      #####\n",
        "                ####(hint: 記得 每次 backward 之後 要 zero_grad 去 清 gradient, 不然 gradient會累加 )######   \n",
        "                ####################################################################################\n",
        "      \n",
        "                ####################################################################################      \n",
        "                #####                            TODO 區塊 （ PART 2 ）                          #####\n",
        "                ####################################################################################\n",
        "\n",
        "        precision_matrices = {n: p for n, p in precision_matrices.items()}\n",
        "        return precision_matrices\n",
        "\n",
        "    def penalty(self, model: nn.Module):\n",
        "        loss = 0\n",
        "        for n, p in model.named_parameters():\n",
        "            _loss = self._precision_matrices[n] * (p - self._means[n]) ** 2\n",
        "            loss += _loss.sum()\n",
        "        return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tp6EbyrhXoAH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def scp_train(model, optimizer, task, total_epochs, summary_epochs, scp_tasks, lambda_scp,alpha=0.65):\n",
        "  losses = []\n",
        "  loss = 0.0\n",
        "  ###############################\n",
        "  #####  TODO 區塊 （PART 2） #####\n",
        "  ###############################\n",
        "  ##  參考 MAS. EWC train 的寫法 ##                 \n",
        "  ###############################\n",
        "  #####  TODO 區塊 （PART 2） #####\n",
        "  ###############################\n",
        "  return model, optimizer, losses"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V1PxP2W3ZcrV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# if __name__ == \"__main__\": \n",
        "#   pass \n",
        "###############################\n",
        "#####  TODO 區塊 （PART 2） #####\n",
        "###############################\n",
        "##     參考 main 區塊一樣       ##                 \n",
        "##     的 code 結合新方法       ##\n",
        "###############################\n",
        "#####  TODO 區塊 （PART 2） #####\n",
        "###############################"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}