train_lstm_mlflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6b32ef6b",
   "metadata": {},
   "source": [
    "# Train LSTM Model by MLFlow and PyTorch Lightning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47efbd58",
   "metadata": {},
   "outputs": [],
   "source": [
    "from argparse import ArgumentParser\n",
    "import os\n",
    "\n",
    "from loguru import logger\n",
    "import mlflow\n",
    "from mlflow.tracking import MlflowClient\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytorch_lightning as pl\n",
    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix\n",
    "import torch\n",
    "\n",
    "from models import HarLSTM, ModelUtils\n",
    "from pl_data import HarDataModule\n",
    "from utils import FeatUtils\n",
    "from chart import plot_feat_tensor_chart, plot_conf_matrix_chart\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22c0bb44",
   "metadata": {},
   "source": [
    "# 1. Prepare features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a968979",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir_path = \"./data/har_dataset\"\n",
    "batch_size = 16\n",
    "data_module = HarDataModule(data_dir_path, \n",
    "                            batch_size=batch_size,\n",
    "                           normalize=\"std\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e37881b",
   "metadata": {},
   "source": [
    "# 2. Define Network parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7f44a1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lstr_args = ['--max_epochs','50',\n",
    "            '--gpus', '1',\n",
    "             '--batch_size', '16',\n",
    "             '--stochastic_weight_avg', 'True',\n",
    "             '--gradient_clip_val', '5',\n",
    "             '--gradient_clip_algorithm', 'norm',\n",
    "            # DEBUGGING https://pytorch-lightning.readthedocs.io/en/latest/common/debugging.html\n",
    "            # don't forget to turn it off after debugging, slows things down a lot.\n",
    "            # '--profiler', 'pytorch', # issue no.3\n",
    "            # '--log_gpu_memory', 'all',\n",
    "            # '--limit_train_batches', '3',\n",
    "            # '--limit_predict_batches', '3',\n",
    "            # '--overfit_batches', '3',\n",
    "            # Inspect gradient norms\n",
    "            # about 10% performance hit, let's do it always anyway.\n",
    "            # '--track_grad_norm', '2',\n",
    "             ]\n",
    "\n",
    "parser = ArgumentParser()\n",
    "parser.add_argument('--batch_size', default=16, type=int)\n",
    "parser = pl.Trainer.add_argparse_args(parser)\n",
    "args = parser.parse_args(lstr_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecafbe69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if GPU is available\n",
    "use_gpu = torch.cuda.is_available()\n",
    "if(use_gpu):\n",
    "    print('Training on GPU!')\n",
    "else: \n",
    "    print('No GPU available, training on CPU; consider making n_epochs very small.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff6a54de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the model w/ hyperparams\n",
    "input_size = 9\n",
    "output_size = 6\n",
    "n_hidden = 128\n",
    "n_layers = 2\n",
    "\n",
    "# training params\n",
    "epochs = 50\n",
    "lr=0.0001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86ad9552",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = HarLSTM(input_size, output_size, n_hidden=n_hidden, n_layers=n_layers)\n",
    "print(\"Model information:\")\n",
    "print(net)\n",
    "trainer = pl.Trainer.from_argparse_args(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e81e78f4",
   "metadata": {},
   "source": [
    "# 3. Train the model by MLFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1e4599f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define helper functions\n",
    "def log_model_params_step(net):\n",
    "    mlflow.log_param(\"model_type\", type(net))\n",
    "    mlflow.log_param(\"n_layers\", net.n_layers)\n",
    "    mlflow.log_param(\"n_hidden\", net.n_hidden)\n",
    "    mlflow.log_param(\"drop_prob\", net.drop_prob)\n",
    "    mlflow.log_param(\"input_size\", net.input_size)\n",
    "\n",
    "def save_scaler_step(scaler, scaler_path=\"scaler.pkl\"):\n",
    "    FeatUtils.save_feat_scaler(scaler, scaler_path)\n",
    "    mlflow.log_artifact(scaler_path, artifact_path=\"model\")\n",
    "    os.remove(scaler_path)\n",
    "\n",
    "def show_data_sample_step(data_module, n_sample=4, fig_dir_path=\"figures/data_sample\"):\n",
    "    \"\"\"Save some data as figures from each set (train/valid/test)\"\"\"\n",
    "    loader_dict = {\n",
    "        \"train\": data_module.train_dataloader(),\n",
    "        \"test\": data_module.test_dataloader(),\n",
    "        \"valid\": data_module.val_dataloader(),\n",
    "    }\n",
    "    \n",
    "    for data_type, loader in loader_dict.items():\n",
    "        inputs, labels = iter(loader).next()\n",
    "\n",
    "        for i in range(n_sample):\n",
    "            chart = plot_feat_tensor_chart(inputs[i], labels[i])\n",
    "            chart_path = f\"train_sample_{i}.html\"\n",
    "            chart.save(chart_path, embed_options={\"renderer\":\"svg\"})\n",
    "            mlflow.log_artifact(chart_path, artifact_path=fig_dir_path)\n",
    "            os.remove(chart_path)\n",
    "\n",
    "def evaluate_model_step(net, data_module, batch_size, fig_dir_path=\"figures/evaluation\", use_gpu=True):\n",
    "    \"\"\"Evaluate a model and save a confusion matrix chart\"\"\"\n",
    "    test_loader = data_module.test_dataloader()\n",
    "    test_loss, test_labels, preds  = ModelUtils.test_net(net, net.criterion, test_loader, batch_size, use_gpu=use_gpu)\n",
    "    \n",
    "    acc = accuracy_score(test_labels, preds)\n",
    "    prec, recall, f1, _ = precision_recall_fscore_support(test_labels, preds, average=\"macro\")\n",
    "    \n",
    "    mlflow.log_metric(\"acc\", acc)\n",
    "    mlflow.log_metric(\"prec\", prec)\n",
    "    mlflow.log_metric(\"recall\", recall)\n",
    "    mlflow.log_metric(\"f1\", f1)\n",
    "    \n",
    "    # Let's save both html and png formats\n",
    "    chart = plot_conf_matrix_chart(test_labels, preds)\n",
    "    html_chart_path = \"conf_matrix.html\"\n",
    "    chart.save(html_chart_path, embed_options={\"renderer\":\"svg\"})\n",
    "    mlflow.log_artifact(html_chart_path, artifact_path=fig_dir_path)\n",
    "    \n",
    "    png_chart_path = \"conf_matrix.png\"\n",
    "    chart.save(png_chart_path)\n",
    "    mlflow.log_artifact(png_chart_path, artifact_path=fig_dir_path)\n",
    "    \n",
    "    os.remove(html_chart_path)\n",
    "    os.remove(png_chart_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "405a07bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_name = \"HAR_LSTM_Experiment\"\n",
    "mlflow_uri = \"http://mlflow_tracker:5000\"\n",
    "mlflow.set_tracking_uri(mlflow_uri)\n",
    "\n",
    "mlflow.set_experiment(experiment_name)\n",
    "\n",
    "tracking_uri = mlflow.get_tracking_uri()\n",
    "print(\"Current tracking uri: {}\".format(tracking_uri))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10fe1ec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow_run_name = \"HAR_LSTM_Training\"\n",
    "mlflow.pytorch.autolog()\n",
    "\n",
    "# Train the model\n",
    "with mlflow.start_run(run_name=mlflow_run_name) as run:\n",
    "    artifact_uri = mlflow.get_artifact_uri()\n",
    "    print(\"Current artifact uri: {}\".format(artifact_uri))\n",
    "    \n",
    "    log_model_params_step(net)\n",
    "    \n",
    "    mlflow.log_param(\"batch_size\", batch_size)\n",
    "    mlflow.log_param(\"train_val_ratio\", data_module.train_val_ratio)\n",
    "    mlflow.log_param(\"scaler\", type(data_module.scaler) if data_module.scaler is not None else None)\n",
    "    trainer.fit(net, datamodule=data_module)\n",
    "    # Run this for calculating \"test_loss\" metric.\n",
    "    # It will be automatically pushed to MLFlow tracker.\n",
    "    trainer.test(ckpt_path=\"best\", datamodule=data_module)\n",
    "    \n",
    "    show_data_sample_step(data_module)\n",
    "    save_scaler_step(data_module.scaler)\n",
    "    evaluate_model_step(net, data_module, batch_size, use_gpu=use_gpu)\n",
    "    \n",
    "    logger.info(\"MLFlow finished!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d724103a",
   "metadata": {},
   "source": [
    "# 4. Load the previously model and test manually"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c90c678",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_id = run.info.run_id\n",
    "model_uri = f\"runs:/{run_id}/model\"\n",
    "loaded_net = mlflow.pytorch.load_model(model_uri=model_uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff47df4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = data_module.test_dataloader()\n",
    "test_loss, test_labels, preds  = ModelUtils.test_net(loaded_net, loaded_net.criterion, test_loader, batch_size, use_gpu=use_gpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f8fb7c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "acc = accuracy_score(test_labels, preds)\n",
    "prec, recall, f1, _ = precision_recall_fscore_support(test_labels, preds, average=\"macro\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50156d50",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"accuracy: {acc}\")\n",
    "print(f\"precision: {prec}\")\n",
    "print(f\"recall: {recall}\")\n",
    "print(f\"f1: {f1}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dc1b761",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}