MAML-omniglot-ADAM-20way-1-shot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-14T03:22:31.272106Z",
     "start_time": "2020-04-14T03:22:30.955368Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import os\n",
    "import zipfile\n",
    "\n",
    "# root_path = './../datasets'\n",
    "# processed_folder =  os.path.join(root_path)\n",
    "\n",
    "# zip_ref = zipfile.ZipFile(os.path.join(root_path,'omniglot.zip'), 'r')\n",
    "# zip_ref.extractall(root_path)\n",
    "# zip_ref.close()\n",
    "root_dir = './../datasets/omniglot/python'\n",
    "root_dir_train = os.path.join(root_dir,'images_background')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-07T06:55:03.418079Z",
     "start_time": "2020-04-07T06:55:03.346124Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据预处理\n",
    "拿到原始数据之后先将下面的代码取消注释，进行数据预处理。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-13T07:07:01.042611Z",
     "start_time": "2020-04-13T07:07:01.035061Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "# # 数据预处理\n",
    "# import torchvision.transforms as transforms\n",
    "# from PIL import Image\n",
    "\n",
    "# '''\n",
    "# an example of img_items:\n",
    "# ( '0709_17.png',\n",
    "#   'Alphabet_of_the_Magi/character01',\n",
    "#   './../datasets/omniglot/python/images_background/Alphabet_of_the_Magi/character01')\n",
    "# '''\n",
    "\n",
    "\n",
    "# root_dir_train = os.path.join(root_dir, 'images_background')\n",
    "# root_dir_test = os.path.join(root_dir, 'images_evaluation')\n",
    "\n",
    "# def find_classes(root_dir_train):\n",
    "#     img_items = []\n",
    "#     for (root, dirs, files) in os.walk(root_dir_train): \n",
    "#         for file in files:\n",
    "#             if (file.endswith(\"png\")):\n",
    "#                 r = root.split('/')\n",
    "#                 img_items.append((file, r[-2] + \"/\" + r[-1], root))\n",
    "#     print(\"== Found %d items \" % len(img_items))\n",
    "#     return img_items\n",
    "\n",
    "# ## 构建一个词典{class:idx}\n",
    "# def index_classes(items):\n",
    "#     class_idx = {}\n",
    "#     count = 0\n",
    "#     for item in items:\n",
    "#         if item[1] not in class_idx:\n",
    "#             class_idx[item[1]] = count\n",
    "#             count += 1\n",
    "#     print('== Found {} classes'.format(len(class_idx)))\n",
    "#     return class_idx\n",
    "        \n",
    "\n",
    "# img_items_train =  find_classes(root_dir_train) # [(file1, label1, root1),..]\n",
    "# img_items_test = find_classes(root_dir_test)\n",
    "\n",
    "# class_idx_train = index_classes(img_items_train)\n",
    "# class_idx_test = index_classes(img_items_test)\n",
    "\n",
    "\n",
    "# def generate_temp(img_items,class_idx):\n",
    "#     temp = dict()\n",
    "#     for imgname, classes, dirs in img_items:\n",
    "#         img = '{}/{}'.format(dirs, imgname)\n",
    "#         label = class_idx[classes]\n",
    "#         transform = transforms.Compose([lambda img: Image.open(img).convert('L'),\n",
    "#                                   lambda img: img.resize((28,28)),\n",
    "#                                   lambda img: np.reshape(img, (28,28,1)),\n",
    "#                                   lambda img: np.transpose(img, [2,0,1]),\n",
    "#                                   lambda img: img/255.\n",
    "#                                   ])\n",
    "#         img = transform(img)\n",
    "#         if label in temp.keys():\n",
    "#             temp[label].append(img)\n",
    "#         else:\n",
    "#             temp[label] = [img]\n",
    "#     print('begin to generate omniglot.npy')\n",
    "#     return temp\n",
    "#     ## 每个字符包含20个样本\n",
    "\n",
    "# temp_train = generate_temp(img_items_train, class_idx_train)\n",
    "# temp_test = generate_temp(img_items_test, class_idx_test)\n",
    "\n",
    "# img_list = []\n",
    "# for label, imgs in temp_train.items():\n",
    "#     img_list.append(np.array(imgs))\n",
    "# img_list = np.array(img_list).astype(np.float) # [[20 imgs],..., 1623 classes in total]\n",
    "# print('data shape:{}'.format(img_list.shape)) # (964, 20, 1, 28, 28)\n",
    "# np.save(os.path.join(root_dir, 'omniglot_train.npy'), img_list)\n",
    "# print('end.')\n",
    "\n",
    "\n",
    "# img_list = []\n",
    "# for label, imgs in temp_test.items():\n",
    "#     img_list.append(np.array(imgs))\n",
    "# img_list = np.array(img_list).astype(np.float) # [[20 imgs],..., 1623 classes in total]\n",
    "# print('data shape:{}'.format(img_list.shape)) # (659, 20, 1, 28, 28)\n",
    "\n",
    "# np.save(os.path.join(root_dir, 'omniglot_test.npy'), img_list)\n",
    "# print('end.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-07T08:02:03.283025Z",
     "start_time": "2020-04-07T08:02:03.276106Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-14T03:23:43.874491Z",
     "start_time": "2020-04-14T03:23:41.975288Z"
    }
   },
   "outputs": [],
   "source": [
    "img_list_train = np.load(os.path.join(root_dir, 'omniglot_train.npy')) # (964, 20, 1, 28, 28)\n",
    "img_list_test = np.load(os.path.join(root_dir, 'omniglot_test.npy')) # (659, 20, 1, 28, 28)\n",
    "\n",
    "x_train = img_list_train\n",
    "x_test = img_list_test\n",
    "# num_classes = img_list.shape[0]\n",
    "datasets = {'train': x_train, 'test': x_test}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-14T03:23:51.352052Z",
     "start_time": "2020-04-14T03:23:50.310935Z"
    },
    "code_folding": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DB: train (964, 20, 1, 28, 28) test (659, 20, 1, 28, 28)\n"
     ]
    }
   ],
   "source": [
    "### 准备数据迭代器\n",
    "n_way = 20\n",
    "k_spt = 1  ## support data 的个数\n",
    "k_query = 15 ## query data 的个数\n",
    "imgsz = 28\n",
    "resize = imgsz\n",
    "task_num = 16\n",
    "batch_size = task_num\n",
    "\n",
    "indexes = {\"train\": 0, \"test\": 0}\n",
    "datasets = {\"train\": x_train, \"test\": x_test}\n",
    "print(\"DB: train\", x_train.shape, \"test\", x_test.shape)\n",
    "\n",
    "\n",
    "def load_data_cache(dataset):\n",
    "    \"\"\"\n",
    "    Collects several batches data for N-shot learning\n",
    "    :param dataset: [cls_num, 20, 84, 84, 1]\n",
    "    :return: A list with [support_set_x, support_set_y, target_x, target_y] ready to be fed to our networks\n",
    "    \"\"\"\n",
    "    #  take 5 way 1 shot as example: 5 * 1\n",
    "    setsz = k_spt * n_way\n",
    "    querysz = k_query * n_way\n",
    "    data_cache = []\n",
    "\n",
    "    # print('preload next 10 caches of batch_size of batch.')\n",
    "    for sample in range(10):  # num of epochs\n",
    "\n",
    "        x_spts, y_spts, x_qrys, y_qrys = [], [], [], []\n",
    "        for i in range(batch_size):  # one batch means one set\n",
    "\n",
    "            x_spt, y_spt, x_qry, y_qry = [], [], [], []\n",
    "            selected_cls = np.random.choice(dataset.shape[0], n_way, replace =  False) \n",
    "\n",
    "            for j, cur_class in enumerate(selected_cls):\n",
    "\n",
    "                selected_img = np.random.choice(20, k_spt + k_query, replace = False)\n",
    "\n",
    "                # 构造support集和query集\n",
    "                x_spt.append(dataset[cur_class][selected_img[:k_spt]])\n",
    "                x_qry.append(dataset[cur_class][selected_img[k_spt:]])\n",
    "                y_spt.append([j for _ in range(k_spt)])\n",
    "                y_qry.append([j for _ in range(k_query)])\n",
    "\n",
    "            # shuffle inside a batch\n",
    "            perm = np.random.permutation(n_way * k_spt)\n",
    "            x_spt = np.array(x_spt).reshape(n_way * k_spt, 1, resize, resize)[perm]\n",
    "            y_spt = np.array(y_spt).reshape(n_way * k_spt)[perm]\n",
    "            perm = np.random.permutation(n_way * k_query)\n",
    "            x_qry = np.array(x_qry).reshape(n_way * k_query, 1, resize, resize)[perm]\n",
    "            y_qry = np.array(y_qry).reshape(n_way * k_query)[perm]\n",
    " \n",
    "            # append [sptsz, 1, 84, 84] => [batch_size, setsz, 1, 84, 84]\n",
    "            x_spts.append(x_spt)\n",
    "            y_spts.append(y_spt)\n",
    "            x_qrys.append(x_qry)\n",
    "            y_qrys.append(y_qry)\n",
    "\n",
    "#         print(x_spts[0].shape)\n",
    "        # [b, setsz = n_way * k_spt, 1, 84, 84]\n",
    "        x_spts = np.array(x_spts).astype(np.float32).reshape(batch_size, setsz, 1, resize, resize)\n",
    "        y_spts = np.array(y_spts).astype(np.int).reshape(batch_size, setsz)\n",
    "        # [b, qrysz = n_way * k_query, 1, 84, 84]\n",
    "        x_qrys = np.array(x_qrys).astype(np.float32).reshape(batch_size, querysz, 1, resize, resize)\n",
    "        y_qrys = np.array(y_qrys).astype(np.int).reshape(batch_size, querysz)\n",
    "#         print(x_qrys.shape)\n",
    "        data_cache.append([x_spts, y_spts, x_qrys, y_qrys])\n",
    "\n",
    "    return data_cache\n",
    "\n",
    "datasets_cache = {\"train\": load_data_cache(x_train),  # current epoch data cached\n",
    "                       \"test\": load_data_cache(x_test)}\n",
    "\n",
    "def next(mode='train'):\n",
    "    \"\"\"\n",
    "    Gets next batch from the dataset with name.\n",
    "    :param mode: The name of the splitting (one of \"train\", \"val\", \"test\")\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    # update cache if indexes is larger than len(data_cache)\n",
    "    if indexes[mode] >= len(datasets_cache[mode]):\n",
    "        indexes[mode] = 0\n",
    "        datasets_cache[mode] = load_data_cache(datasets[mode])\n",
    "\n",
    "    next_batch = datasets_cache[mode][indexes[mode]]\n",
    "    indexes[mode] += 1\n",
    "\n",
    "    return next_batch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-14T03:27:10.793796Z",
     "start_time": "2020-04-14T03:27:10.756874Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from copy import deepcopy,copy\n",
    "        \n",
    "\n",
    "class BaseNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BaseNet, self).__init__()\n",
    "        self.vars = nn.ParameterList()  ## 包含了所有需要被优化的tensor\n",
    "        self.vars_bn = nn.ParameterList()\n",
    "        \n",
    "        # 第1个conv2d\n",
    "        weight = nn.Parameter(torch.ones(64, 1, 3, 3))\n",
    "        nn.init.kaiming_normal_(weight)\n",
    "        bias = nn.Parameter(torch.zeros(64))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "        # 第1个BatchNorm层\n",
    "        weight = nn.Parameter(torch.ones(64))\n",
    "        bias = nn.Parameter(torch.zeros(64))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "        running_mean = nn.Parameter(torch.zeros(64), requires_grad= False)\n",
    "        running_var = nn.Parameter(torch.zeros(64), requires_grad= False)\n",
    "        self.vars_bn.extend([running_mean, running_var])\n",
    "        \n",
    "        # 第2个conv2d\n",
    "        weight = nn.Parameter(torch.ones(64, 64, 3, 3))\n",
    "        nn.init.kaiming_normal_(weight)\n",
    "        bias = nn.Parameter(torch.zeros(64))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "        # 第2个BatchNorm层\n",
    "        weight = nn.Parameter(torch.ones(64))\n",
    "        bias = nn.Parameter(torch.zeros(64))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "        running_mean = nn.Parameter(torch.zeros(64), requires_grad= False)\n",
    "        running_var = nn.Parameter(torch.zeros(64), requires_grad= False)\n",
    "        self.vars_bn.extend([running_mean, running_var])\n",
    "        \n",
    "        # 第3个conv2d\n",
    "        weight = nn.Parameter(torch.ones(64, 64, 3, 3))\n",
    "        nn.init.kaiming_normal_(weight)\n",
    "        bias = nn.Parameter(torch.zeros(64))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "        # 第3个BatchNorm层\n",
    "        weight = nn.Parameter(torch.ones(64))\n",
    "        bias = nn.Parameter(torch.zeros(64))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "        running_mean = nn.Parameter(torch.zeros(64), requires_grad= False)\n",
    "        running_var = nn.Parameter(torch.zeros(64), requires_grad= False)\n",
    "        self.vars_bn.extend([running_mean, running_var])\n",
    "        \n",
    "        # 第4个conv2d\n",
    "        weight = nn.Parameter(torch.ones(64, 64, 3, 3))\n",
    "        nn.init.kaiming_normal_(weight)\n",
    "        bias = nn.Parameter(torch.zeros(64))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "        # 第4个BatchNorm层\n",
    "        weight = nn.Parameter(torch.ones(64))\n",
    "        bias = nn.Parameter(torch.zeros(64))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "        running_mean = nn.Parameter(torch.zeros(64), requires_grad= False)\n",
    "        running_var = nn.Parameter(torch.zeros(64), requires_grad= False)\n",
    "        self.vars_bn.extend([running_mean, running_var])\n",
    "        \n",
    "        ##linear\n",
    "        weight = nn.Parameter(torch.ones([20,64]))\n",
    "        bias = nn.Parameter(torch.zeros(20))\n",
    "        self.vars.extend([weight,bias])\n",
    "        \n",
    "    def forward(self, x, params = None, bn_training=True):\n",
    "        '''\n",
    "        :bn_training: set False to not update\n",
    "        :return: \n",
    "        '''\n",
    "        if params is None:\n",
    "            params = self.vars\n",
    "        \n",
    "        weight, bias = params[0], params[1]  # 第1个CONV层\n",
    "        x = F.conv2d(x, weight, bias, stride = 1, padding = 1)\n",
    "        weight, bias = params[2], params[3]  # 第1个BN层\n",
    "        running_mean, running_var = self.vars_bn[0], self.vars_bn[1]\n",
    "        x = F.batch_norm(x, running_mean, running_var, weight=weight,bias =bias, training= bn_training, momentum = 1)\n",
    "        x = F.relu(x, inplace = [True])  #第1个relu\n",
    "        x = F.max_pool2d(x,kernel_size=2)  #第1个MAX_POOL层  \n",
    "        \n",
    "        \n",
    "        \n",
    "        weight, bias = params[4], params[5]  # 第2个CONV层\n",
    "        x = F.conv2d(x, weight, bias, stride = 1, padding = 1)\n",
    "        weight, bias = params[6], params[7]  # 第2个BN层\n",
    "        running_mean, running_var = self.vars_bn[2], self.vars_bn[3]\n",
    "        x = F.batch_norm(x, running_mean, running_var, weight=weight,bias =bias, training= bn_training, momentum=1)\n",
    "        x = F.relu(x, inplace = [True])  #第2个relu\n",
    "        x = F.max_pool2d(x,kernel_size=2)  #第2个MAX_POOL层   \n",
    "        \n",
    "        \n",
    "        weight, bias = params[8], params[9]  # 第3个CONV层\n",
    "        x = F.conv2d(x, weight, bias, stride = 1, padding = 1)\n",
    "        weight, bias = params[10], params[11]  # 第3个BN层\n",
    "        running_mean, running_var = self.vars_bn[4], self.vars_bn[5]\n",
    "        x = F.batch_norm(x, running_mean, running_var, weight=weight,bias =bias, training= bn_training,momentum=1)\n",
    "        x = F.relu(x, inplace = [True])  #第3个relu,\n",
    "        x = F.max_pool2d(x,kernel_size=2)  #第3个MAX_POOL层\n",
    "        \n",
    "        \n",
    "        weight, bias = params[12], params[13]  # 第4个CONV层\n",
    "        x = F.conv2d(x, weight, bias, stride = 1, padding = 1)\n",
    "        weight, bias = params[14], params[15]  # 第4个BN层\n",
    "        running_mean, running_var = self.vars_bn[6], self.vars_bn[7]\n",
    "        x = F.batch_norm(x, running_mean, running_var, weight=weight,bias =bias, training= bn_training)\n",
    "        x = F.max_pool2d(x,kernel_size=2)  #第4个MAX_POOL层\n",
    "        \n",
    "        x = F.relu(x, inplace = [True])  #第4个relu\n",
    "        \n",
    "        x = x.view(x.size(0), -1) ## flatten\n",
    "        weight, bias = params[-2], params[-1]  # linear\n",
    "        x = F.linear(x, weight, bias)\n",
    "        \n",
    "        output = x\n",
    "        \n",
    "        return output\n",
    "    \n",
    "    \n",
    "    def parameters(self):\n",
    "        \n",
    "        return self.vars\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-29T12:00:30.197710Z",
     "start_time": "2020-02-29T12:00:30.186076Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-02-29T05:41:40.773998Z",
     "start_time": "2020-02-29T05:41:40.762077Z"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-14T03:27:13.001807Z",
     "start_time": "2020-04-14T03:27:12.969355Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "class MetaLearner(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MetaLearner, self).__init__()\n",
    "        self.update_step = 5 ## task-level inner update steps\n",
    "        self.update_step_test = 5\n",
    "        self.net = BaseNet()\n",
    "        self.meta_lr = 0.001\n",
    "        self.base_lr = 0.1\n",
    "        self.meta_optim = torch.optim.Adam(self.net.parameters(), lr = self.meta_lr)\n",
    "#         self.meta_optim = torch.optim.SGD(self.net.parameters(), lr = self.meta_lr, momentum = 0.9, weight_decay=0.0005)\n",
    "        \n",
    "    def forward(self,x_spt, y_spt, x_qry, y_qry):\n",
    "        # 初始化\n",
    "        task_num, ways, shots, h, w = x_spt.size()\n",
    "        query_size = x_qry.size(1) # 75 = 15 * 5\n",
    "        loss_list_qry = [0 for _ in range(self.update_step + 1)]\n",
    "        correct_list = [0 for _ in range(self.update_step + 1)]\n",
    "        \n",
    "        for i in range(task_num):\n",
    "            ## 第0步更新\n",
    "            y_hat = self.net(x_spt[i], params = None, bn_training=True) # (ways * shots, ways)\n",
    "            loss = F.cross_entropy(y_hat, y_spt[i]) \n",
    "            grad = torch.autograd.grad(loss, self.net.parameters())\n",
    "            tuples = zip(grad, self.net.parameters()) ## 将梯度和参数\\theta一一对应起来\n",
    "            # fast_weights这一步相当于求了一个\\theta - \\alpha*\\nabla(L)\n",
    "            fast_weights = list(map(lambda p: p[1] - self.base_lr * p[0], tuples))\n",
    "            # 在query集上测试，计算准确率\n",
    "            # 这一步使用更新前的数据\n",
    "            with torch.no_grad():\n",
    "                y_hat = self.net(x_qry[i], self.net.parameters(), bn_training = True)\n",
    "                loss_qry = F.cross_entropy(y_hat, y_qry[i])\n",
    "                loss_list_qry[0] += loss_qry\n",
    "                pred_qry = F.softmax(y_hat, dim=1).argmax(dim=1)  # size = (75)\n",
    "                correct = torch.eq(pred_qry, y_qry[i]).sum().item()\n",
    "                correct_list[0] += correct\n",
    "            \n",
    "            # 使用更新后的数据在query集上测试。\n",
    "            with torch.no_grad():\n",
    "                y_hat = self.net(x_qry[i], fast_weights, bn_training = True)\n",
    "                loss_qry = F.cross_entropy(y_hat, y_qry[i])\n",
    "                loss_list_qry[1] += loss_qry\n",
    "                pred_qry = F.softmax(y_hat, dim=1).argmax(dim=1)  # size = (75)\n",
    "                correct = torch.eq(pred_qry, y_qry[i]).sum().item()\n",
    "                correct_list[1] += correct   \n",
    "            \n",
    "            for k in range(1, self.update_step):\n",
    "                \n",
    "                y_hat = self.net(x_spt[i], params = fast_weights, bn_training=True)\n",
    "                loss = F.cross_entropy(y_hat, y_spt[i])\n",
    "                grad = torch.autograd.grad(loss, fast_weights)\n",
    "                tuples = zip(grad, fast_weights) \n",
    "                fast_weights = list(map(lambda p: p[1] - self.base_lr * p[0], tuples))\n",
    "                \n",
    "                if k < self.update_step - 1:\n",
    "                    with torch.no_grad():\n",
    "                        y_hat = self.net(x_qry[i], params = fast_weights, bn_training = True)\n",
    "                        loss_qry = F.cross_entropy(y_hat, y_qry[i])\n",
    "                        loss_list_qry[k+1] += loss_qry\n",
    "                else:\n",
    "                    y_hat = self.net(x_qry[i], params = fast_weights, bn_training = True)\n",
    "                    loss_qry = F.cross_entropy(y_hat, y_qry[i])\n",
    "                    loss_list_qry[k+1] += loss_qry\n",
    "                \n",
    "                with torch.no_grad():\n",
    "                    pred_qry = F.softmax(y_hat,dim=1).argmax(dim=1)\n",
    "                    correct = torch.eq(pred_qry, y_qry[i]).sum().item()\n",
    "                    correct_list[k+1] += correct\n",
    "#         print('hello')\n",
    "                \n",
    "        loss_qry = loss_list_qry[-1] / task_num\n",
    "        self.meta_optim.zero_grad() # 梯度清零\n",
    "        loss_qry.backward()\n",
    "        self.meta_optim.step()\n",
    "        \n",
    "        accs = np.array(correct_list) / (query_size * task_num)\n",
    "        loss = np.array(loss_list_qry) / ( task_num)\n",
    "        return accs,loss\n",
    "\n",
    "    \n",
    "    \n",
    "    def finetunning(self, x_spt, y_spt, x_qry, y_qry):\n",
    "        assert len(x_spt.shape) == 4\n",
    "        \n",
    "        query_size = x_qry.size(0)\n",
    "        correct_list = [0 for _ in range(self.update_step_test + 1)]\n",
    "        \n",
    "        new_net = deepcopy(self.net)\n",
    "        y_hat = new_net(x_spt)\n",
    "        loss = F.cross_entropy(y_hat, y_spt)\n",
    "        grad = torch.autograd.grad(loss, new_net.parameters())\n",
    "        fast_weights = list(map(lambda p:p[1] - self.base_lr * p[0], zip(grad, new_net.parameters())))\n",
    "        \n",
    "        # 在query集上测试，计算准确率\n",
    "        # 这一步使用更新前的数据\n",
    "        with torch.no_grad():\n",
    "            y_hat = new_net(x_qry,  params = new_net.parameters(), bn_training = True)\n",
    "            pred_qry = F.softmax(y_hat, dim=1).argmax(dim=1)  # size = (75)\n",
    "            correct = torch.eq(pred_qry, y_qry).sum().item()\n",
    "            correct_list[0] += correct\n",
    "\n",
    "        # 使用更新后的数据在query集上测试。\n",
    "        with torch.no_grad():\n",
    "            y_hat = new_net(x_qry, params = fast_weights, bn_training = True)\n",
    "            pred_qry = F.softmax(y_hat, dim=1).argmax(dim=1)  # size = (75)\n",
    "            correct = torch.eq(pred_qry, y_qry).sum().item()\n",
    "            correct_list[1] += correct\n",
    "\n",
    "        for k in range(1, self.update_step_test):\n",
    "            y_hat = new_net(x_spt, params = fast_weights, bn_training=True)\n",
    "            loss = F.cross_entropy(y_hat, y_spt)\n",
    "            grad = torch.autograd.grad(loss, fast_weights)\n",
    "            fast_weights = list(map(lambda p:p[1] - self.base_lr * p[0], zip(grad, fast_weights)))\n",
    "            \n",
    "            y_hat = new_net(x_qry, fast_weights, bn_training=True)\n",
    "            \n",
    "            with torch.no_grad():\n",
    "                pred_qry = F.softmax(y_hat, dim=1).argmax(dim=1)\n",
    "                correct = torch.eq(pred_qry, y_qry).sum().item()\n",
    "                correct_list[k+1] += correct\n",
    "                \n",
    "        del new_net\n",
    "        accs = np.array(correct_list) / query_size\n",
    "        return accs\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-14T03:27:15.093572Z",
     "start_time": "2020-04-14T03:27:15.090246Z"
    }
   },
   "outputs": [],
   "source": [
    "# net = torch.load('./trained_models/MTL-5000epochs.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2020-04-14T03:27:13.031Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0\n",
      "[0.05       0.08625    0.1075     0.13270833 0.161875   0.19520833]\n",
      "在mean process之前： (992, 6)\n",
      "测试集准确率: [0.04962 0.1088  0.1405  0.1738  0.2063  0.2367 ]\n"
     ]
    }
   ],
   "source": [
    "## omniglot\n",
    "import random\n",
    "random.seed(1337)\n",
    "np.random.seed(1337)\n",
    "\n",
    "import time\n",
    "device = torch.device('cuda:2')\n",
    "\n",
    "meta = MetaLearner().to(device)\n",
    "\n",
    "epochs = 60001\n",
    "for step in range(epochs):\n",
    "    start = time.time()\n",
    "    x_spt, y_spt, x_qry, y_qry = next('train')\n",
    "    x_spt, y_spt, x_qry, y_qry = torch.from_numpy(x_spt).to(device),\\\n",
    "                                 torch.from_numpy(y_spt).to(device),\\\n",
    "                                 torch.from_numpy(x_qry).to(device),\\\n",
    "                                 torch.from_numpy(y_qry).to(device)\n",
    "    accs,loss = meta(x_spt, y_spt, x_qry, y_qry)\n",
    "    end = time.time()\n",
    "    if step % 100 == 0:\n",
    "        print(\"epoch:\" ,step)\n",
    "        print(accs)\n",
    "#         print(loss)\n",
    "        \n",
    "    if step % 1000 == 0:\n",
    "        accs = []\n",
    "        for _ in range(1000//task_num):\n",
    "            # db_train.next('test')\n",
    "            x_spt, y_spt, x_qry, y_qry = next('test')\n",
    "            x_spt, y_spt, x_qry, y_qry = torch.from_numpy(x_spt).to(device),\\\n",
    "                                         torch.from_numpy(y_spt).to(device),\\\n",
    "                                         torch.from_numpy(x_qry).to(device),\\\n",
    "                                         torch.from_numpy(y_qry).to(device)\n",
    "\n",
    "            \n",
    "            for x_spt_one, y_spt_one, x_qry_one, y_qry_one in zip(x_spt, y_spt, x_qry, y_qry):\n",
    "                test_acc = meta.finetunning(x_spt_one, y_spt_one, x_qry_one, y_qry_one)\n",
    "                accs.append(test_acc)\n",
    "        print('在mean process之前：',np.array(accs).shape)\n",
    "        accs = np.array(accs).mean(axis=0).astype(np.float16)\n",
    "        print('测试集准确率:',accs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-01T03:00:56.266331Z",
     "start_time": "2020-03-01T03:00:56.205955Z"
    }
   },
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ML3.6",
   "language": "python",
   "name": "ml3.6"
  },
  "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.6.9"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}