Semi-supervise learning without adversarial examples (Freezed training settings).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",
    "from pytorch_stats_loss import torch_wasserstein_loss, torch_energy_loss\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training with 2000 samples\n",
      "Training with 2000 unsup 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",
    "# Unsup data\n",
    "N_res = 1000\n",
    "idx_res = np.array(range(X_star.shape[0]-1))[~idx]\n",
    "idx_res = np.random.choice(idx_res.shape[0], N_res, replace=True)\n",
    "X_res = X_star[idx_res, :]\n",
    "print(f\"Training with {N} unsup samples\")\n",
    "X_u_train = np.vstack([X_u_train, X_res])\n",
    "u_train = np.vstack([u_train, torch.rand(X_res.shape[0], 1) - 1000])\n",
    "# del X_res\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",
    "        # pls init the self.model before\n",
    "        self.model = model\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=False, mini=None, maxi=None):\n",
    "        super(SemiSupModel, self).__init__()\n",
    "        self.network = network\n",
    "        self.selector = selector\n",
    "        self.normalize_derivative_features = normalize_derivative_features\n",
    "        self.mini = mini\n",
    "        self.maxi = maxi\n",
    "    def forward(self, X_u_train):\n",
    "        X_selector, y_selector = self.network.get_selector_data(*dimension_slicing(X_u_train))\n",
    "        if self.normalize_derivative_features:\n",
    "            X_selector = (X_selector-self.mini)/(self.maxi-self.mini)\n",
    "        unsup_loss = self.selector.loss(X_selector, y_selector)\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",
    "\n",
    "### Version without normalized derivatives ###\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, \n",
    "#                              mini=None, \n",
    "#                              maxi=None)\n",
    "\n",
    "### Version with normalized derivatives ###\n",
    "referenced_derivatives = np.load(\"./saved_path_inverse_burger/data/derivatives-25600-V1-with-1000unlabledsamples.npy\")\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=True, \n",
    "                             mini=to_tensor(np.min(referenced_derivatives, axis=0), False), \n",
    "                             maxi=to_tensor(np.max(referenced_derivatives, axis=0), False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pcgrad_closure():\n",
    "    global N, X_u_train, u_train\n",
    "    uf, unsup_loss = semisup_model(X_u_train)\n",
    "    losses = [F.mse_loss(uf[:N, :], 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",
    "    global N, X_u_train, u_train\n",
    "    optimizer2.zero_grad()\n",
    "    mse_loss = F.mse_loss(semisup_model.network(*dimension_slicing(X_u_train))[:N, :], 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": "172204954ee34f0386b3859dc405bcdc",
       "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.35E-06\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAX4AAAENCAYAAAAIbA6TAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/Il7ecAAAACXBIWXMAAAsTAAALEwEAmpwYAAAt0ElEQVR4nO3deXxU5dn/8c81kxWyAQkQCPtOWCVCAcUVRawg7lT7KLWitWqpLbVqVXgqP23Vam21io8IdbeIgoCKC+ICBcMiq+yLgUDCEpJA9ly/P2aIAZKQQCZnluv9es0rc+4zZ843Q7hm5j7n3LeoKsYYY0KHy+kAxhhjGpYVfmOMCTFW+I0xJsRY4TfGmBBjhd8YY0KMFX5jjAkxPiv8IhIlIstE5DsRWScik73t00Vku4is8t76+SqDMcaYk4X58LmLgAtVNV9EwoGvReRD77qJqjrTh/s2xhhTDZ8VfvVcGZbvXQz33uxqMWOMcZj48spdEXEDy4HOwHOqep+ITAcG4/lG8BnwR1Utqul5EhMTtX379j7LaYwxwWj58uX7VTXpxHafFv6KnYgkAO8BdwMHgL1ABDAV2Kqq/1vFNuOB8QBt27YdsHPnTp/nNMaYYCIiy1U17cT2BjmrR1VzgIXACFXNVI8i4BVgYDXbTFXVNFVNS0o66Q3LGGPMafLlWT1J3k/6iEg0MBz4XkSSvW0CXAms9VUGY4wxJ/PlWT3JwAxvP78LeEdV54rI5yKSBAiwCrjDhxmMMcacwJdn9awG+lfRfqGv9mmM+VFJSQkZGRkUFhY6HcX4WFRUFCkpKYSHh9fq8b78xG+McVBGRgaxsbG0b98eT8+qCUaqyoEDB8jIyKBDhw612saGbDAmSBUWFtKsWTMr+kFORGjWrFmdvtkF9Sf+dXsOk51XRPPYKJrHRdK0UQQul/0nMKHDin5oqOu/c1AX/tf+u4s3l+2qWHa7hMSYCFrERZHSJJo2TRrRpqnn1rZpI9o0iSbMbV+CTIhShaVLITMTkpNh0CCo5zeOZ555hvHjx9OoUaN6fd66yMnJ4Y033uDOO+9skP21b9+e9PR0EhMTGTJkCIsXLz6t55k+fTqXXHIJrVq1OuNMQV34f3txF64Z0Jqs3CKy8orIyiskO6+IvblFfJ+Zx6frsyguK694fESYi85JMXRPjqVHyzj6t02gd0o8kWFuB38LYxrA/Plw++2QkwMuF5SXQ0ICvPgijBxZb7t55plnuOmmmxwv/M8///wZFf7S0lLCwupePk+36IOn8Pfq1csK/6k0j4uieVxUtevLy5V9eYX8cLCAnQeOsDkrn+/35vHNlv3MWrEb8LwZ9E2JZ0C7ppzdvgmDOjYjJjKoXzYTaubPh2uugYKC49vz8z3tM2fWufgfOXKE6667joyMDMrKynjooYfYt28fe/bs4YILLiAxMZGFCxeyYMECHnnkEYqKiujUqROvvPIKMTExLF++nHvvvZf8/HwSExOZPn06ycnJnH/++fTt25dFixZRWlrKtGnTGDhwIEeOHOHuu+9m7dq1lJSUMGnSJEaPHs26desYN24cxcXFlJeX8+677/LQQw+xdetW+vXrx/Dhw3niiSeOy/7nP/+Z1157jaSkJNq0acOAAQP4/e9/z/nnn0+/fv34+uuvGTt2LF27duXRRx+luLiYZs2a8frrr9OiRQsOHDjA2LFj2b17N4MHD6by6AgxMTHk53uGMHviiSd45513KCoqYsyYMUyePJkdO3Zw2WWXcc4557B48WJat27N7NmzmTdvHunp6dx4441ER0ezZMkSoqOjT+/fGzxHhP39NmDAAG1o+/MK9eO1mTpl3nq98rmvtdP987TdfXO164Pz9Y5X03Xe6j1aUFza4LmMqa3169ef+kHl5aqtW6t6OnqqvqWkeB5XBzNnztRf/vKXFcs5OTmqqtquXTvNzs5WVdXs7Gw999xzNT8/X1VVH3/8cZ08ebIWFxfr4MGDNSsrS1VV33rrLR03bpyqqp533nkVz7to0SJNTU1VVdX7779fX331VVVVPXTokHbp0kXz8/P1rrvu0tdee01VVYuKivTo0aO6ffv2iu1OtGzZMu3bt68WFBRobm6udu7cWZ944omKff/qV7+qeOzBgwe13Pu6vPTSS3rvvfeqqurdd9+tkydPVlXVuXPnKlDxOzdu3FhVVT/++GO97bbbtLy8XMvKyvTyyy/XRYsW6fbt29XtduvKlStVVfXaa6+t+L3OO+88/fbbb6t9zav69wbStYqaah9dq9EsJpJLUltySWpLAAqKy1i56xAfr9vLvDWZfLh2L40j3FzcswVj+rdmWJckO3BsAs/SpXD4cM2PycmBZcs8ff611Lt3b373u99x33338dOf/pRzzz33pMf897//Zf369QwdOhSA4uJiBg8ezMaNG1m7di3Dhw8HoKysjOTk5Irtxo4dC8CwYcPIzc0lJyeHBQsWMGfOHJ588knAc0bTrl27GDx4MFOmTCEjI4OrrrqKLl261Jj7m2++YfTo0URFRREVFcUVV1xx3Prrr7++4n5GRgbXX389mZmZFBcXV5xK+eWXXzJr1iwALr/8cpo0aXLSfhYsWMCCBQvo399zqVN+fj6bN2+mbdu2dOjQgX79+gEwYMAAduzYUWPm02GFv5aiI9wM6ZzIkM6JPHxFKku3HeCD1Xv4cO1eZq/aQ8fExowb2p5rBrQhOsKOCZgAkZnp6dOvicsFe/bU6Wm7du3KihUrmD9/Pn/605+46KKLePjhh497jKoyfPhw3nzzzePa16xZQ2pqKkuWLKnyuU88g0VEUFXeffddunXrdty6Hj16MGjQIObNm8fIkSN58cUX6dixY51+l8oaN25ccf/uu+/m3nvvZdSoUXzxxRdMmjSp1s+jqtx///3cfvvtx7Xv2LGDyMjIimW3203BiV1w9cBOYTkNbpcwpHMij13Vh2UPXMzfb+hHbFQYD81exzl/+Zznv9hCXmGJ0zGNObXkZM+B3JqUl0MdDyju2bOHRo0acdNNNzFx4kRWrFgBQGxsLHl5eQD85Cc/4ZtvvmHLli2A57jApk2b6NatG9nZ2RWFv6SkhHXr1lU899tvvw3A119/TXx8PPHx8Vx66aX84x//qOhPX7lyJQDbtm2jY8eO3HPPPYwePZrVq1cfl+FEQ4cO5YMPPqCwsJD8/Hzmzp1b7e94+PBhWrduDcCMGTMq2ocNG8Ybb7wBwIcffsihQ4dO2vbSSy9l2rRpFf39u3fvJisrq8bXtKbcdWWf+M9QRJiL0f1aM6pvK9J3HuK5hVv460cbeeGLrYwb2oFbz+1AXFTtLqM2psENGgTx8Z4DudVJSICBVQ6iW601a9YwceJEXC4X4eHh/Otf/wJg/PjxjBgxglatWrFw4UKmT5/O2LFjKSryTMnx6KOP0rVrV2bOnMk999zD4cOHKS0tZcKECaSmpgKe4Qn69+9PSUkJ06ZNA+Chhx5iwoQJ9OnTh/Lycjp06MDcuXN55513ePXVVwkPD6dly5Y88MADNG3alKFDh9KrVy8uu+yy4w7unn322YwaNYo+ffrQokULevfuTXx8fJW/46RJk7j22mtp0qQJF154Idu3bwfgkUceYezYsaSmpjJkyBDatm170raXXHIJGzZsYPDgwYDnoO9rr72G2119b8Ett9zCHXfcUS8HdxtkPP4zlZaWpunp6U7HqLU1GYf5x+ebWbB+H4kxEfxhRHeuOSvFjgGYBrVhwwZ69Ohx6gdWd1YPQHT0aZ3V4yvnn38+Tz75JGlpJw0xX2/y8/OJiYnh6NGjDBs2jKlTp3LWWWf5bH/1pap/b0fH4w81vVPimfo/aXxw1zm0bdqIP8xczVX/Wsx3P+Q4Hc2Yk40c6SnuKSkQEwNxcZ6fKSl+VfQbyvjx4+nXrx9nnXUWV199dUAU/bqyT/w+Vl6uvLdyN499+D0HjhRxy5D23DeiO1HhdgDY+FatP/Efo+o5e2fPHk+f/sCB9X7lrvGdunzitz5+H3O5hKsHpHBJague/Hgjr3yzg0Wbsvnbdf3o1ybB6XjG/EikTqdsmsBlXT0NJDYqnMmje/HarYMoKC7jmn8t5sVFWykv9/9vXCZwBcI3enPm6vrvbIW/gZ3TJZGPJgxjeM8WPPbh9/zy3+kcOlLsdCwThKKiojhw4IAV/yCn3vH4o6KqH57mRNbH7xBV5d9LdjJl3gZaJUQx4xcDades8ak3NKaWbAau0FHdDFzV9fFb4XfY8p0HuXVGOmEuYdotZ9MnJcHpSMaYIGGnc/qpAe2aMvOOIUSFu7lh6n/5YmPNV+8ZY8yZssLvBzo3j2HWnUPokNiYW2ek8/7K3U5HMsYEMZ8VfhGJEpFlIvKdiKwTkcne9g4islREtojI2yIS4asMgaR5bBRv3z6Yge2bcu87q5i3OtPpSMaYIOXLT/xFwIWq2hfoB4wQkZ8AfwGeVtXOwCHgVh9mCCgxkWG8fEsaA9o14TdvrWTBur1ORzLGBCGfFX7vPADHRn4K994UuBCY6W2fAVzpqwyBqFFEGNNuOZvU1vH8+o0VLN12wOlIxpgg49M+fhFxi8gqIAv4BNgK5KhqqfchGUDrarYdLyLpIpKenZ3ty5h+JzYqnH+PG0ibpo24/bXl7DxwxOlIxpgg4tPCr6plqtoPSAEGAt3rsO1UVU1T1bSkpCRfRfRb8Y3CmXbz2QD8Yvq3HC6w8f2NMfWjQc7qUdUcYCEwGEgQkWNjBKUAdgpLNdonNuaFmwaw6+BR7npjBaVlp5gwwxhjasGXZ/UkiUiC9340MBzYgOcN4Brvw24GZvsqQzD4ScdmTLmyN19t3s/kD9Y7HccYEwR8OTpnMjBDRNx43mDeUdW5IrIeeEtEHgVWAi/7MENQuO7sNmzJzmfql9vo2yaBawakOB3JGBPAfFb4VXU10L+K9m14+vtNHdw3ojurM3L40/tr6N06nm4tY52OZIwJUHblboBwu4Rnb+hPTGQ4d76+nCNFpafeyBhjqmCFP4A0j4vi7zf0Y9v+Izw8e53TcYwxAcoKf4AZ2jmRuy/ozLsrMvhorQ3rYIypOyv8Aejui7rQq3UcD7y3lqw8G2vdGFM3VvgDULjbxdPX9SO/qJQHZq2xGZaMMXVihT9AdWkRy30juvPphixmrbBr4IwxtWeFP4CNG9Kes9omMGX+BnKO2ry9xpjascIfwFwuYcqY3hwuKOEvH33vdBxjTICwwh/geiTHces5HXhz2Q+k7zjodBxjTACwwh8EJlzchdYJ0Tz43lpKbCA3Y8wpWOEPAo0iwpg0KpWN+/J4+evtTscxxvg5K/xBYnjPFgzv2YJnPt1ExqGjTscxxvgxK/xBZNKoVAAe+9AO9BpjqmeFP4i0Tojm9mGdmLc6k2Xb7UCvMaZqVviDzB3ndSI5PorJH6yjrNyu6DXGnMwKf5CJjnDzx8u6s25PLjOX/+B0HGOMH7LCH4RG9W3FgHZNeOLjjeQV2iTtxpjjWeEPQiLCI1f0ZH9+Mf9cuMXpOMYYP2OFP0j1SfHMzTvt6+3sPHDE6TjGGD9ihT+I/eHSboS7XTaOjzHmOD4r/CLSRkQWish6EVknIr/xtk8Skd0issp7G+mrDKGueVwUtw/rxPw1e20cH2NMhTAfPncp8DtVXSEiscByEfnEu+5pVX3Sh/s2XrcN68Aby3by6Nz1vDfAjezdC8nJMGgQiDgdzxjjAJ8VflXNBDK99/NEZAPQ2lf7M1VrFBHGE4130/mB31NWVkBYmBvKyyEhAV58EUbaFy5jQk2D9PGLSHugP7DU23SXiKwWkWki0qQhMoSs+fM59/5f0SpvP2FHj0BuLuTnQ0YGXHMNzJ/vdEJjTAPzeeEXkRjgXWCCquYC/wI6Af3wfCN4qprtxotIuoikZ2dn+zpmcFKF8eORgoKq1xcUwO23ex5njAkZPi38IhKOp+i/rqqzAFR1n6qWqWo58BIwsKptVXWqqqapalpSUpIvYwavpUvh8OGaH5OTA8uWNUgcY4x/8OVZPQK8DGxQ1b9Vak+u9LAxwFpfZQh5mZngOsU/scsFe/Y0TB5jjF/w5Vk9Q4GfA2tEZJW37QFgrIj0AxTYAdzuwwyhLTnZcyC3JuXl0KpVw+QxxvgFX57V8zVQ1fmCdjSxoQwaBPHxnoO51UlIgIFV9rYZY4KUXbkbzERg6lSIjq5ydVlUtOeUTjuf35iQYoU/2I0cCTNnQkoKxMRAXBwaE8O+uCSm3DKZ8hGXOZ3QGNPAfNnHb/zFyJGwa5fn7J09e5BWrfja3YppM1fTd/UeRvez6+qMCSVW+EOFiKfP32tMuTJt8Q7++tFGLk1tSVS428FwxpiGZF09IcrlEh4c2YPdOQW88s0Op+MYYxqQFf4QNqRzIhd1b87zC7dwIL/I6TjGmAZihT/E3T+yB0dLyvj7Z5udjmKMaSBW+ENc5+Yx/GxgW15fuostWTWc72+MCRpW+A2/ubgL0eFuHv/QZuoyJhRY4TckxkRy5wWd+HTDPpZsPeB0HGOMj1nhNwD8YmgHWidEM2X+esrLbZhmY4KZFX4DQFS4m4mXdmPt7lzeX7Xb6TjGGB+ywm8qjOrbij4p8Tzx8UYKS8qcjmOM8REr/KbCsYu6Mg8X8tzCLU7HMcb4iBV+c5xBHZsxpn9rXli0lS1ZeU7HMcb4gBV+c5IHL+9Bo4gwHnhvLWrz8RoTdKzwm5MkxkTywMjuLNt+kP8sz3A6jjGmnlnhN1W6dkAbBrZvypR5G8jKK3Q6jjGmHlnhN1VyuYTHru5NYUkZD1qXjzFBxQq/qVanpBgmXtqNT9bvY/aqPU7HMcbUEyv8pkbjhnYgrV0THpmzjn251uVjTDDwWeEXkTYislBE1ovIOhH5jbe9qYh8IiKbvT+b+CqDOXNul/DXa/pQVFrGA7PWWJePMUHAl5/4S4HfqWpP4CfAr0WkJ/BH4DNV7QJ85l02fqxjUgwTL+3OZ99nMWuFDedgTKDzWeFX1UxVXeG9nwdsAFoDo4EZ3ofNAK70VQZTf8YNac/A9k2ZNGcdGYeOOh3HGHMGGqSPX0TaA/2BpUALVc30rtoLtGiIDObMuFzCU9f1BWDCW6soLSt3OJEx5nT5vPCLSAzwLjBBVXMrr1NPh3GVncYiMl5E0kUkPTs729cxTS20adqIP1/Zi/Sdh3hu4Van4xhjTpNPC7+IhOMp+q+r6ixv8z4RSfauTwayqtpWVaeqapqqpiUlJfkypqmDK/u35sp+rXj2880s33nI6TjGmNPgy7N6BHgZ2KCqf6u0ag5ws/f+zcBsX2UwvvG/V/aiVUIUE95eSV5hidNxjDF15MtP/EOBnwMXisgq720k8DgwXEQ2Axd7l00AiYsK55nr+7Mnp5CH3rereo0JNGG+emJV/RqQalZf5Kv9moYxoF0TfnNRF/72ySaGdE7kurQ2TkcyxtSSXblrTtuvL+jMkE7NeHj2Wjbvs7H7jQkUVvjNaXO7hGeu70dMZBi/fmMFBcU2XaMxgcAKvzkjzeOiePr6fmzOymfyB+ucjmOMqQUr/OaMndsliTvP78Rb3/7A7FU2pIMx/s4Kv6kXv724K2ntmvDArDVs33/E6TjGmBrUqvCLSGMRcXnvdxWRUd6Ls4wBIMzt4tmx/QkPc/Hr11dQWGL9/cb4q9p+4v8SiBKR1sACPOfnT/dVKBOYWiVE8+Q1fVmfmctj8zc4HccYU43aFn5R1aPAVcDzqnotkOq7WCZQXdyzBb88pwMzluzko7WZp97AGNPgal34RWQwcCMwz9vm9k0kE+j+MKI7fVPimThzNT8ctCGcjfE3tS38E4D7gfdUdZ2IdAQW+iyVCWgRYS7++bOzALjrzZUUl9oQzsb4k1oVflVdpKqjVPUv3oO8+1X1Hh9nMwGsTdNG/OXqPnz3Qw5PfPy903GMMZXU9qyeN0QkTkQaA2uB9SIy0bfRTKAb2TuZn/+kHS99tZ3PNuxzOo4xxqu2XT09vZOoXAl8CHTAc2aPMTV68PIe9EyO43f/+Y7MwwVOxzHGUPvCH+49b/9KYI6qllDNzFnGVBYV7uafP+tPcWk5v//Pd5SX25+NMU6rbeF/EdgBNAa+FJF2QG6NWxjj1TEphj9d3pNvthxg+uIdTscxJuTV9uDus6raWlVHqsdO4AIfZzNBZOzANlzUvTmPf/S9DeFsjMNqe3A3XkT+dmzycxF5Cs+nf2NqRUR4/Oo+xESGMeHtVXaKpzEOqm1XzzQgD7jOe8sFXvFVKBOckmIjeeyq3qzbk8szn25yOo4xIau2hb+Tqj6iqtu8t8lAR18GM8Hp0tSWXJeWwguLtrJ850Gn4xgTkmpb+AtE5JxjCyIyFLBz88xpefiKVJLjo5k4c7WN4mmMA2pb+O8AnhORHSKyA/gncHtNG4jINBHJEpG1ldomichuEVnlvY087eQmYMVEhvGXq/uwLfsIT39iXT7GNLTantXznar2BfoAfVS1P3DhKTabDoyoov1pVe3nvc2vU1oTNM7pksjYgW146attrNx1yOk4xoSUOs3Apaq53it4Ae49xWO/BKwT11TrgZE9aBkXZV0+xjSwM5l6UU5zu7tEZLW3K6jJGezfBLjYqHAeu7oPW7LyefazzU7HMSZknEnhP51r7/8FdAL6AZnAU9U9UETGH7tuIDs7+/QSGr93XtekirN81u4+7HQcY0JCjYVfRPJEJLeKWx7Qqq47U9V9qlqmquXAS8DAGh47VVXTVDUtKSmprrsyAeTBy3vSLCaS+2etobTMLuwyxtdqLPyqGquqcVXcYlU1rK47E5HkSotj8AzxbEJcfHQ4j1zRkzW7DzNjyU6n4xgT9M6kq6dGIvImsAToJiIZInIr8FcRWSMiq/GM9fNbX+3fBJbLeydzQbcknlqwkd05domIMb7ks8KvqmNVNVlVw1U1RVVfVtWfq2pvVe3jndHLZuM2gGcsn/8d3QtVeGT2WlRt+GZjfMVnhd+YumrTtBH3Du/Kpxuy+GjtXqfjGBO0rPAbvzJuaHtSW8XxyJx15BaWOB3HmKBkhd/4lTC3i8eu6s3+/CKe+Gij03GMCUpW+I3f6ZOSwP8Mbs9rS3eyfKcN52BMfbPCb/zS7y/tRsu4KB6YtYYSO7ffmHplhd/4pZjIMCaPSmXjvjxe+mqb03GMCSpW+I3fuiS1JZemtuDZzzbzw8GjTscxJmhY4Td+7ZErUnGJMPmDdU5HMSZoWOE3fq1VQjQTLu7CpxuyWLDOzu03pj5Y4Td+b9zQDnRrEcvkD9ZztLjU6TjGBDwr/MbvhbtdTBnTi905Bfzdxu035oxZ4TcBIa19U65LS+Hlr7azcW+e03GMCWhW+E3A+ONlPYiJCuOh920QN2POhBV+EzCaNo7g/su6s2zHQWYuz3A6jjEBywq/CSjXDmjDgHZNeOzD7zl0pNjpOMYEJCv8JqC4XMKjV/bicEEJf/34e6fjGBOQrPCbgNMjOY5fDG3Pm8t+sEHcjDkNVvhNQJpwcVeS46P40/trbYJ2Y+rICr8JSI0jw3jkip5syMxl+uIdTscxJqBY4TcB69LUllzQLYmnP9lE5mGboN2Y2rLCbwKWiDB5VC9Ky5U/z13vdBxjAobPCr+ITBORLBFZW6mtqYh8IiKbvT+b+Gr/JjS0bdaIuy/szPw1e/liY5bTcYwJCL78xD8dGHFC2x+Bz1S1C/CZd9mYM3LbsI50TGrMw7PXUVhS5nQcY/yezwq/qn4JHDyheTQww3t/BnClr/ZvQkdkmJtHR/di18GjPLdwi9NxjPF7Dd3H30JVM7339wItGnj/JkgN6ZzImP6teWHRVrZk2SBuxtTEsYO76hllq9qRtkRkvIiki0h6dnZ2AyYzgerBy3vQKCKMB9+zQdyMqUlDF/59IpIM4P1Z7dE4VZ2qqmmqmpaUlNRgAU3gSoyJ5I+XdWfpdhvEzZiaNHThnwPc7L1/MzC7gfdvgtz1aW1Ia9eE/zd/AwdtEDdjquTL0znfBJYA3UQkQ0RuBR4HhovIZuBi77Ix9cblEqaM6U1eYSmPzd/gdBxj/FKYr55YVcdWs+oiX+3TGIBuLWP55bkdeWHRVq4ZkMKgjs2cjmSMX7Erd01Q+s1FXUhpEs2D76+luNQGcTOmMiv8JihFR7j58+hebMnK54VFW52OY4xfscJvgtYF3Zvz0z7J/OPzzWzaZ+f2G3OMFX4T1CaPSiU2KpyJ//nOxu03xssKvwlqzWIimTQqle8yDjPtm+1OxzHGL1jhN0Hvij7JXNyjBU8t2MT2/UecjmOM46zwm6AnIkwZ04uIMBf3zVxNebkN52BCmxV+ExJaxEXx0E97smzHQWYs2eF0HGMcZYXfhIxrB6RwYffmPPbh93aWjwlpVvhNyBAR/nJ1H2Ijw7jnzZUUldqkLSY0WeE3ISUpNpK/XtOH7/fm8eTHG52OY4wjrPCbkHNRjxbc9JO2vPTVdr7Zst/pOMY0OCv8JiQ9OLInnZIa87t3viPnqA3fbEKLFX4TkqIj3Pz9hv4cOFLEhLdX2SmeJqRY4Tchq1freB65IpUvNmbzzKebnI5jTIOxwm9C2o2D2nLtgBSe/XwLn6zf53QcYxqEFX4T0kSEP1/Zi16t4/jt26vYkJnrdCRjfM4Kvwl5UeFupv48jcaRbsa98i2ZhwucjmSMT1nhNwZolRDNK7cMJL+olHGvfEteYYnTkYzxGSv8xnj1bBXH8zeexZasfG77dzoFxXZlrwlOVviNqWRY1ySeuq4vS7cfZPyr6RSWWPE3wceRwi8iO0RkjYisEpF0JzIYU53R/Vrzl6v78NXm/fzqteU2po8JOk5+4r9AVfupapqDGYyp0nVpbZgyphcLN2Zz6/R0jhSVOh3JmHpjXT3GVOPGQe148tq+LNl2gJ/931IOHbGhHUxwcKrwK7BARJaLyPiqHiAi40UkXUTSs7OzGzieMR7XDEjhhZsGsCEzl6tfWMwOm7rRBAGnCv85qnoWcBnwaxEZduIDVHWqqqapalpSUlLDJzTGa3jPFrz6i4EcOlLMlc9/w+KtNqKnCWyOFH5V3e39mQW8Bwx0IocxtTWoYzPe//VQEmMi+Z+Xl/HG0l1ORzLmtDV44ReRxiISe+w+cAmwtqFzGFNX7Zo1ZtadQxjaOZEH3lvDw7PX2hk/JiA58Ym/BfC1iHwHLAPmqepHDuQwps7iosJ5+eY0fnlOB/69ZCdXPb+Y7dbvbwKMqPr/OORpaWmanm6n+xv/8sn6fUyc+R0lpeU8OqYXY/qnOB3JmOOIyPKqTpm30zmNOU3De7Zg/j3nktoqnt++/R33vLmSrLxCp2MZc0pW+I05A60SonnjtkH89uKufLR2Lxc9tYh/L9lBaVm509GMqZYVfmPOUJjbxW8u7sKHE86lT0o8D89exyVPf8nsVbspsykdjR+ywm9MPemUFMNrtw7ihZsGEO528Zu3VjHimS+ZuTyD4lL7BmD8hx3cNcYHysuVeWsyeW7hFr7fm0dyfBS3ntOBGwa2JSYyzOl4JkRUd3DXCr8xPqSqLNqUzYuLtrFk2wHiosL4+eB23DKkA0mxkU7HM0HOCr8xDlv1Qw4vLtrKR+v2Eu52cc2AFMaf25H2iY2djmaClBV+Y/zE9v1HmPrlNt5dkUFJWTmX9WrJHed1ok9KgtPRTJCxwm+Mn8nKK2TG4h28umQnuYWlnNU2gbED2/LTPq2IjnA7Hc8EASv8xvip/KJS3v72B95YupOt2UeIjQpjTP/W3HB2W3q2inM6nglgVviN8XOqyrc7DvHmsl3MW5NJcWk5fdskcMPZbRiR2pImjSOcjmgCjBV+YwJIztFiZq3YzZvLdrE5K58wlzC4UzMu753MJaktaWpvAqYWrPAbE4BUlbW7c5m3JpP5azLZdfAoItA3JYFhXZM4r2sSfVPiCXPbtZjmZFb4jQlwqsq6Pbl8umEfizZl890POZQrxEWFcW6XJIZ1TWRY1ySS46Odjmr8hBV+Y4JMztFivt6yny83ZbNoUzb7cosA6NoihnO7JHF2+yac1bYJzeOiHE5qnGKF35ggpqps2pfPok1ZfLlpP8t2HKwYHyilSTT92zYhtVUcPZLj6JkcZ1cNhwgr/MaEkKLSMtbtyWXFzkOs2HWIVbty2HP4x7kCEmMi6dkqjh7JsfRM9rwhtGvWiAi3i325RRwtLqV1k2giw+x6gkBWXeG30aKMCUKRYW7Oauvp6jkm52gxGzLzWJ+Zy4bMXNbvyWXa1v2UlFX94U8EWsVHkxgbSXx0OHFRYcRFh3vve39Gh5EYE0lSbCTNYyOJiQxDRBrq1zSnyQq/MSEioVEEgzs1Y3CnZhVtxaXlbM3OZ+PePH44eJSSciUpNpJG4W52HTzKzgNHOHi0hMMFJWQcPMrhAs/90mrmGYgKd9E8NoomjcKJ875BxEWHeX96bkkxESTFRpIUE0VibASNIqwMNTR7xY0JYRFhLnp4u3pqS1UpKCkjt6CUnIJiDuQXk5VXSHZeEdl5RWTlFZFztITcwhL25BSQW1hKbkEJRdXMSdA4wu15I4iNrPj2kBRzwnJsJM1iIqzrqZ44UvhFZATwd8AN/J+qPu5EDmNM3YkIjSLCaBQRRsv42p8xVFhSRm5BCdn5njeI/fnFFW8W2flF7M8rYnNWPou3HuBwQUmVzxEfHV7xxtCkceUup2PfMMIqluMrfeOwN4zjNXjhFxE38BwwHMgAvhWROaq6vqGzGGMaTlS4m6hwd61OLy0qLeOA941hv/eNouK+d3nzvvyKrqfqvk0cE+72vFk1jnDTKDKMxpHe+xFhxER627zLjSI8OSPDXESGu4gM894PcxN1bDncdVJbRJgLtyswjm848Yl/ILBFVbcBiMhbwGjACr8xBvAcnG6VEE2rhNpdjFZYUkZuYQm5BSUcLij1/iypaDtSXMbRolKOFJdxxPvzaFEph44WcKSolKPFpRwpKqOgpOyMcrsEwlwuwtxCmEsId3veDMLdP7ZVrHe7vMsnr3e7BbcIbpcwfljHOnXF1YYThb818EOl5QxgkAM5jDFBouLbROyZXaxWVu45flFcWk5RaRmFJZ6fRSXlFJX+eL/wxLbScgpLyigrV0rKlNKyckrLldLyckrLvG3l3rYyb1u5UlZeTkmZcrS41Lvux8eVlytlqtxQ0KaeXqUf+e3BXREZD4wHaNu2rcNpjDGhwO0Sz5zIQX59mxMjO+0GKr+FpXjbjqOqU1U1TVXTkpKSGiycMcYEOycK/7dAFxHpICIRwA3AHAdyGGNMSGrwrh5VLRWRu4CP8ZzOOU1V1zV0DmOMCVWO9PGr6nxgvhP7NsaYUGezNxhjTIixwm+MMSHGCr8xxoQYK/zGGBNiAmIiFhE5DGwGEoH9ddw8Hjhcx/WnaqvqflVt9Z23unUntle3XFNuf35t/SlvTcsnZjyTrKfKezp/C8H02lZu84f/Z/762rZT1ZMvhFJVv78BU70/009327qsP1VbVferaavXvNWtO7G9uuWacvvza+tPeWtariLjaWf1xd9CML22vswbjK/tibdA6er5wIfbVrX+VG1V3a9ufV3VtG11605sr275VLnrqqFe28r3nc5b0/KJGc8k66m2P52/hWB6bWuzz7rmOdW6QH5tjxMQXT3HiEi6VjF/pL8KpLyBlBUCK28gZQXL60v+kjVQPvEfM9XpAHUUSHkDKSsEVt5AygqW15f8ImtAfeI3xhhz5gLtE78xxpgzZIXfGGNCjBV+Y4wJMUFT+EXkXBF5QUT+T0QWO52nJiLiEpEpIvIPEbnZ6TynIiLni8hX3tf3fKfznIqINBaRdBH5qdNZTkVEenhf15ki8iun85yKiFwpIi+JyNsiconTeWoiIh1F5GURmel0lup4/1ZneF/TGxtqv35R+EVkmohkicjaE9pHiMhGEdkiIn+s6TlU9StVvQOYC8zw56x4JpdPAUrwzDnsM/WUV4F8IAof5q2nrAD3Ae/4JuVxuerj73aD9+/2OmBoAOR9X1VvA+4ArvfzrNtU9VZfZaxOHbNfBcz0vqajGixkXa8i88UNGAacBayt1OYGtgIdgQjgO6An0BtPca98a15pu3eAWH/OCvwRuN277Ux/f20Bl3e7FsDrfp51OJ5Z3W4Bfurvr613m1HAh8DPAiGvd7ungLMCJKtP/4+dYfb7gX7ex7zRUBn9YrJ1Vf1SRNqf0DwQ2KKq2wBE5C1gtKo+BlT5FV5E2gKHVTXPn7OKSAZQ7F0s81VWqL/X1usQPpyGup5e2/OBxnj+UxWIyHxVLffXvN7nmQPMEZF5wBu+yFpfeUVEgMeBD1V1hT9ndUpdsuP5Bp0CrKIBe2D8ovBXozXwQ6XlDGDQKba5FXjFZ4mqV9ess4B/iMi5wJe+DFaNOuUVkauAS4EE4J8+TXayOmVV1QcBROQWYL+vin4N6vrano/n634kzsxKV9e/3buBi4F4Eemsqi/4MtwJ6vraNgOmAP1F5H7vG4RTqsv+LPBPEbmcMx/io9b8ufDXmao+4nSG2lDVo3jepAKCqs7C82YVMFR1utMZakNVvwC+cDhGranqs3iKld9T1QN4jkX4LVU9Aoxr6P36xcHdauwG2lRaTvG2+aNAygqBlTeQsoLl9aVAynoiv8ruz4X/W6CLiHQQkQg8B+zmOJypOoGUFQIrbyBlBcvrS4GU9UT+lb0hj3bXcBT8TSCTH09vvNXbPhLYhOdo+INO5wy0rIGWN5CyWl7LGsjZbZA2Y4wJMf7c1WOMMcYHrPAbY0yIscJvjDEhxgq/McaEGCv8xhgTYqzwG2NMiLHCbwKaiOQ38P4adK4HEUkQkTsbcp8m+FnhN6YSEalx/CpVHdLA+0wArPCbemWF3wQdEekkIh+JyHLxzBzW3dt+hYgsFZGVIvKpiLTwtk8SkVdF5BvgVe/yNBH5QkS2icg9lZ473/vzfO/6mSLyvYi87h2yGBEZ6W1bLiLPisjcKjLeIiJzRORz4DMRiRGRz0RkhYisEZHR3oc+DnQSkVUi8oR324ki8q2IrBaRyb58LU2QcvryZrvZ7UxuQH4VbZ8BXbz3BwGfe+83gYqr1X8JPOW9PwlYDkRXWl6MZ6jkROAAEF55f8D5wGE8g225gCXAOXhmKfsB6OB93JvA3Coy3oLncv6m3uUwIM57PxHYAgjQnuMn9LgEmOpd58Iz6cgwp/8d7BZYt6AaltkYEYkBhgD/8X4Ahx8nj0kB3haRZDyzIG2vtOkcVS2otDxPVYuAIhHJwjP72InTTi5T1QzvflfhKdL5wDZVPfbcbwLjq4n7iaoePBYd+H8iMgwoxzN+e4sqtrnEe1vpXY4BuuDMvA4mQFnhN8HGBeSoar8q1v0D+JuqzvFOgDKp0rojJzy2qNL9Mqr+v1Kbx9Sk8j5vBJKAAapaIiI78Hx7OJEAj6nqi3XclzEVrI/fBBVVzQW2i8i14JkqUET6elfH8+MY6Df7KMJGoGOlqfdqOyF5PJDlLfoXAO287XlAbKXHfQz8wvvNBhFpLSLNzzy2CSX2id8EukbeOYyP+RueT8//EpE/AeHAW3gmt56EpwvoEPA50KG+w6hqgff0y49E5Aiecdhr43XgAxFZA6QD33uf74CIfCMia/HMcztRRHoAS7xdWfnATUBWff8uJnjZsMzG1DMRiVHVfO9ZPs8Bm1X1aadzGXOMdfUYU/9u8x7sXYenC8f6441fsU/8xhgTYuwTvzHGhBgr/MYYE2Ks8BtjTIixwm+MMSHGCr8xxoQYK/zGGBNi/j8EVdE1mOOgkAAAAABJRU5ErkJggg==\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",
    "\n",
    "else: suggested_lr = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deleting the fake labels\n",
      "Epoch 0:  31.576194763183594\n",
      "Epoch 100:  7.650530815124512\n",
      "Epoch 200:  3.6829519271850586\n",
      "Epoch 300:  1.9426982402801514\n",
      "Epoch 400:  1.1838877201080322\n",
      "Epoch 500:  0.8849645256996155\n",
      "Epoch 600:  0.703170895576477\n",
      "Epoch 700:  0.5933613777160645\n",
      "Epoch 800:  0.5372388958930969\n",
      "Epoch 900:  0.5393721461296082\n",
      "Epoch 1000:  0.5682770609855652\n",
      "Epoch 1100:  0.5368601679801941\n",
      "Epoch 1200:  0.5273119807243347\n",
      "Epoch 1300:  0.5492504835128784\n",
      "Epoch 1400:  0.5099127888679504\n",
      "Epoch 1500:  0.4665961265563965\n",
      "Epoch 1600:  0.4223608076572418\n",
      "Epoch 1700:  0.40641260147094727\n",
      "Epoch 1800:  0.38611000776290894\n",
      "Epoch 1900:  0.40525883436203003\n"
     ]
    }
   ],
   "source": [
    "print(\"Deleting the fake labels\")\n",
    "u_train = u_train[:N, :]\n",
    "\n",
    "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": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.00022033529239706695\n",
      "Epoch 10:  5.986796622892143e-07\n",
      "Epoch 20:  5.307817332322884e-07\n",
      "Epoch 30:  4.350840754341334e-07\n",
      "Finishing the second stage\n",
      "Testing\n",
      "Test MSE: 7.505475991820276e-07\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",
    "# 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": 13,
   "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_trained2000labeledsamples_trained1000unlabeledsamples.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(7.5055e-07, grad_fn=<MseLossBackward>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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_trained2000labeledsamples_trained1000unlabeledsamples.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": 15,
   "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-3000-V1-with-1000unlabledsamples.npy\", derivatives_train)\n",
    "# np.save(\"./saved_path_inverse_burger/data/dynamics-3000-V1-with-1000unlabledsamples.npy\", dynamics_train)\n",
    "# np.save(\"./saved_path_inverse_burger/data/derivatives-25600-V1-with-1000unlabledsamples.npy\", derivatives_test)\n",
    "# np.save(\"./saved_path_inverse_burger/data/dynamics-25600-V1-with-1000unlabledsamples.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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}