Selector Network for semi-supervise learning (old version without 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):\n",
    "        super().__init__()\n",
    "        # Nonlinear model, Training with PDE reg.\n",
    "        layers = [nn.Linear(X_train_dim, 50), nn.Tanh(), nn.Linear(50, 1)]\n",
    "        self.nonlinear_model = nn.Sequential(*layers)\n",
    "        self.nonlinear_model.apply(self.xavier_init)\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": [
    "semisup_model = SemiSupModel(network=Network(model=simple_solver_model(50)), \n",
    "                             selector=SeclectorNetwork(X_train_dim=6), \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": [],
   "source": [
    "params = semisup_model.parameters()\n",
    "\n",
    "### For SGD and Adam ###\n",
    "# learning_rate1, learning_rate2 = 1e-7, 5e-2\n",
    "\n",
    "### For LBFGS (a good choice already!!!) ###\n",
    "learning_rate1, learning_rate2 = 8e-2, 5e-2 # (1e-1, 5e-2) is also OK!\n",
    "\n",
    "choice = 'LBFGS'; 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",
    "    lr_finder = LRFinder(semisup_model, optimizer=optimizer1, closure=pcgrad_update, criterion=criterion, device=\"cpu\")\n",
    "    lr_finder.range_test(trainloader, val_loader=None, end_lr=100, num_iter=300)\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",
    "    # Set the learing_rate to the suggested one.\n",
    "    optimizer1 = lr_finder.optimizer\n",
    "    for g in optimizer1.param_groups:\n",
    "        g['lr'] = suggested_lr\n",
    "\n",
    "epochs1 = 1500; epochs2 = 500;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.283119797706604\n",
      "Epoch 10:  0.2770058810710907\n",
      "Epoch 20:  0.2770058810710907\n",
      "Finishing the first stage\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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:  0.015079818665981293\n",
      "Epoch 10:  5.55653468836681e-06\n",
      "Epoch 20:  1.692606588221679e-06\n",
      "Epoch 30:  1.3741183693127823e-06\n",
      "Epoch 40:  1.314015776188171e-06\n",
      "Epoch 50:  1.230398424922896e-06\n",
      "Epoch 60:  1.2222469649714185e-06\n",
      "Epoch 70:  1.198522227241483e-06\n",
      "Epoch 80:  1.1328039590807748e-06\n",
      "Epoch 90:  9.887070291370037e-07\n",
      "Epoch 100:  9.69153234109399e-07\n",
      "Epoch 110:  9.69047050602967e-07\n",
      "Finishing the second stage\n",
      "Testing\n",
      "Test MSE: 2.0664865587605163e-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",
    "# 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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BEST-2000: 1e-06 (LBFGS)\n",
    "# torch.save(semisup_model.network.state_dict(), \"./saved_path_inverse_burger/nn_nonlinear_semisup_without_physical_reg_trained2000samples.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading the best model and testing\n",
    "# semisup_model.network.load_state_dict(torch.load(\"./saved_path_inverse_burger/nn_nonlinear_semisup_without_physical_reg_trained2000samples.pth\"))\n",
    "# F.mse_loss(semisup_model.network(*dimension_slicing(X_star)).detach(), u_star)"
   ]
  },
  {
   "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
}