Selector Network for semi-supervise learning (new version with Batchnorm).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[MLENS] backend: threading\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2 \n",
    "%reload_ext autoreload\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# always import gbm_algos first !\n",
    "import xgboost, lightgbm, catboost\n",
    "\n",
    "# Core\n",
    "import numpy as np\n",
    "import scipy.io as io\n",
    "from torch.autograd import grad\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "# Sklearn\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from mlens.ensemble import SuperLearner\n",
    "\n",
    "# Let's do facy optimizers\n",
    "from optimizers import Lookahead, AdamGC, SGDGC\n",
    "# Modify at /usr/local/lib/python3.9/site-packages/torch_lr_finder/lr_finder.py\n",
    "from torch_lr_finder import LRFinder\n",
    "from onecyclelr import OneCycleLR\n",
    "import pcgrad\n",
    "\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training with 2000 samples\n"
     ]
    }
   ],
   "source": [
    "DATA_PATH = \"/Users/pongpisit/Desktop/research/pinn/Solving-Differential-Equations-with-Neural-Networks/SymbolicMathematics/data/burgers_shock.mat\"\n",
    "data = io.loadmat(DATA_PATH)\n",
    "\n",
    "t = data['t'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "Exact = np.real(data['usol']).T\n",
    "\n",
    "X, T = np.meshgrid(x,t)\n",
    "\n",
    "X_star = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))\n",
    "u_star = Exact.flatten()[:,None]              \n",
    "\n",
    "# Doman bounds\n",
    "lb = X_star.min(0)\n",
    "ub = X_star.max(0)\n",
    "\n",
    "N = 2000\n",
    "print(f\"Training with {N} samples\")\n",
    "idx = np.random.choice(X_star.shape[0], N, replace=False)\n",
    "X_u_train = X_star[idx, :]\n",
    "u_train = u_star[idx,:]\n",
    "\n",
    "# Convert to torch.tensor\n",
    "X_u_train = torch.tensor(X_u_train).float().requires_grad_(True)\n",
    "u_train = torch.tensor(u_train).float().requires_grad_(True)\n",
    "X_star = torch.tensor(X_star).float().requires_grad_(True)\n",
    "u_star = torch.tensor(u_star).float().requires_grad_(True)\n",
    "\n",
    "feature_names=['uf', 'u_x',  'u_xx', 'u_tt', 'u_xt', 'u_tx']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Network(nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super(Network, self).__init__()\n",
    "        self.model = model\n",
    "#         self.model.apply(self.xavier_init)\n",
    "        # For tracking\n",
    "        self.index2features = ('uf', 'u_x',  'u_xx', 'u_tt', 'u_xt', 'u_tx')\n",
    "        self.uf = None\n",
    "        \n",
    "    def xavier_init(self, m):\n",
    "        if type(m) == nn.Linear:\n",
    "            torch.nn.init.xavier_uniform_(m.weight)\n",
    "            m.bias.data.fill_(0.01)\n",
    "        \n",
    "    def forward(self, x, t):\n",
    "        self.uf = self.model(torch.cat([x, t], dim=1))\n",
    "        return self.uf\n",
    "    \n",
    "    def get_selector_data(self, x, t):\n",
    "        uf = self.forward(x, t)\n",
    "        \n",
    "        ### PDE Loss calculation ###\n",
    "        # first-order derivatives\n",
    "        u_t = self.gradients(uf, t)[0]\n",
    "        u_x = self.gradients(uf, x)[0]\n",
    "        # Homo second-order derivatives\n",
    "        u_tt = self.gradients(u_t,t)[0]\n",
    "        u_xx = self.gradients(u_x, x)[0]\n",
    "        # Hetero second-order derivatives\n",
    "        u_xt = self.gradients(u_t, x)[0]\n",
    "        u_tx = self.gradients(u_x, t)[0]\n",
    "        \n",
    "        X_selector = torch.cat([uf, u_x, u_xx, u_tt, u_xt, u_tx], dim=1)\n",
    "        y_selector = u_t\n",
    "        \n",
    "        return X_selector, y_selector\n",
    "    \n",
    "    def gradients(self, func, x):\n",
    "        return grad(func, x, create_graph=True, retain_graph=True, grad_outputs=torch.ones(func.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Does the SeclectorNetwork has to be a neural networks ???\n",
    "class SeclectorNetwork(nn.Module):\n",
    "    def __init__(self, X_train_dim, bn=None):\n",
    "        super().__init__()\n",
    "        # Nonlinear model, Training with PDE reg.\n",
    "        self.nonlinear_model = TorchMLP(dimensions=[X_train_dim, 50, 50, 1], activation_function=nn.Tanh, bn=bn, dropout=nn.Dropout(p=0.1), inp_drop=False)\n",
    "        \n",
    "    def xavier_init(self, m):\n",
    "        if type(m) == nn.Linear:\n",
    "            torch.nn.init.xavier_uniform_(m.weight)\n",
    "            m.bias.data.fill_(0.01)\n",
    "        \n",
    "    def forward(self, inn):\n",
    "        ut_approx = self.nonlinear_model(inn)\n",
    "        return ut_approx\n",
    "    \n",
    "    def loss(self, X_input, y_input):\n",
    "        ut_approx = self.forward(X_input)\n",
    "        mse_loss = F.mse_loss(ut_approx, y_input, reduction='mean')\n",
    "        return mse_loss\n",
    "\n",
    "class SemiSupModel(nn.Module):\n",
    "    def __init__(self, network, selector, normalize_derivative_features):\n",
    "        super(SemiSupModel, self).__init__()\n",
    "        self.network = network\n",
    "        self.selector = selector\n",
    "        self.normalize_derivative_features = normalize_derivative_features\n",
    "    def forward(self, X_u_train):\n",
    "        inn = X_u_train\n",
    "        if self.normalize_derivative_features:\n",
    "            inn = minmax_normalize(inn)\n",
    "        unsup_loss = self.selector.loss(*self.network.get_selector_data(*dimension_slicing(inn)))\n",
    "        return self.network.uf, unsup_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# network = Network(model=TorchMLP(dimensions=[6, 50, 50, 50 ,50, 50, 1], bn=nn.BatchNorm1d))\n",
    "# selector = SeclectorNetwork(X_train_dim=6, bn=nn.LayerNorm)\n",
    "semisup_model = SemiSupModel(network=Network(model=TorchMLP(dimensions=[2, 50, 50, 50 ,50, 50, 1], activation_function=nn.Tanh, bn=nn.LayerNorm, dropout=None)),\n",
    "                             selector=SeclectorNetwork(X_train_dim=6, bn=nn.LayerNorm),\n",
    "                             normalize_derivative_features=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pcgrad_closure():\n",
    "    uf, unsup_loss = semisup_model(X_u_train)\n",
    "    losses = [F.mse_loss(uf, u_train), unsup_loss]\n",
    "    updated_grads = []\n",
    "    \n",
    "    for i in range(2):\n",
    "        optimizer1.zero_grad()\n",
    "        losses[i].backward(retain_graph=True)\n",
    "\n",
    "        g_task = []\n",
    "        for param in semisup_model.parameters():\n",
    "            if param.grad is not None:\n",
    "                g_task.append(Variable(param.grad.clone(), requires_grad=False))\n",
    "            else:\n",
    "                g_task.append(Variable(torch.zeros(param.shape), requires_grad=False))\n",
    "        # appending the gradients from each task\n",
    "        updated_grads.append(g_task)\n",
    "\n",
    "    updated_grads = list(pcgrad.pc_grad_update(updated_grads))[0]\n",
    "    for idx, param in enumerate(semisup_model.parameters()):\n",
    "        param.grad = (updated_grads[0][idx]+updated_grads[1][idx])\n",
    "        \n",
    "    return sum(losses)\n",
    "\n",
    "def closure():\n",
    "    optimizer2.zero_grad()\n",
    "    mse_loss = F.mse_loss(semisup_model.network(*dimension_slicing(X_u_train)), u_train)\n",
    "    mse_loss.backward(retain_graph=True)\n",
    "    return mse_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learning rate finding\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a8e9971cf6a4f9189b3b7072046cf51",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=300.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stopping early, the loss has diverged\n",
      "\n",
      "Learning rate search finished. See the graph with {finder_name}.plot()\n",
      "LR suggestion: steepest gradient\n",
      "Suggested LR: 7.24E-05\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYgAAAELCAYAAADDZxFQAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/Il7ecAAAACXBIWXMAAAsTAAALEwEAmpwYAAAuoUlEQVR4nO3deXxU9bnH8c+TPSEhgSQECEtYEvZNIghoQMUFXKjWjau1UitarUppvWpbq/Zq673aXq87aBH3StFaFFBsRVBAZd8X2YSQQEJC9j157h8z0hgDJGFOTiZ53q/XvMjMnDnznWHgm985Z35HVBVjjDGmrgC3AxhjjGmZrCCMMcbUywrCGGNMvawgjDHG1MsKwhhjTL2sIIwxxtTLsYIQke4islREtonIVhG5u55lJohIvohs8F5+51QeY4wxjRPk4LqrgF+q6joRiQLWisjHqrqtznKfqeqlDV1pXFycJiUl+TKnMca0emvXrj2qqvGNeYxjBaGqmUCm9+dCEdkOJAJ1C6JRkpKSWLNmjQ8SGmNM2yEi3zT2Mc2yD0JEkoARwJf13D1GRDaKyGIRGdQceYwxxpyak5uYABCRSOAdYIaqFtS5ex3QU1WLRGQy8B6QXM86pgPTAXr06OFsYGOMMYDDIwgRCcZTDm+o6rt171fVAlUt8v68CAgWkbh6lputqqmqmhof36hNaMYYY5rIsRGEiAjwF2C7qv75BMt0Bo6oqorIKDyFleNUJmPaqsrKStLT0ykrK3M7inFYWFgY3bp1Izg4+LTX5eQmpnHAj4DNIrLBe9uvgR4AqvoCcBXwMxGpAkqB69SmlzXG59LT04mKiiIpKQnP726mNVJVcnJySE9Pp1evXqe9PiePYvocOOknUVWfAZ5xKoMxxqOsrMzKoQ0QEWJjY8nOzvbJ+hzfSd2SVNco+aWV5BZXUFpRTURoIJGhQRSXV3Egt4SosCDiI8OIiwohIqRNvTWmDbByaBt8+ffcZv4X/GBTBne+tZ6GbsBqFxJIXFQo7UKCKK+qJjwkkK7R4QxJjCY8JJAAEQIEAgOEuMhQuneMoEdsBO3DTn+7nzGuU4Uvv4TMTOjSBUaPBgcK5sknn2T69OlERET4fN0NlZeXx5tvvsntt9/eLM/37Xe54uLiGDt2LCtXrmzSeubOncuFF15I165dfZzw39pMQfRLiOLOc/vSoV0IHdt5RgglFVUUlVcREhhAUlw7isqrOFpYTnZROUcLK8guKqe4vIqw4ADKKmv4OquIJduOnPR5okKD6BwdRteYcHp0jKBnbATdvX/26BhhIxPT8i1aBLfeCnl5EBAANTUQEwOzZsHkyT59qieffJIbbrjB9YJ47rnnTqsgqqqqCApq/L/tppYDeApi8ODBVhC+kJwQxcwL+532ekorqqmsqUFroEaVqhrlSEEZB3NLOHishIy8MjLzSzmUV8r6A8coKKv6zuPjo0Lp6R1tdI0OJzbSU1hxkaHERoYQ2y6UuMgQ2xxg3LFoEVx1FZSWfvf2oiLP7fPnN6kkiouLueaaa0hPT6e6upoHHniAI0eOkJGRwbnnnktcXBxLly5lyZIlPPjgg5SXl9OnTx9efvllIiMjWbt2LTNnzqSoqIi4uDjmzp1Lly5dmDBhAsOGDWPZsmVUVVUxZ84cRo0aRXFxMXfeeSdbtmyhsrKShx56iClTprB161amTZtGRUUFNTU1vPPOOzzwwAPs2bOH4cOHc8EFF/D4449/J/t//dd/8frrrxMfH0/37t0ZOXIkv/rVr5gwYQLDhw/n888/Z+rUqaSkpPDII49QUVFBbGwsb7zxBgkJCeTk5DB16lQOHTrEmDFjqH0cTmRkJEVFRQA8/vjjzJs3j/Lycq644goefvhh9u/fz6RJkzj77LNZuXIliYmJ/OMf/2DhwoWsWbOG66+/nvDwcFatWkV4eHjj/75PRVX96jJy5Ej1J8eKy3XDgWP6jw2H9Ol/7dJfzdug17ywUs/6wz+1130faM97v38Z/LsP9crnVuh972zUp/65SxdsOKS7swq1qrrG7Zdj/NS2bdtOvVBNjWpioqpnA1P9l27dPMs10vz58/WnP/3p8et5eXmqqtqzZ0/Nzs5WVdXs7Gw955xztKioSFVVH3vsMX344Ye1oqJCx4wZo1lZWaqq+te//lWnTZumqqrjx48/vt5ly5bpoEGDVFX1/vvv19dee01VVY8dO6bJyclaVFSkP//5z/X1119XVdXy8nItKSnRffv2HX9cXV999ZUOGzZMS0tLtaCgQPv27auPP/748ef+2c9+dnzZ3NxcrfG+Ny+++KLOnDlTVVXvvPNOffjhh1VV9YMPPlDg+Gtu166dqqp+9NFHesstt2hNTY1WV1frJZdcosuWLdN9+/ZpYGCgrl+/XlVVr7766uOva/z48bp69ep6c9f39w2s0Ub+f9tmRhBuiYkIISYihGHdY753X02NkldaSU5ROTnFFeQUVZBVWMae7CJ2HS7io61HyC2uOL58eHAgqUkdSEuOZ3y/eJI7RdpIw/jOl19Cfv7Jl8nLg6++8uyTaIQhQ4bwy1/+knvvvZdLL72Uc84553vLfPHFF2zbto1x48YBUFFRwZgxY9i5cydbtmzhggsuAKC6upouXbocf9zUqVMBSEtLo6CggLy8PJYsWcKCBQt44oknAM9RXAcOHGDMmDE8+uijpKenc+WVV5Kc/L2JG75jxYoVTJkyhbCwMMLCwrjsssu+c/+11157/Of09HSuvfZaMjMzqaioOH6Y6fLly3n3Xc/3hC+55BI6dOjwvedZsmQJS5YsYcSIEQAUFRXx9ddf06NHD3r16sXw4cMBGDlyJPv37z9pZl+ygnBRQIDQ0btP5EQf07LKavZkF7Eto4CtGQWs2H2URxdt59FF2+kSHcbYPnGM6tWBc/t3olNUWLPmN61MZqZnn8PJBARARkajV52SksK6detYtGgRv/3tbzn//PP53e++O7u/qnLBBRfw1ltvfef2zZs3M2jQIFatWlXvuuv+kiQiqCrvvPMO/fp9d7PygAEDGD16NAsXLmTy5MnMmjWL3r17N/r1fKtdu3bHf77zzjuZOXMml19+OZ9++ikPPfRQg9ejqtx///3ceuut37l9//79hIaGHr8eGBhIad3Nfw6yEwa1cGHBgQzqGs3Vqd156PJBfDxzPCvvO4/HrhzC8O4xLN2Zxb3vbGb0H/7FVc+v5IVle1i5+yhlldVuRzf+pksXzw7pk6mpgSbsFM3IyCAiIoIbbriBe+65h3Xr1gEQFRVFYWEhAGeddRYrVqxg9+7dgGe/xa5du+jXrx/Z2dnHC6KyspKtW7ceX/fbb78NwOeff050dDTR0dFcdNFFPP3008e3969fvx6AvXv30rt3b+666y6mTJnCpk2bvpOhrnHjxvH+++9TVlZGUVERH3zwwQlfY35+PomJiQC88sorx29PS0vjzTffBGDx4sUcO3bse4+96KKLmDNnzvH9EYcOHSIrK+uk7+nJcvuKjSD8UNeYcK4b1YPrRvVAVdl1pIiPth7mwy2HeWzxDsBzmO6kIV24aWwSgxOjXU5s/MLo0RAd7dkhfSIxMTBqVKNXvXnzZu655x4CAgIIDg7m+eefB2D69OlcfPHFdO3alaVLlzJ37lymTp1KeXk5AI888ggpKSnMnz+fu+66i/z8fKqqqpgxYwaDBnkmfw4LC2PEiBFUVlYyZ84cAB544AFmzJjB0KFDqampoVevXnzwwQfMmzeP1157jeDgYDp37syvf/1rOnbsyLhx4xg8eDCTJk36zk7qM888k8svv5yhQ4eSkJDAkCFDiI6u/9/TQw89xNVXX02HDh0477zz2LdvHwAPPvggU6dOZdCgQYwdO7beCUcvvPBCtm/fzpgxYwDPzuvXX3+dwMDAE76nN910E7fddpujO6nl24b1F6mpqWrngzixnKJyNh3KZ/HmTD7YlElJRTVnJnXgJ+N6ccHABIICbdDYFm3fvp0BAwacesETHcUEEB7e5KOYnDJhwgSeeOIJUlNTHXuOoqIiIiMjKSkpIS0tjdmzZ3PGGWc49ny+UN/ft4isVdVGvVE2gmhlYiNDObdfJ87t14nfXDKQeasP8sqq/fzsjXUkxoTz47E9uTa1B9ER9oU+U4/Jkz0l0Ezfg/AH06dPZ9u2bZSVlfHjH/+4xZeDL9kIog2orlE+3naEl1fs48t9uYQHB3LVyG7cNC6JPvGRbsczzaDBI4hvqXqOVsrI8OxzGDXKkW9SG2fYCMI0WGCAcPHgzlw8uDNbM/J5ecV+3l59kNe++IbxKfHcmtabsX2/dxoO05aJNPpQVtP62AbpNmZQ12ieuHoYK+47j19MTGFbZgH/8dKX/PSVNew/Wux2POMgf9taYJrGl3/PVhBtVHxUKHdPTOaz/zyXey/uz6o9R7nwf5fz2OIdFJVXnXoFxq+EhYWRk5NjJdHKqfd8EGFhvvlOlO2DMABkFZTxPx/tZP7adOIiQ5kxMZlrz+xOsB311CrYGeXajhOdUa4p+yCsIMx3bDiYx6MLt7F6/zH6d47i0SsGM7JnR7djGWNOU1MKwn49NN8xvHsM824dwws3jKSgtJIfPr+K+9/dREFZpdvRjDHNzArCfI+I56inj2eOZ3pab+atSeeHz63kYG6J29GMMc3ICsKcULvQIH49eQCv3TyKIwVl/ODZFaz95vvzyBhjWicrCHNKY/vE8fc7xhEZFsTUF7/gwy2ZbkcyxjQDKwjTIH3iI3nv9nEM7tqe299Yx1tfHXA7kjHGYVYQpsE6tAvh9Z+OJi0lnvvf3cwzn3xtx9Ub04pZQZhGiQgJ4sUbU7liRCJPLNnFw+9vo6bGSsKY1sjmYjKNFhwYwJ+uHkZsuxBe+nwfucUVPHH1MEKC7PcNY1oTKwjTJAEBwm8uGUBcVCiPLd5BfmklL96YaiVhTCti/5pNk4kIt43vwx+vHMKyXdncM3+jbW4yphWxEYQ5bVNH9SC3uILHP9pJ5/Zh3D+5EecdMMa0WFYQxidun9CHw/llzFq+l4T2Yfzk7F5uRzLGnCYrCOMTIsJDlw8iu7Cc/1q4jU7tQ7l0aFe3YxljToPtgzA+ExggPHndcFJ7dmDm2xtZd8Cm5TDGn1lBGJ8KCw7kxRtT6Rwdxu2vr+NoUbnbkYwxTWQFYXwuJiKE5284g2MlFdz55nqqqmvcjmSMaQIrCOOIQV2jefSKIazam8PjS3a6HccY0wSOFYSIdBeRpSKyTUS2isjd9SwjIvKUiOwWkU0icoZTeUzzu2pkN64f3YNZy/ayeLPNAGuMv3FyBFEF/FJVBwJnAXeIyMA6y0wCkr2X6cDzDuYxLvjdZQMZ1j2GX/1tI3uyi9yOY4xpBMcKQlUzVXWd9+dCYDuQWGexKcCr6vEFECMiXZzKZJpfaFAgz19/BqHBgfz8zfWUVVa7HckY00DNsg9CRJKAEcCXde5KBA7Wup7O90vE+LmuMeE8ftVQtmcW8N8f7nA7jjGmgRwvCBGJBN4BZqhqQRPXMV1E1ojImuzsbN8GNM3i/AEJ3DQ2iZdX7OeTHUfcjmOMaQBHC0JEgvGUwxuq+m49ixwCute63s1723eo6mxVTVXV1Pj4eGfCGsfdN6k//TtH8au/bSKroMztOMaYU3DyKCYB/gJsV9U/n2CxBcCN3qOZzgLyVdUOd2mlwoIDeeY/RlBcXsWv/77ZzkZnTAvn5AhiHPAj4DwR2eC9TBaR20TkNu8yi4C9wG7gReB2B/OYFqBvpyjuuagf/9yexbvrvjdYNMa0II5N1qeqnwNyimUUuMOpDKZlmjauFx9tPcxD729lbN9YukSHux3JGFMP+ya1aXaBAcLjVw2jqlq59x3b1GRMS2UFYVyRFNeO+yb1Z/mubP6+3jY1GdMSWUEY1/zorJ6c0SOGRxZuJ7e4wu04xpg6rCCMawIChD9cOYSC0kr+uGi723GMMXVYQRhX9e/cnlvSevO3tems2pPjdhxjTC1WEMZ1d52XTPeO4fzm75spr7K5moxpKawgjOvCQwJ55AdD2Hu0mBeX73U7jjHGywrCtAjjU+KZPKQzT3+ym4O5JW7HMcZgBWFakAcuHUiACI/ZjK/GtAhWEKbF6BIdzvS03izclMnab465HceYNs8KwrQo09N6Ex8Vyh8WbbdvWBvjMisI06K0Cw3ilxeksPabY3y09bDbcYxp06wgTItzdWp3UhIieWzxDiqqatyOY0ybZQVhWpzAAOH+SQPYn1PCvDUHT/0AY4wjrCBMizShXzwje3bg2aW77ctzxrjECsK0SCLCLyamkJlfxrzVNoowxg1WEKbFGtc3ljOTOvDs0j2UVdoowpjmZgVhWqxvRxGHC8r461cH3I5jTJtjBWFatDF9YhnVqyPPfWqjCGOamxWEadG+HUVkFZbz5pc2ijCmOVlBmBZvTJ9YxvSO5fllNoowpjlZQRi/8IsLUsguLOf1L75xO4oxbYYVhPELo3p1ZFzfWF5YtoeSiiq34xjTJlhBGL/xi4kpHC2qsFGEMc3ECsL4jdSkjpyTHMesZXttFGFMM7CCMH5lxsQUcooreHWVjSKMcZoVhPErI3t2IC0lntnL91JcbqMIY5xkBWH8zi8mJpNbXMErq/a7HcWYVs0KwvidET06cG4/zyiisKzS7TjGtFpWEMYvzZiYQl5JJa+s3O92FGNaLSsI45eGdY/h/P6dePGzfRTYKMIYR1hBGL81Y2IK+aWVzF2x3+0oxrRKVhDGbw3pFs3EAQm89Nle8kttFGGMr1lBGL82Y2IyBWVV9u1qYxzgWEGIyBwRyRKRLSe4f4KI5IvIBu/ld05lMa3X4MRoJvSL5y+f7aXssxXw97/DF1+AqtvRjPF7QQ6uey7wDPDqSZb5TFUvdTCDaQN+zT4in/g5gY+VQnAg1NRATAzMmgWTJ7sdzxi/5dgIQlWXA7lOrd8YABYtIuWOaXQtPEpwaTEUFEBREaSnw1VXwaJFbic0xm+5vQ9ijIhsFJHFIjLoRAuJyHQRWSMia7Kzs5szn2nJVGH6dCgtrf/+0lK49Vbb3GRME7lZEOuAnqo6DHgaeO9EC6rqbFVNVdXU+Pj45spnWrovv4T8/JMvk5cHX33VLHGMaW1cKwhVLVDVIu/Pi4BgEYlzK4/xQ5mZEHCKj3BAAGRkNE8eY1oZ1wpCRDqLiHh/HuXNkuNWHuOHunTx7JA+mZoa6Nq1efIY08o4dhSTiLwFTADiRCQdeBAIBlDVF4CrgJ+JSBVQClynahuLTSOMHg3R0Z6d0icSEwOjRjVbJGNaE8cKQlWnnuL+Z/AcBmtM04jA7Nmeo5Xq2VGt4eHIrFme5Ywxjeb2UUzGnJ7Jk2H+fOjWDSIjoX17KsPbkREVx46n59j3IIw5DU5+Uc6Y5jF5Mhw44DlaKSODmk4JXP5JMf21Pa+7nc0YP2YFYVoHEc8+CSAUuKVqD39cvIMNB/MY3j3G1WjG+CvbxGRapevP6kl0eDDPLt3tdhRj/JYVhGmVIkODuGlsEh9vO8LOw4VuxzHGL1lBmFbrprFJRIQE8tynNoowpikaVBAi0k5EArw/p4jI5SIS7Gw0Y05Ph3Yh3HBWT97fmME3OcVuxzHG7zR0BLEcCBORRGAJ8CM803kb06L99OxeBAUG8MKyPW5HMcbvNLQgRFVLgCuB51T1auCEs68a01J0ah/GNandmL82ncz8E8z6aoypV4MLQkTGANcDC723BToTyRjfujWtDzUKLy7f53YUY/xKQwtiBnA/8HdV3SoivYGljqUyxoe6d4xgyvCuvPnVN+QUlbsdxxi/0aCCUNVlqnq5qv63d2f1UVW9y+FsxvjM7RP6UF5Vw8sr9rsdxRi/0dCjmN4UkfYi0g7YAmwTkXucjWaM7/TtFMXFgzrzyqr9FJRVuh3HGL/Q0E1MA1W1APgBsBjohedIJmP8xu0T+lJYVsVrq75xO4oxfqGhBRHs/d7DD4AFqloJ2LkbjF8Z0i2aCf3ieemzvRSXV7kdx5gWr6EFMQvYD7QDlotIT6DAqVDGOOXu85M5VlLJqzaKMOaUGrqT+ilVTVTVyerxDXCuw9mM8bkRPTqQlhLPizaKMOaUGrqTOlpE/iwia7yXP+EZTRjjd+4+P5nc4gpe/8JGEcacTEM3Mc0BCoFrvJcC4GWnQhnjpJE9O3BOchyzl++lpMJGEcacSEMLoo+qPqiqe72Xh4HeTgYzxkl3n59MTnEFb3xxwO0oxrRYDS2IUhE5+9srIjIOsIltjN9KTerI2D6xzFq+h9KKarfjGNMiNbQgbgOeFZH9IrIfeAa41bFUxjSDu89P5mhRBW98afsijKlPQ49i2qiqw4ChwFBVHQGc52gyYxw2uncsY3rH8sIy2xdhTH0adUY5VS3wfqMaYKYDeYxpVjMvTOFoUTmvrLRRhDF1nc4pR8VnKYxxyZlJHRmfEs8Ly/bYHE3G1HE6BWFTbZhW4VcX9iO/tJKXPrPzRRhT20kLQkQKRaSgnksh0LWZMhrjqCHdopk0uDN/+WwvucUVbscxpsU4aUGoapSqtq/nEqWqQc0V0hinzbwghZLKajt3tTG1nM4mJmNajeSEKK4YnsgrK/dzpKDM7TjGtAhWEMZ4zZiYQnWN8swnu92OYkyLYAVhjFeP2AiuObM7b311gG9yit2OY4zrrCCMqeXu85MJDgzgfz7a6XYUY1xnBWFMLQntw5ie1puFmzJZ+80xt+MY4yrHCkJE5ohIlohsOcH9IiJPichuEdkkImc4lcWYxpie1pv4qFAeXbgNVfu6j2m7nBxBzAUuPsn9k4Bk72U68LyDWYxpsHahQfzyghTWHchj8ZbDbscxxjWOFYSqLgdyT7LIFOBV7ylMvwBiRKSLU3mMaYyrU7vTLyGKxxbvoKzSpgM3bZOb+yASgYO1rqd7bzPGdYEBwu8uG8iB3BJmL9/rdhxjXOEXO6lFZPq358POzs52O45pI8b1jeOSoV14duluDuaWuB3HmGbnZkEcArrXut7Ne9v3qOpsVU1V1dT4+PhmCWcMwG8vGUBggPDw+9vcjmJMs3OzIBYAN3qPZjoLyFfVTBfzGPM9XaLDufv8ZP65/Qj/2n7E7TjGNCsnD3N9C1gF9BORdBG5WURuE5HbvIssAvYCu4EXgdudymLM6Zg2rhd9O0Xy4IKtduY506Y4NiOrqk49xf0K3OHU8xvjKyFBAfzhiiFcM2sVf1qyiwcuHeh2JGOahV/spDbGbaN6deT60T14ecU+NhzMczuOMc3CCsKYBrpvUn86RYVx7/xNVFTVuB3HGMdZQRjTQFFhwTzyg8HsPFLI85/aiYVM62cFYUwjTByYwJThXXn6k6/ZnJ7vdhxjHGUFYUwj/f7ywcRFhjLj7fU2DYdp1awgjGmk6Ihgnrh6GHuyi3ls8Q634xjjGCsIY5rg7OQ4po1LYu7K/Xz2tU3/YlonKwhjmujei/vTt1Mkv/rbRvJKKtyOY4zPWUEY00RhwYE8ee1wcooqeOAfW92OY4zPWUEYcxoGJ0bziwtSeH9jBv/YUO9ck8b4LSsIY07TrWm9GdmzA799bwsZeaVuxzHGZ6wgjDlNQYEB/PmaYdTUKHe9tZ6qavuWtWkdrCCM8YGese34w5VDWPPNMf708S634xjjE1YQxvjIlOGJTB3Vg+c/3cOnO7PcjmPMabOCMMaHHrxsIP07RzFz3kbbH2H8nhWEMT4UFhzIs9efQWVVDdNfW0NphU3FYfyXFYQxPtYnPpL/mzqcrRkF3DN/I55zYxnjf6wgjHHAef0T+M+L+vPBpkyes6nBjZ9y7JSjxrR1t43vzY7DBTz+0U56x7Vj0pAubkcyplFsBGGMQ0SE//7hUM7oEcOMtzew/sAxtyMZ0yhWEMY4KCw4kBdvTCWhfRi3vLqGg7klbkcypsGsIIxxWGxkKHNuOpOKqhpuevkrcorK3Y5kTINYQRjTDPp2iuTFG1NJP1bKjXO+Ir+00u1IxpySFYQxzWR071hm/Wgku44U8pO5qympqHI7kjEnZQVhTDOa0K8TT103gvUHjjHt5dUUlVtJmJbLCsKYZjZpSBf+99rhrPnmGP/x4hfkFtvZ6EzLZAVhjAumDE9k1g0j2XG4kGtnreJwfpnbkYz5HisIY1wycWACr0wbRUZeKVfPWsk3OcVuRzLmO6wgjHHRmD6xvDX9LIrKqrjqhVXsPFzodiRjjrOCMMZlQ7vFMO/WMQQIXDNrFWu/yXU7kjGAFYQxLUJyQhTzbxtLh4hgps7+knfXpbsdyRgrCGNaiu4dI3jvjnGM7NmBmfM28j8f7qCmxqYKN+6xgjCmBYmJCOHVm0cxdVR3nvt0Dz97Y619oc64xgrCmBYmODCAP1wxhAcuHcjH244w5ZkVbM8scDuWaYMcLQgRuVhEdorIbhG5r577bxKRbBHZ4L381Mk8xvgLEeHms3vxyk9GcaykkinPruDlFfvs7HSmWTlWECISCDwLTAIGAlNFZGA9i76tqsO9l5ecymOMPzonOZ6PZpzDOX3jePj9bUybu5rsQpsN1jQPJ0cQo4DdqrpXVSuAvwJTHHw+Y1ql2MhQXvpxKr+fMohVe3KY9H/LWbojy+1Ypg1wsiASgYO1rqd7b6vrhyKySUTmi0h3B/MY47dEhBvHJPH+nWcTFxnKtLmreWjBVsoqq92OZloxt3dSvw8kqepQ4GPglfoWEpHpIrJGRNZkZ2c3a0BjWpKUhCjeu2Mc08YlMXflfi5/5nM2HsxzO5ZppZwsiENA7RFBN+9tx6lqjqp+u0H1JWBkfStS1dmqmqqqqfHx8Y6ENcZfhAUH8uBlg5g77UwKSqu44rkVPLZ4h40mjM85WRCrgWQR6SUiIcB1wILaC4hIl1pXLwe2O5jHmFZlQr9OLJmZxjWp3Xlh2R4ufnI5S3favgnjO44VhKpWAT8HPsLzH/88Vd0qIr8Xkcu9i90lIltFZCNwF3CTU3mMaY3ahwXz2A+H8vrNowkQYdrLq5n+6hoO5pa4Hc20AuJvx1WnpqbqmjVr3I5hTItTXlXNXz7fx9P/2o2i3DGhL7ek9SYsONDtaKYFEJG1qpramMe4vZPaGOMjoUGB3D6hL//85XjO69+JP328i4ufXM77GzNsTifTJFYQxrQyiTHhPHf9SF79ySiCAwO48631XPx/y1m4KdOKwjSKFYQxrVRaSjwfzkjjqakjqK5R7nhzHZOf+ozFm60oTMPYPghj2oDqGuWDTRn837++Zm92Mf07R3Hb+D5cMrQLwYH2e2Jb0JR9EFYQxrQh1TXKgo2HeG7pHr7OKiIxJpybz+7FdaO6ExES5HY84yArCGNMg9TUKEt3ZvHCsj2s3n+MmIhgbjyrJzeOTSIuMtTteMYBVhDGmEZb+00us5btZcm2I4QEBXDliER+cnYvUhKi3I5mfMgKwhjTZLuzinh5xT7eWZdOWWUNo3t15IdndGPSkM5EhQW7Hc+cJisIY8xpyy2u4K2vDjB/bTr7jhYTFhzARYM688MzujGubxyBAeJ2RNMEVhDGGJ9RVdYfzOPddem8vzGT/NJKEtqH8oPhifxwZDfbBOVnrCCMMY4or6rmk+1ZvLMunU93ZlNVowxObM+VI7oxZXhXYm3HdotnBWGMcdzRonLe35jBu+sOsflQPkEBQlpKPJcO7cLEgQm0t/0VLZIVhDGmWe08XOjdBJVBRn4ZIYEBpKXEc8nQzkwckGA7t1sQKwhjjCtqapQN6Xks3JTJos2ZZOaXERIUQFpyHBcP7sIFAxKIjrCycJMVhDHGdTU1np3bCzdl8uGWTDLyywgKEMb0ieXCgQmcnRxPUmwEInY0VENVVtfwyY4senSMYECX9k1ahxWEMaZFUVU2peezeMthPtySyf4cz4mMEmPCSU3qQEpCFGcmdWRIYjThIW3rvBWH88tYviubz3cfpVqVXrHtmNAvnoiQIBQlNCiQDQfz+GJvDp/uzOZoUTk3junJ76cMbtLzWUEYY1osVWV/Tgmff+35T3HLoQIO5ZUCIALdOoTTKy6SnKJyCsuqiIkIpndcO/p1bk+/zpGkJESRGBPu1yOPrMIy7p2/iV1Hio6/9k5RobQLDeJgbglV9cyy2yEimLF94rjyjETGp8QT1MTJFa0gjDF+Ja+kgtX7j7E9s4BdRwrZd7SY2MhQOkQEk1tcwZ6sIjLyy44vHxkaxMAu7RnVqyNDu0UzsGt7vyqN215byyc7s7h0SBf6dY5ifL94+iVEISLkl1ayak8OoKhCSUU1gxOjSe4USYAPvpxoBWGMaXXySyv5+kghO48UsutwIRvS89lyKJ9q72/bMRHBjOsbx7n9OjE+JZ74qJb3nYzdWUW8sy6d5z/dw39e3I/bJ/Rt9gxNKQib39cY06JFhweTmtSR1KSOx28rqahi5+FCtmUWsOFAHst2ZbNwUyYAQxKjObdfPBP6d2JYtxjXpwbJLa7gsqc/p7SymrSUeKaf09vVPI1hIwhjjN9TVbZlFvDpzmyW7shi3YFj1Khn+31aSjzn9utEWko8HduFNHu299YfYsbbG3h7+lmM7h3b7M//LRtBGGPaJBFhUNdoBnWN5o5z+5JXUsFnXx9l6c4slu3M5h8bMhCBYd1iSEuJZ2yfWEb0iCE0yPkjpz7ZkUVcZAhn1hoB+QsrCGNMqxMTEcJlw7py2bCu1NQoWzLyWbojm6U7s3jmk6956l9fExoUQGpSB8b2iWNMn1iGJkY3+QihE6muUZbtymbigASf7GhublYQxphWLSBAGNothqHdYrh7YjIFZZV8tTeXlXtyWLU3h8c/2gl4jpAa1asjY/vEMqZPLAM6tz/t/9Q3HDxGfmkl5/Xv5IuX0uysIIwxbUr7sGAmDkxg4sAEAHKKyvlyXy4r9xxl5Z4cPtmRBXiOjjqrVyxj+8Yytk8sfeIjG3047fsbMwkOFM5OjvP562gOVhDGmDYtNjKUyUO6MHlIF8DzDedVe4+ycncOK/fk8OHWwwDER4UypncsZ/SIYXiPDgzoEnXSfRjF5VW8szadyUO6EB3un/NQWUEYY0wtnaPDuGJEN64Y0Q1V5WBu6fHRxRd7c1iwMQOAkMAABnRtz4juMQz3Xnp0jDi+Weq9DYcoLK/ixjFJLr6a02OHuRpjTAOpKpn5ZWw4mHf8sjk9n9LKagDCgwPp06kdnaLCWL0/lx4dI/jgzrNbxDe97TBXY4xxkIjQNSacrjHhxzdJVVXXsOtIEZvS89h1pIjd2UWkHyvh/P6duHV8nxZRDk1lBWGMMachKDCAgV3bM7Br06bhbsl8e9CvMcaYVsMKwhhjTL2sIIwxxtTLCsIYY0y9HC0IEblYRHaKyG4Rua+e+0NF5G3v/V+KSJKTeYwxxjScYwUhIoHAs8AkYCAwVUQG1lnsZuCYqvYF/hf4b6fyGGOMaRwnRxCjgN2quldVK4C/AlPqLDMFeMX783zgfPHng4aNMaYVcbIgEoGDta6ne2+rdxlVrQLyge+dUUNEpovIGhFZk52d7VBcY4wxtfnFF+VUdTYwG0BEskWkGDjahFVF4ymhpixT9/aTXa/v52//jKNp2U+W7VT3n2722j83Nb+vsp8sY+3r/pC97nUnPjf2mT/19bbyme/Z0ODHqaojF2AM8FGt6/cD99dZ5iNgjPfnIO8LkQase00TM81u6jJ1bz/Z9fp+rvVnk7I3JL9T2X2R31fZT5axvve7JWdvjs+NfebtM9/U915VHd3EtBpIFpFeIhICXAcsqLPMAuDH3p+vAj5R76tyyPunsUzd2092vb6fG/Lcp3KqdTiVvSHPfSq+yl73thO9Fn/IXve6E58b+8yf+rp95k/A0dlcRWQy8CQQCMxR1UdF5Pd4Wm2BiIQBrwEjgFzgOlXd24D1rtFGzkrYUvhzdvDv/JbdHf6cHfw7/+lmd3QfhKouAhbVue13tX4uA65uwqpnn2Y0N/lzdvDv/JbdHf6cHfw7/2ll97vzQRhjjGkeNtWGMcaYellBGGOMqZcVhDHGmHq1uoIQkXNE5AUReUlEVrqdpzFEJEBEHhWRp0Xkx6d+RMshIhNE5DPvez/B7TxNISLtvN/Yv9TtLI0hIgO87/t8EfmZ23kaQ0R+ICIveiftvNDtPI0hIr1F5C8iMt/tLA3l/Yy/4n3Prz/V8i2qIERkjohkiciWOrefdFbY2lT1M1W9DfiAf8/z5DhfZMczN1U3oBLP1CTNwkfZFSgCwmjG7OCz/AD3AvOcSVk/H33mt3s/89cA45zMW5uPsr+nqrcAtwHXOpm3Nh9l36uqNzub9NQa+VquBOZ73/PLT7ny0/mWna8vQBpwBrCl1m2BwB6gNxACbMQzO+wQPCVQ+9Kp1uPmAVH+lB24D7jV+9j5fpY9wPu4BOANf/vcABfg+TLnTcCl/pTd+5jLgcXAf/hbdu/j/gSc4afZm+3fqg9ey/3AcO8yb55q3S1qLiZVXV7POSGOzwoLICJ/Baao6h+BejcFiEgPIF9VC53MW5svsotIOlDhvVrtYNzv8NX77nUMCHUk6An46L2fALTD84+oVEQWqWqNk7nBd++9qi4AFojIQuBNByPXfk5fvO8CPAYsVtV1Dkc+zsefeVc15rXgGd13AzbQgC1ILaogTqC+WWFHn+IxNwMvO5ao4Rqb/V3gaRE5B1juZLAGaFR2EbkSuAiIAZ5xNFnDNCq/qv4GQERuAo42RzmcRGPf+wl4Nh2EUueLqS5o7Gf+TmAiEC0ifVX1BSfDnUJj3/dY4FFghIjc7y2SluJEr+Up4BkRuYQGTMfhDwXRaKr6oNsZmkJVS/CUm99R1XfxFJxfU9W5bmdoLFX9FPjU5RhNoqpP4flPy++oag6efSd+Q1WLgWkNXb5F7aQ+gUNA91rXu3lv8weW3T3+nN+yu8Ofs9flk9fiDwXRkFlhWyrL7h5/zm/Z3eHP2evyzWtxc+97PXvj3wIy+fdhnjd7b58M7MKzV/43bue07C3r4s/5Lbtlb8mvxSbrM8YYUy9/2MRkjDHGBVYQxhhj6mUFYYwxpl5WEMYYY+plBWGMMaZeVhDGGGPqZQVhWg0RKWrm52vW842ISIyI3N6cz2naNisIY05ARE46V5mqjm3m54wBrCBMs7GCMK2aiPQRkQ9FZK14znjX33v7ZSLypYisF5F/ikiC9/aHROQ1EVkBvOa9PkdEPhWRvSJyV611F3n/nOC9f76I7BCRN7zTWCMik723rRWRp0Tkg3oy3iQiC0TkE+BfIhIpIv8SkXUisllEpngXfQzoIyIbRORx72PvEZHVIrJJRB528r00bU+rnM3VmFpmA7ep6tciMhp4DjgP+Bw4S1VVRH4K/CfwS+9jBgJnq2qpiDwE9AfOBaKAnSLyvKpW1nmeEcAgIANYAYwTkTXALCBNVfeJyFsnyXkGMFRVc72jiCtUtUBE4oAvRGQBnhNKDVbV4QDiOUVnMp65/wXP+SDSVNXtqeJNK2EFYVotEYkExgJ/8/5CD/8+mVE34G0R6YLnjFv7aj10gaqW1rq+UFXLgXIRycJz1ry6p1X9SlXTvc+7AUjCcwrWvar67brfAqafIO7Hqpr7bXTgDyKSBtTgmds/oZ7HXOi9rPdej8RTGFYQxiesIExrFgDkffsbdx1PA39W1QXeE+48VOu+4jrLltf6uZr6/900ZJmTqf2c1wPxwEhVrRSR/XjO9V2XAH9U1VmNfC5jGsT2QZhWS1ULgH0icjV4Tm8pIsO8d0fz7/nxf+xQhJ1A71qng7y2gY+LBrK85XAu0NN7eyGezVzf+gj4iXekhIgkikin049tjIeNIExrEiGe83p/6894fht/XkR+CwQDf8VzAveH8Gx6OgZ8AvTydRjvPozbgQ9FpBjPHP0N8QbwvohsBtYAO7zryxGRFSKyBc85nO8RkQHAKu8mtCLgBiDL16/FtE023bcxDhKRSFUt8h7V9Czwtar+r9u5jGkI28RkjLNu8e603opn05HtLzB+w0YQxhhj6mUjCGOMMfWygjDGGFMvKwhjjDH1soIwxhhTLysIY4wx9bKCMMYYU6//B0X8kpMQ3RRyAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "params = semisup_model.parameters()\n",
    "\n",
    "### For SGD and Adam ###\n",
    "learning_rate1, learning_rate2 = 1e-7, 1e-1\n",
    "\n",
    "### For LBFGS (a good choice already!!!) ###\n",
    "# print(\"Using LBFGS's learning rate set\")\n",
    "# learning_rate1, learning_rate2 = 8e-2, 5e-2 # (1e-1, 5e-2) is also OK!\n",
    "\n",
    "choice = 'Adam'; auto_lr = True\n",
    "if choice == 'LBFGS':\n",
    "    optimizer1 = torch.optim.LBFGS(params, lr=learning_rate1, \n",
    "                                   max_iter=100, max_eval=125, \n",
    "                                  history_size=120, line_search_fn='strong_wolfe')\n",
    "if choice == 'Adam':\n",
    "    optimizer1 = AdamGC(params, lr=learning_rate1, use_gc=True, gc_conv_only=False, gc_loc=False)\n",
    "if choice == 'SGD':\n",
    "    optimizer1 = SGDGC(params, lr=learning_rate1, use_gc=True, nesterov=True, momentum=0.95)\n",
    "\n",
    "if choice != 'LBFGS' and auto_lr:\n",
    "    print('Learning rate finding')\n",
    "    bs = 4000; bs = X_u_train.shape[0] if bs>X_u_train.shape[0] else bs\n",
    "    criterion = LadderLoss(return_list=True)\n",
    "    trainloader = get_dataloader(X_u_train, u_train, bs=bs)\n",
    "    \n",
    "    lr_finder = LRFinder(semisup_model, optimizer=optimizer1, \n",
    "                         closure=pcgrad_update, criterion=criterion, device=\"cpu\")\n",
    "    lr_finder.range_test(trainloader, val_loader=None, end_lr=100, num_iter=300)\n",
    "    \n",
    "    # to inspect the loss-learning rate graph\n",
    "    suggested_lr, _ = lr_finder.plot()\n",
    "    # To prevent divergence during the second stage training.\n",
    "    # suggested_lr = min(suggested_lr, 5e-3)\n",
    "    lr_finder.reset(); plt.show()\n",
    "else:\n",
    "    suggested_lr = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the learing_rate to the suggested one.\n",
    "# suggested_lr = float(input())\n",
    "\n",
    "if suggested_lr:\n",
    "    optimizer1 = lr_finder.optimizer\n",
    "    for g in optimizer1.param_groups:\n",
    "        g['lr'] = suggested_lr\n",
    "        \n",
    "epochs1 = 2000; epochs2 = 500;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  2.15313982963562\n",
      "Epoch 100:  0.2881850600242615\n",
      "Epoch 200:  0.24482086300849915\n",
      "Epoch 300:  0.23417049646377563\n",
      "Epoch 400:  0.2043256163597107\n",
      "Epoch 500:  0.18732070922851562\n",
      "Epoch 600:  0.17506664991378784\n",
      "Epoch 700:  0.15822215378284454\n",
      "Epoch 800:  0.15274736285209656\n",
      "Epoch 900:  0.13736051321029663\n",
      "Epoch 1000:  0.12564720213413239\n",
      "Epoch 1100:  0.1172359436750412\n",
      "Epoch 1200:  0.10317790508270264\n",
      "Epoch 1300:  0.09375813603401184\n",
      "Epoch 1400:  0.08131937682628632\n",
      "Epoch 1500:  0.07160459458827972\n",
      "Epoch 1600:  0.06343422085046768\n",
      "Epoch 1700:  0.05505188927054405\n",
      "Epoch 1800:  0.05087268352508545\n",
      "Epoch 1900:  0.04529455304145813\n"
     ]
    }
   ],
   "source": [
    "semisup_model.train()\n",
    "curr_loss = 1000; F_print = 10 if choice == 'LBFGS' else 100\n",
    "\n",
    "# Stage I\n",
    "for i in range(epochs1):\n",
    "    optimizer1.step(pcgrad_closure)\n",
    "    l = pcgrad_closure()\n",
    "    if (i % F_print) == 0:\n",
    "        if l.item() != curr_loss:\n",
    "            curr_loss = l.item()\n",
    "        else:\n",
    "            print(\"Epoch {}: \".format(i), curr_loss)\n",
    "            print(\"Finishing the first stage\")\n",
    "            break\n",
    "        print(\"Epoch {}: \".format(i), curr_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.0002471697225701064\n",
      "Epoch 10:  1.0192807167186402e-06\n",
      "Epoch 20:  8.126207831082866e-07\n",
      "Epoch 30:  7.739533316453162e-07\n",
      "Epoch 40:  6.066410378480214e-07\n",
      "Epoch 50:  5.455259497466614e-07\n",
      "Epoch 60:  5.2041633580302e-07\n",
      "Finishing the second stage\n",
      "Testing\n",
      "Test MSE: 4.511709448706824e-06\n"
     ]
    }
   ],
   "source": [
    "optimizer2 = torch.optim.LBFGS(semisup_model.network.parameters(), \n",
    "                              lr=learning_rate2, max_iter=100, max_eval=125, \n",
    "                              history_size=120, line_search_fn='strong_wolfe')\n",
    "\n",
    "curr_loss = 1000\n",
    "# Stage II\n",
    "for i in range(epochs2):\n",
    "    optimizer2.step(closure)\n",
    "    l = closure()\n",
    "    if (i % 10) == 0:\n",
    "        if l.item() != curr_loss:\n",
    "            curr_loss = l.item()\n",
    "        else:\n",
    "            print(\"Finishing the second stage\")\n",
    "            break\n",
    "        print(\"Epoch {}: \".format(i), curr_loss)\n",
    "\n",
    "print(\"Testing\")\n",
    "semisup_model.network.eval()\n",
    "\n",
    "# should be able to reach the order of 1e-6.\n",
    "# So that I can use this algo instead of the ladder networks\n",
    "# Compare btw the two semi-supervise learning?\n",
    "print('Test MSE:', F.mse_loss(semisup_model.network(*dimension_slicing(X_star)).detach(), u_star).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BEST-2000: 1e-06 (LBFGS)\n",
    "# torch.save(semisup_model.state_dict(), \"./saved_path_inverse_burger/semisup_model_with_LayerNormDropout_without_physical_reg_trained2000samples.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading the best model and testing\n",
    "# semisup_model.load_state_dict(torch.load(\"./saved_path_inverse_burger/semisup_model_with_LayerNormDropout_without_physical_reg_trained2000samples.pth\"), strict=False)\n",
    "# semisup_model.eval()\n",
    "# F.mse_loss(semisup_model.network(*dimension_slicing(X_star)).detach(), u_star)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# derivatives_test, dynamics_test = semisup_model.network.get_selector_data(*dimension_slicing(X_star))\n",
    "# derivatives_train, dynamics_train = semisup_model.network.get_selector_data(*dimension_slicing(X_u_train))\n",
    "\n",
    "# derivatives_test, dynamics_test = to_numpy(derivatives_test), to_numpy(dynamics_test)\n",
    "# derivatives_train, dynamics_train = to_numpy(derivatives_train), to_numpy(dynamics_train)\n",
    "\n",
    "# np.save(\"./saved_path_inverse_burger/data/derivatives-2000-V3.npy\", derivatives_train)\n",
    "# np.save(\"./saved_path_inverse_burger/data/dynamics-2000-V3.npy\", dynamics_train)\n",
    "# np.save(\"./saved_path_inverse_burger/data/derivatives-25600-V3.npy\", derivatives_test)\n",
    "# np.save(\"./saved_path_inverse_burger/data/dynamics-25600-V3.npy\", dynamics_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}