Network_Training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import random\n",
    "from torchvision import transforms\n",
    "from normalize import Normalize, MapToRange\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from torch import nn\n",
    "from torch_nn import *\n",
    "\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"OPENBLAS_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"VECLIB_MAXIMUM_THREADS\"] = \"1\"\n",
    "os.environ[\"NUMEXPR_NUM_THREADS\"] = \"1\"\n",
    "torch.set_num_threads(1)\n",
    "\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Preprocessing the dataset**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrajectoryDataset(Dataset):\n",
    "    \n",
    "    def __init__(self, dataset, indices):\n",
    "        self.dataset = dataset\n",
    "        self.indices = indices\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.indices)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        i, j = self.indices[index]        \n",
    "        X = torch.tensor([\n",
    "            self.dataset['dx'][i, j],\n",
    "            self.dataset['dy'][i, j],\n",
    "            self.dataset['dz'][i, j],\n",
    "            self.dataset['vx'][i, j],\n",
    "            self.dataset['vy'][i, j],\n",
    "            self.dataset['vz'][i, j],\n",
    "            self.dataset['phi'][i, j],\n",
    "            self.dataset['theta'][i, j],\n",
    "            self.dataset['psi'][i, j],\n",
    "            self.dataset['p'][i, j],\n",
    "            self.dataset['q'][i, j],\n",
    "            self.dataset['r'][i, j],\n",
    "            self.dataset['omega'][i, j, 0],\n",
    "            self.dataset['omega'][i, j, 1],\n",
    "            self.dataset['omega'][i, j, 2],\n",
    "            self.dataset['omega'][i, j, 3],\n",
    "#             self.dataset['Mx_ext'][i],\n",
    "#             self.dataset['My_ext'][i],\n",
    "#             self.dataset['Mz_ext'][i]\n",
    "        ], dtype=torch.float32)\n",
    "        \n",
    "        U = torch.tensor([\n",
    "            self.dataset['u'][i, j, 0],\n",
    "            self.dataset['u'][i, j, 1],\n",
    "            self.dataset['u'][i, j, 2],\n",
    "            self.dataset['u'][i, j, 3]\n",
    "        ], dtype=torch.float32)\n",
    "        \n",
    "        return X, U\n",
    "    \n",
    "# trajectories containing 199 points\n",
    "dataset_path = 'datasets/HOVER_TO_HOVER_NOMINAL.npz'\n",
    "\n",
    "dataset = dict()\n",
    "print('loading dataset...')\n",
    "with np.load(dataset_path) as full_dataset:\n",
    "    # total number of trajectories\n",
    "    num = len(full_dataset['dx'])\n",
    "    print(num, 'trajectories')\n",
    "    dataset = {key: full_dataset[key] for key in [\n",
    "        't', 'dx', 'dy', 'dz', 'vx', 'vy', 'vz', 'phi', 'theta', 'psi', 'p', 'q', 'r','omega', 'u', 'omega_min','omega_max', 'k_omega', 'Mx_ext', 'My_ext', 'Mz_ext'\n",
    "    ]}\n",
    "\n",
    "# train/test split\n",
    "batchsize_train = 256\n",
    "batchsize_val = 4096\n",
    "train_trajectories = range(int(0.8*num))\n",
    "test_trajectories = list(set(range(num)) - set(train_trajectories))\n",
    "\n",
    "train_indices = [(i, j) for i in train_trajectories for j in range(199)]\n",
    "train_set = TrajectoryDataset(dataset, train_indices)\n",
    "train_loader = DataLoader(train_set, batch_size=batchsize_train, shuffle=True, num_workers=1)\n",
    "\n",
    "test_indices = [(i, j) for i in test_trajectories for j in range(199)]\n",
    "test_set = TrajectoryDataset(dataset, test_indices)\n",
    "test_loader = DataLoader(test_set, batch_size=batchsize_val, shuffle=True, num_workers=1)\n",
    "\n",
    "print('ready')\n",
    "\n",
    "print('Amount of testing trajectories: ',len(test_trajectories),f'(Batchsize: {batchsize_val})')\n",
    "print('Amount of training trajectories: ',len(train_trajectories),f'(Batchsize: {batchsize_train})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(test_trajectories))\n",
    "print(len(train_trajectories))\n",
    "\n",
    "print(dataset['omega_min'])\n",
    "print(dataset['omega_max'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Calculating mean and standard deviation for normalization**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "X_mean = torch.zeros(16)\n",
    "X_std = torch.zeros(16)\n",
    "\n",
    "N=10000\n",
    "\n",
    "for i, data in tqdm(enumerate(test_set)):\n",
    "    X = data[0]\n",
    "    X_mean += X\n",
    "    if i>=N:\n",
    "        break\n",
    "X_mean = X_mean/N\n",
    "\n",
    "print('mean:')\n",
    "print(X_mean)\n",
    "    \n",
    "for i, data in tqdm(enumerate(test_set)):\n",
    "    X = data[0]\n",
    "    X_std += (X-X_mean)**2\n",
    "    if i>=N:\n",
    "        break\n",
    "\n",
    "X_std = torch.sqrt(X_std/N)\n",
    "print('std:')\n",
    "print(X_std)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Defining the neural network**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    Normalize(mean=X_mean, std=X_std),\n",
    "    nn.Linear(16, 120),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(120, 120),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(120, 120),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(120, 4),\n",
    "    nn.Sigmoid()\n",
    ")\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Testing**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1 = torch.randn(16)\n",
    "print(model(x1))\n",
    "# print([param.shape for param in model.parameters()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Define a Loss function and optimizer**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "# optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)\n",
    "learning_rate = 0.0001\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)\n",
    "\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.9, patience=1, verbose=True, threshold=0.001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Training loop**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import time\n",
    "import copy\n",
    "\n",
    "loss_list = []\n",
    "loss_val_list = []\n",
    "best_loss = 0.1\n",
    "first = True\n",
    "start_time = time.time()\n",
    "\n",
    "# loop over the dataset multiple times\n",
    "num_epochs = 100\n",
    "\n",
    "nn_model_name = f\"{dataset_path[9:-4]}_{batchsize_train}_{batchsize_val}_{learning_rate}_{num_epochs}\"\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    \n",
    "    if first:\n",
    "        time_remaining = '-'\n",
    "    else:\n",
    "        time_estimate = epoch_time*(num_epochs-epoch+1)\n",
    "        if time_estimate > 60:\n",
    "            if time_estimate > 3600:\n",
    "                time_remaining = str(round(time_estimate/3600,2))+' h'\n",
    "            else:\n",
    "                time_remaining = str(round(time_estimate/60,2))+' min'\n",
    "        else:\n",
    "            time_remaining = str(round(time_estimate,0))+' s'\n",
    "        \n",
    "    first = False\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs}, Current learning rate: {optimizer.state_dict()['param_groups'][0]['lr']}, Time remaining: {time_remaining}\")\n",
    "\n",
    "    start_time_epoch = time.time()\n",
    "    \n",
    "    loop = tqdm(enumerate(train_loader), total=len(train_loader), leave=False)\n",
    "    \n",
    "    for i, (data, targets) in loop:\n",
    "        \n",
    "        # Zero the parameter gradients\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # Forward pass\n",
    "        outputs = model(data)\n",
    "        \n",
    "        # Loss\n",
    "        loss = criterion(outputs, targets)\n",
    "        \n",
    "        # Backward pass\n",
    "        loss.backward()\n",
    "        \n",
    "        # Update weights\n",
    "        optimizer.step()\n",
    "        \n",
    "        # Update progressbar\n",
    "        loop.set_description(f\"Epoch [{epoch+1}/{num_epochs}]\")\n",
    "        loop.set_postfix(loss=loss.item())\n",
    "        loss_list.append(loss.item())\n",
    "\n",
    "    # Validate\n",
    "    with torch.no_grad():\n",
    "        # Get a random batch from the test dataset\n",
    "        data_val, targets_val = next(iter(test_loader))\n",
    "\n",
    "        # Forward pass\n",
    "        outputs_val = model(data_val)\n",
    "\n",
    "        # Loss\n",
    "        loss_val = criterion(outputs_val, targets_val)\n",
    "\n",
    "        if loss_val < best_loss:\n",
    "            # Save best model\n",
    "            best_model = copy.deepcopy(model)\n",
    "            \n",
    "            # Backup\n",
    "            torch.save(model, 'neural_networks/tmp_benchmark.pt')\n",
    "            \n",
    "            best_loss = loss_val\n",
    "            print(\"Best model updated!\")\n",
    "\n",
    "        # Scheduler (reduce learning rate if loss stagnates)\n",
    "        scheduler.step(loss_val)\n",
    "        \n",
    "        loss_val_list.append(loss_val.item())\n",
    "\n",
    "    print(f'loss = {loss:.8f}, loss validation = {loss_val:.8f} '+r' (control error: +/-'+str(round(100*np.sqrt(float(loss_val)),2))+'%)\\n')\n",
    "\n",
    "    epoch_time = (time.time() - start_time_epoch)\n",
    "\n",
    "    loop.close()\n",
    "    \n",
    "# Compute excecution time\n",
    "execution_time = (time.time() - start_time)    \n",
    "print(f\"Total training time: {round(execution_time,2)}s\")\n",
    "\n",
    "# Save best model and copy for maptorange network\n",
    "torch.save(best_model, f'neural_networks/{nn_model_name}.pt')\n",
    "best_model_for_maptorange = torch.load('neural_networks/tmp_benchmark.pt')\n",
    "print(best_model_for_maptorange)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Testing performance**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "loader = test_loader\n",
    "# loop over the test dataset\n",
    "loop = tqdm(enumerate(loader), total=len(loader), leave=False)\n",
    "running_loss = 0\n",
    "\n",
    "for i, (data, targets) in loop:\n",
    "    outputs = model(data)\n",
    "    loss = criterion(outputs, targets)\n",
    "    \n",
    "    running_loss += loss.item()\n",
    "    \n",
    "    # update progressbar\n",
    "    loop.set_postfix(loss=loss.item())\n",
    "\n",
    "loop.close()\n",
    "print('average loss =', running_loss/len(loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Saving model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model, 'neural_networks/HOVER_TO_HOVER_NOMINAL.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.load('neural_networks/HOVER_TO_HOVER_NOMINAL.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_processed_output =nn.Sequential(\n",
    "    *model,\n",
    "    MapToRange(dataset['omega_min'], dataset['omega_max'])\n",
    ")\n",
    "\n",
    "print(model_processed_output)\n",
    "torch.save(model_processed_output, 'neural_networks/HOVER_TO_HOVER_NOMINAL_.pt')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}