train_cont_forecasting_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from state_quantization.dataset import load_dataset\n",
    "from state_quantization.dataset import DynamicsModelDataset\n",
    "from state_quantization.train import train_model, test_step\n",
    "from state_quantization.eval import eval_model, compare_models\n",
    "from state_quantization.forcasting_models import LSTMForcasting, LSTMForcastingConstOut\n",
    "from state_quantization.trainer import NNTrainer\n",
    "from torch.optim.lr_scheduler import MultiStepLR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda:0\" if use_cuda else \"cpu\")\n",
    "print(f\"Using Device: {device}\")\n",
    "torch.backends.cudnn.benchmark = True"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load Dataset"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "dataset_input_key = 'merged_input'\n",
    "dataset_output_key = 'merged_output'\n",
    "dataset_file_path = 'tmp/ib-out/ib-samples-la.npy'\n",
    "normalized_data_params_save_path = 'tmp/transformer/NormalizeInputConfigs.pkl'\n",
    "dataset_device = 'cpu'\n",
    "y_indexes = [4, 6]\n",
    "\n",
    "train_dataset, val_dataset = load_dataset(file_path=dataset_file_path, input_key=dataset_input_key,\n",
    "                                          output_key=dataset_output_key, dataset_class=DynamicsModelDataset,\n",
    "                                          normalize=True, device=dataset_device, y_clip_range=y_indexes,\n",
    "                                          normalized_input_data_params_save_path=normalized_data_params_save_path)\n",
    "\n",
    "batch_size = 8000\n",
    "params = {'batch_size': batch_size,\n",
    "          'shuffle': True,\n",
    "          'num_workers': 0,\n",
    "          'drop_last': True,\n",
    "          'pin_memory': not train_dataset.x.is_cuda}\n",
    "print(params)\n",
    "train_loader = DataLoader(train_dataset, **params)\n",
    "val_loader = DataLoader(val_dataset, **params)\n",
    "train_dataset.y[0]"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Define model"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model_dict_path = 'state_quantization/model_no_disc_dict'\n",
    "model_path = 'state_quantization/model_no_disc'\n",
    "num_of_features = train_dataset.get_features_size()\n",
    "seq_len = train_dataset.get_seq_len()\n",
    "hidden_size = 20\n",
    "out_size = train_dataset.get_output_feature_size()\n",
    "print(f'Out Size:{out_size}')\n",
    "look_ahead = train_dataset.get_look_ahead_size()\n",
    "n_layers = 1\n",
    "dropout = 0.1\n",
    "print(look_ahead)\n",
    "model = forcasting_model = LSTMForcasting(features=num_of_features, hidden_size=hidden_size, out_size=out_size, seq_len=seq_len,\n",
    "                                  look_ahead=look_ahead, dropout=dropout, n_layers=n_layers).to(device=device)\n",
    "\n",
    "load_to_gpu = model.is_cuda() and not train_dataset.x.is_cuda\n",
    "print(load_to_gpu)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Train Model"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "lr = 1e-3\n",
    "\n",
    "lr_milestones=[24, 54]\n",
    "gamma = 0.1\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "lr_scheduler = MultiStepLR(optimizer, milestones=lr_milestones, gamma=gamma)\n",
    "trainer = NNTrainer(model=model,train_loader=train_loader, test_loader=val_loader,load_to_gpu=load_to_gpu,optimizer=optimizer, lr_scheduler=lr_scheduler)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "trainer.train(n_epochs=50)\n",
    "# train_model(model=model, train_loader=train_loader, test_loader=val_loader, n_epochs=100, learning_rate=1e-3,\n",
    "#             load_to_gpu=load_to_gpu, gamma=0.1, lr_milestones=[25, 200])"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Save Model"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\n",
    "torch.save(model.state_dict(), model_dict_path)\n",
    "#torch.save(model, model_path)\n",
    "torch.cuda.empty_cache()\n",
    "del model"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load Model for evaluation"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\n",
    "model_classes = [LSTMForcasting, LSTMForcastingConstOut]\n",
    "models = []\n",
    "for model_class in model_classes:\n",
    "    models.append(model_class(features=num_of_features, hidden_size=hidden_size, out_size=out_size,\n",
    "                              seq_len=seq_len, n_layers=n_layers, dropout=dropout,\n",
    "                              look_ahead=look_ahead).to(device))\n",
    "    models[-1].load_state_dict(torch.load(model_dict_path))"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "for la_model in models:\n",
    "    print(la_model.__class__.__name__)\n",
    "    test_step(model=la_model, data_loader=val_loader, loss_function=torch.nn.MSELoss(), load_to_gpu=load_to_gpu)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "#value_keys = ['setpoint', 'velocity', 'gain', 'shift', 'fatigue', 'consumption']\n",
    "value_keys = ['fatigue', 'consumption']\n",
    "\n",
    "models = []\n",
    "for model_class in model_classes:\n",
    "    models.append(model_class(features=num_of_features, hidden_size=hidden_size, out_size=out_size,\n",
    "                              seq_len=seq_len, n_layers=n_layers, dropout=dropout,\n",
    "                              look_ahead=look_ahead).to(dataset_device))\n",
    "    models[-1].load_state_dict(torch.load(model_dict_path))\n",
    "\n",
    "compare_models(models, x=val_dataset.x[78], y=val_dataset.y[78], value_keys=value_keys)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}