ECG_classification_tutorial.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "cell_execution_strategy": "setup"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "### Mount Google Drive"
      ],
      "metadata": {
        "id": "m8Nu4pwjGLIi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UNTpsolU4Iq5",
        "outputId": "dbb3cc6e-c9e8-46b5-c167-a83a23ffd432"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Change Directory\n",
        "%cd '/content/drive/My Drive/dsp/project'"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6Pi0AAtH6B62",
        "outputId": "0567f048-0851-4c3e-8123-ea720f5487cd"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/My Drive/dsp/project\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Install Required Libraries\n"
      ],
      "metadata": {
        "id": "_PbpzSbq8PWo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pip install tensorboardX"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "scgWhdfeG85h",
        "outputId": "c6f844c2-bcb2-45e1-daa0-9afbb2b2e5ae"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting tensorboardX\n",
            "  Downloading tensorboardX-2.6.2.2-py2.py3-none-any.whl.metadata (5.8 kB)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from tensorboardX) (1.26.4)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from tensorboardX) (24.1)\n",
            "Requirement already satisfied: protobuf>=3.20 in /usr/local/lib/python3.10/dist-packages (from tensorboardX) (3.20.3)\n",
            "Downloading tensorboardX-2.6.2.2-py2.py3-none-any.whl (101 kB)\n",
            "\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/101.7 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m101.7/101.7 kB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: tensorboardX\n",
            "Successfully installed tensorboardX-2.6.2.2\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import argparse\n",
        "import os\n",
        "import time\n",
        "import torch\n",
        "import numpy as np\n",
        "from tensorboardX import SummaryWriter\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import os\n",
        "import argparse\n",
        "import torch\n",
        "import time\n",
        "import numpy as np\n",
        "from sklearn import preprocessing\n",
        "from imblearn.over_sampling import RandomOverSampler\n",
        "from torch.utils.data import DataLoader\n",
        "import random\n",
        "from sklearn.metrics import f1_score"
      ],
      "metadata": {
        "id": "p1gBPtu5G9D2"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Set Training Parameters\n"
      ],
      "metadata": {
        "id": "wvrZvH-48uLT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Options:\n",
        "    def __init__(self):\n",
        "        pass\n",
        "\n",
        "    def init(self, parser):\n",
        "        # Global settings\n",
        "        parser.add_argument('--batch_size', type=int, default=256,\n",
        "                            help='Batch size for training and validation.')\n",
        "        parser.add_argument('--nepoch', type=int, default=50,\n",
        "                            help='Number of training epochs.')\n",
        "        parser.add_argument('--lr_initial', type=float, default=1e-4,\n",
        "                            help='Initial learning rate for the optimizer.')\n",
        "        parser.add_argument('--decay_epoch', type=int, default=20,\n",
        "                            help='Epoch at which to start decaying the learning rate.')\n",
        "\n",
        "        # Device settings\n",
        "        parser.add_argument('--device', type=str, default='cuda',\n",
        "                            help='Device to use for training (\"cuda\" for GPU, \"cpu\" for CPU).')\n",
        "\n",
        "        # Model settings\n",
        "        parser.add_argument('--classes', type=int, default=5,\n",
        "                            help='Number of output classes for classification.')\n",
        "\n",
        "        # Pretrained model settings\n",
        "        parser.add_argument('--log_name', type=str, default='241111',\n",
        "                            help='Identifier for logging and checkpointing.')\n",
        "        parser.add_argument('--pretrained', type=bool, default=False,\n",
        "                            help='Whether to load a pretrained model (True/False).')\n",
        "        parser.add_argument('--pretrained_model', type=str,\n",
        "                            default='./log/241111/models/ckpt_opt.pt',\n",
        "                            help='Path to the pretrained model weights file.')\n",
        "\n",
        "        # Dataset settings\n",
        "        parser.add_argument('--fs', type=int, default=360,\n",
        "                            help='Sampling frequency of the ECG data.')\n",
        "        parser.add_argument('--path_train_data', type=str,\n",
        "                            default='./dataset/train_data.npy',\n",
        "                            help='Path to save the training data.')\n",
        "        parser.add_argument('--path_train_labels', type=str,\n",
        "                            default='./dataset/train_labels.npy',\n",
        "                            help='Path to save the training labels.')\n",
        "        parser.add_argument('--path_val_data', type=str,\n",
        "                            default='./dataset/val_data.npy',\n",
        "                            help='Path to save the validation data.')\n",
        "        parser.add_argument('--path_val_labels', type=str,\n",
        "                            default='./dataset/val_labels.npy',\n",
        "                            help='Path to save the validation labels.')\n",
        "        parser.add_argument('--path_test_data', type=str,\n",
        "                            default='./dataset/test_data.npy',\n",
        "                            help='Path to save the test data.')\n",
        "        parser.add_argument('--path_test_labels', type=str,\n",
        "                            default='./dataset/test_labels.npy',\n",
        "                            help='Path to save the test labels.')\n",
        "\n",
        "\n",
        "        return parser\n",
        "\n",
        "\n",
        "parser = argparse.ArgumentParser(description='Options for ECG classification')\n",
        "opt = Options().init(parser).parse_known_args()\n",
        "print(opt)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wh1qL8YxG9L6",
        "outputId": "5e2d089c-4a76-4e7b-87ba-f5d4ec3dfb64"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(Namespace(batch_size=256, nepoch=50, lr_initial=0.0001, decay_epoch=20, device='cuda', classes=5, log_name='241111', pretrained=False, pretrained_model='./log/241111/models/ckpt_opt.pt', fs=360, path_train_data='./dataset/train_data.npy', path_train_labels='./dataset/train_labels.npy', path_val_data='./dataset/val_data.npy', path_val_labels='./dataset/val_labels.npy', path_test_data='./dataset/test_data.npy', path_test_labels='./dataset/test_labels.npy'), ['-f', '/root/.local/share/jupyter/runtime/kernel-29866c30-64e9-4695-bdf8-621fcb820a2c.json'])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Helper Functions"
      ],
      "metadata": {
        "id": "JjZPK3Ls82Er"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# For dataset\n",
        "class ECGDataloader():  # 1110 - 4096 samples\n",
        "    def __init__(self, data, label):\n",
        "        self.data = data\n",
        "        self.label = label\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        return (torch.tensor(self.data[index], dtype=torch.float), torch.tensor(self.label[index], dtype=torch.float))\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "# For dataset\n",
        "def label2index(i):\n",
        "    m = {'N': 0, 'S': 1, 'V': 2, 'F': 3, 'Q': 4}  # uncomment for 5 classes\n",
        "    return m[i]\n",
        "\n",
        "\n",
        "# Create a new directory.\n",
        "def mkdir(path):\n",
        "    if not os.path.exists(path):\n",
        "        os.makedirs(path)\n",
        "\n",
        "\n",
        "# Normalize the ECG data using Z-score normalization.\n",
        "def normalize_ecg(ecg_data):\n",
        "    mean = np.mean(ecg_data, axis=0, keepdims=True)\n",
        "    std = np.std(ecg_data, axis=0, keepdims=True)\n",
        "    return (ecg_data - mean) / (std + 1e-8)  # Prevent division by zero\n",
        "\n",
        "\n",
        "# for using pre-training weights\n",
        "def optimizer_to(optim, device):\n",
        "    for param in optim.state.values():\n",
        "        # Not sure there are any global tensors in the state dict\n",
        "        if isinstance(param, torch.Tensor):\n",
        "            param.data = param.data.to(device)\n",
        "            if param._grad is not None:\n",
        "                param._grad.data = param._grad.data.to(device)\n",
        "        elif isinstance(param, dict):\n",
        "            for subparam in param.values():\n",
        "                if isinstance(subparam, torch.Tensor):\n",
        "                    subparam.data = subparam.data.to(device)\n",
        "                    if subparam._grad is not None:\n",
        "                        subparam._grad.data = subparam._grad.data.to(device)\n",
        "\n",
        "\n",
        "# Calculate total number of parameters in a model.\n",
        "def cal_total_params(our_model):\n",
        "    total_parameters = 0\n",
        "    for variable in our_model.parameters():\n",
        "        shape = variable.size()\n",
        "        variable_parameters = 1\n",
        "        for dim in shape:\n",
        "            variable_parameters *= dim\n",
        "        total_parameters += variable_parameters\n",
        "\n",
        "    return total_parameters\n",
        "\n",
        "\n",
        "# Display a progress bar during training/validation.\n",
        "class Bar(object):\n",
        "    def __init__(self, dataloader):\n",
        "        if not hasattr(dataloader, 'dataset'):\n",
        "            raise ValueError('Attribute `dataset` not exists in dataloder.')\n",
        "        if not hasattr(dataloader, 'batch_size'):\n",
        "            raise ValueError('Attribute `batch_size` not exists in dataloder.')\n",
        "\n",
        "        self.dataloader = dataloader\n",
        "        self.iterator = iter(dataloader)\n",
        "        self.dataset = dataloader.dataset\n",
        "        self.batch_size = dataloader.batch_size\n",
        "        self._idx = 0\n",
        "        self._batch_idx = 0\n",
        "        self._time = []\n",
        "        self._DISPLAY_LENGTH = 50\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.dataloader)\n",
        "\n",
        "    def __iter__(self):\n",
        "        return self\n",
        "\n",
        "    def __next__(self):\n",
        "        if len(self._time) < 2:\n",
        "            self._time.append(time.time())\n",
        "\n",
        "        self._batch_idx += self.batch_size\n",
        "        if self._batch_idx > len(self.dataset):\n",
        "            self._batch_idx = len(self.dataset)\n",
        "\n",
        "        try:\n",
        "            batch = next(self.iterator)\n",
        "            self._display()\n",
        "        except StopIteration:\n",
        "            raise StopIteration()\n",
        "\n",
        "        self._idx += 1\n",
        "        if self._idx >= len(self.dataloader):\n",
        "            self._reset()\n",
        "\n",
        "        return batch\n",
        "\n",
        "    def _display(self):\n",
        "        if len(self._time) > 1:\n",
        "            t = (self._time[-1] - self._time[-2])\n",
        "            eta = t * (len(self.dataloader) - self._idx)\n",
        "        else:\n",
        "            eta = 0\n",
        "\n",
        "        rate = self._idx / len(self.dataloader)\n",
        "        len_bar = int(rate * self._DISPLAY_LENGTH)\n",
        "        bar = ('=' * len_bar + '>').ljust(self._DISPLAY_LENGTH, '.')\n",
        "        idx = str(self._batch_idx).rjust(len(str(len(self.dataset))), ' ')\n",
        "\n",
        "        tmpl = '\\r{}/{}: [{}] - ETA {:.1f}s'.format(\n",
        "            idx,\n",
        "            len(self.dataset),\n",
        "            bar,\n",
        "            eta\n",
        "        )\n",
        "        print(tmpl, end='')\n",
        "        if self._batch_idx == len(self.dataset):\n",
        "            print()\n",
        "\n",
        "    def _reset(self):\n",
        "        self._idx = 0\n",
        "        self._batch_idx = 0\n",
        "        self._time = []\n",
        "\n",
        "\n",
        "# Define a custom writer class that extends SummaryWriter to log training/validation metrics.\n",
        "class Writer(SummaryWriter):\n",
        "    def __init__(self, logdir):\n",
        "        super(Writer, self).__init__(logdir)\n",
        "\n",
        "    # Method to log training loss.\n",
        "    def log_train_loss(self, loss_type, train_loss, step):\n",
        "        self.add_scalar('train_{}_loss'.format(loss_type), train_loss, step)\n",
        "\n",
        "    # Method to log validation loss.\n",
        "    def log_valid_loss(self, loss_type, valid_loss, step):\n",
        "        self.add_scalar('valid_{}_loss'.format(loss_type), valid_loss, step)\n",
        "\n",
        "    # Method to log other performance metrics (e.g., accuracy, F1-score).\n",
        "    def log_score(self, metrics_name, metrics, step):\n",
        "        # Add a scalar value to the writer with the given metric name.\n",
        "        self.add_scalar(metrics_name, metrics, step)\n",
        "\n",
        "def save_checkpoint(exp_log_dir, model, epoch):\n",
        "    save_dict = {\n",
        "        \"model\": model.state_dict(),\n",
        "        'epoch': epoch\n",
        "    }\n",
        "    save_path = os.path.join(exp_log_dir, \"ckpt_opt.pt\")\n",
        "\n",
        "    torch.save(save_dict, save_path)"
      ],
      "metadata": {
        "id": "xLWVggP6HBmh"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Define DNN Model\n"
      ],
      "metadata": {
        "id": "WV2obcO4865u"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class SimpleCNN(nn.Module):\n",
        "    def __init__(self, opt, in_ch=1, out_ch=64, in_len=360):\n",
        "        super(SimpleCNN, self).__init__()\n",
        "\n",
        "        # Convolutional Layer 1 (Reduced model complexity)\n",
        "        self.conv1 = nn.Conv1d(in_channels=in_ch, out_channels=out_ch, kernel_size=5, stride=2, padding=2)\n",
        "        self.bn1 = nn.BatchNorm1d(out_ch)\n",
        "\n",
        "        # Fully Connected Layer\n",
        "        self.fc1 = nn.Linear(out_ch * (in_len // 2), opt.classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(self.bn1(self.conv1(x)))\n",
        "\n",
        "        # Flatten before fully connected layer\n",
        "        x = x.view(x.size(0), -1)\n",
        "\n",
        "        # Fully connected layer\n",
        "        x = self.fc1(x)\n",
        "\n",
        "        return x"
      ],
      "metadata": {
        "id": "FbkIYctIHF8m"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training the Model\n"
      ],
      "metadata": {
        "id": "jtHoqZT_8_pt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "class Trainer:\n",
        "    def __init__(self, opt):\n",
        "        self.opt = opt\n",
        "        self.model = SimpleCNN(opt).to(opt.device)\n",
        "        self.loss_fn = torch.nn.CrossEntropyLoss()\n",
        "        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=opt.lr_initial)\n",
        "        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=opt.decay_epoch, gamma=0.1)\n",
        "        self.writer = Writer(self._get_tboard_dir())\n",
        "        self.train_loader, self.valid_loader = self._load_data()\n",
        "        self.log_file_path = os.path.join(self._get_tboard_dir(), 'training_log.txt')\n",
        "\n",
        "        # Load pretrained model if specified\n",
        "        if self.opt.pretrained:\n",
        "            self._load_pretrained_model()\n",
        "\n",
        "    # Z-score normalization\n",
        "    def _normalize_ecg(self, ecg_data, axis=1):\n",
        "        mean = np.mean(ecg_data, axis=axis, keepdims=True)\n",
        "        std = np.std(ecg_data, axis=axis, keepdims=True)\n",
        "        return (ecg_data - mean) / (std + 1e-8)  # Prevent division by zero\n",
        "\n",
        "    def _load_data(self):\n",
        "        # Load and preprocess data\n",
        "        train_data = np.load(self.opt.path_train_data)  # Load ECG train data\n",
        "        train_labels = np.load(self.opt.path_train_labels)  # Load train labels\n",
        "\n",
        "        val_data = np.load(self.opt.path_val_data)  # Load ECG validation data\n",
        "        val_labels = np.load(self.opt.path_val_labels)  # Load validation labels\n",
        "\n",
        "        Y_train = np.array([label2index(i) for i in train_labels])  # Convert labels to indices\n",
        "        Y_val = np.array([label2index(i) for i in val_labels])  # Convert labels to indices\n",
        "\n",
        "        # Normalize data along the time axis\n",
        "        train_data = self._normalize_ecg(train_data)\n",
        "        val_data = self._normalize_ecg(val_data)\n",
        "\n",
        "        # Expand dimensions to match model input requirements\n",
        "        X_train, X_val = np.expand_dims(train_data, 1), np.expand_dims(val_data, 1)\n",
        "\n",
        "        # Create DataLoader for training and validation\n",
        "        train_loader = DataLoader(ECGDataloader(X_train, Y_train), batch_size=self.opt.batch_size, shuffle=True,\n",
        "                                  num_workers=0, pin_memory=True, drop_last=True)\n",
        "        valid_loader = DataLoader(ECGDataloader(X_val, Y_val), batch_size=self.opt.batch_size, shuffle=False,\n",
        "                                  num_workers=0)\n",
        "\n",
        "        return train_loader, valid_loader\n",
        "\n",
        "\n",
        "    def _get_tboard_dir(self):\n",
        "        # Initialize directories for logging and model storage\n",
        "        log_dir = os.path.join(os.getcwd(), 'log', f'{self.opt.log_name}')\n",
        "        mkdir(log_dir)\n",
        "        mkdir(os.path.join(log_dir, 'logs'))\n",
        "        mkdir(os.path.join(log_dir, 'models'))\n",
        "        return os.path.join(log_dir, 'logs')\n",
        "\n",
        "    def _load_pretrained_model(self):\n",
        "        # Load pretrained model weights if specified\n",
        "        print('Loading the pretrained model...')\n",
        "        chkpt = torch.load(self.opt.pretrained_model)\n",
        "        self.model.load_state_dict(chkpt['model'])\n",
        "        self.optimizer.load_state_dict(chkpt['optimizer'])\n",
        "        optimizer_to(self.optimizer, self.opt.device)\n",
        "        print('Resuming Start Epoch:', chkpt['epoch'] + 1)\n",
        "\n",
        "    def train(self):\n",
        "        # Print the total number of parameters in the model\n",
        "        print(\n",
        "            f'Total parameters: {cal_total_params(self.model):,} ({cal_total_params(self.model) / 1e6:.2f}M)')\n",
        "        best_f1 = 0\n",
        "        for epoch in range(1, self.opt.nepoch + 1):\n",
        "            start_time = time.time()\n",
        "            self.model.train()\n",
        "            train_loss = 0\n",
        "\n",
        "            # Training loop\n",
        "            for X, Y in Bar(self.train_loader):\n",
        "                X, Y = X.float().to(self.opt.device), Y.long().to(self.opt.device)  # Move data to device\n",
        "\n",
        "                # Forward pass and optimization\n",
        "                outputs = self.model(X)  # Get model predictions\n",
        "                loss = self.loss_fn(outputs, Y)  # Calculate loss\n",
        "                self.optimizer.zero_grad()  # Clear previous gradients\n",
        "                loss.backward()  # Backpropagate to calculate gradients\n",
        "                self.optimizer.step()  # Update model parameters\n",
        "                train_loss += loss.item()\n",
        "\n",
        "            avg_train_loss = train_loss / len(self.train_loader)  # Calculate average training loss\n",
        "            self.writer.log_train_loss('total', avg_train_loss, epoch)  # Log training loss\n",
        "\n",
        "            # Validation\n",
        "            accuracy, f1, avg_val_loss = self._evaluate(self.valid_loader, epoch)  # Evaluate model on validation set\n",
        "            if f1 > best_f1:  # Save the best model based on accuracy\n",
        "                best_f1 = f1\n",
        "                save_checkpoint(self._get_model_dir(), self.model, epoch)\n",
        "\n",
        "            self.writer.log_score('F1-score', f1, epoch)  # Log validation accuracy\n",
        "            self.scheduler.step()  # Update learning rate scheduler\n",
        "\n",
        "            # Logging\n",
        "            log_message = (\n",
        "                f'EPOCH[{epoch}] Train Loss: {avg_train_loss:.6f} | Validation Loss: {avg_val_loss:.6f} | Validation F1-score: {f1:.6f} | Time: {time.time() - start_time:.3f}s'\n",
        "            )\n",
        "            print(log_message)\n",
        "            self._log_to_file(log_message)\n",
        "\n",
        "        print('Training completed.')\n",
        "\n",
        "    def _evaluate(self, dataloader, epoch):\n",
        "        # Evaluate the model on the given dataloader\n",
        "        self.model.eval()\n",
        "        pred_labels, true_labels = [], []\n",
        "        total_loss = 0\n",
        "        with torch.no_grad():\n",
        "            for X, Y in Bar(dataloader):\n",
        "                X, Y = X.float().to(self.opt.device), Y.long().to(self.opt.device)  # Move data to device\n",
        "                pred = self.model(X)  # Get model predictions\n",
        "                loss = self.loss_fn(pred, Y)  # Calculate loss\n",
        "                total_loss += loss.item()\n",
        "\n",
        "                # Get predicted class directly from raw logits\n",
        "                pred_classes = torch.argmax(pred, dim=1)\n",
        "                pred_labels.extend(pred_classes.cpu().numpy())\n",
        "                true_labels.extend(Y.cpu().numpy())\n",
        "\n",
        "        # Calculate accuracy\n",
        "        pred_labels = np.array(pred_labels)\n",
        "        true_labels = np.array(true_labels)\n",
        "        accuracy = np.mean(pred_labels == true_labels)  # Calculate accuracy\n",
        "        avg_valid_loss = total_loss / len(dataloader)  # Calculate average validation loss\n",
        "        f1 = f1_score(true_labels, pred_labels, average='macro')  # Calculate F1-score\n",
        "        self.writer.log_valid_loss('total', avg_valid_loss, epoch)  # Log validation loss\n",
        "        return accuracy, f1, avg_valid_loss\n",
        "\n",
        "    def _get_model_dir(self):\n",
        "        # Get directory path for saving models\n",
        "        log_dir = os.path.join(os.getcwd(), 'log', f'{self.opt.log_name}')\n",
        "        return os.path.join(log_dir, 'models')\n",
        "\n",
        "    def _log_to_file(self, message):\n",
        "        # Write log message to file\n",
        "        with open(self.log_file_path, 'a') as f:\n",
        "            f.write(message + '\\n')\n",
        "\n",
        "# Parse command-line arguments\n",
        "opt = Options().init(argparse.ArgumentParser(description='ECG Classification')).parse_known_args()\n",
        "print(opt[0])\n",
        "\n",
        "# Set random seeds for reproducibility\n",
        "torch.manual_seed(1234)\n",
        "np.random.seed(1234)\n",
        "random.seed(1234)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed_all(1234)\n",
        "\n",
        "# Initialize trainer and start training\n",
        "trainer = Trainer(opt[0])\n",
        "trainer.train()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nbSmPHxXHLk-",
        "outputId": "d68ef1b7-9ba6-4443-a8f4-df7000bc1aa1"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Namespace(batch_size=256, nepoch=50, lr_initial=0.0001, decay_epoch=20, device='cuda', classes=5, log_name='241111', pretrained=False, pretrained_model='./log/241111/models/ckpt_opt.pt', fs=360, path_train_data='./dataset/train_data.npy', path_train_labels='./dataset/train_labels.npy', path_val_data='./dataset/val_data.npy', path_val_labels='./dataset/val_labels.npy', path_test_data='./dataset/test_data.npy', path_test_labels='./dataset/test_labels.npy')\n",
            "Total parameters: 58,117 (0.06M)\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[1] Train Loss: 0.242441 | Validation Loss: 0.148391 | Validation F1-score: 0.737243 | Time: 2.272s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[2] Train Loss: 0.136498 | Validation Loss: 0.125626 | Validation F1-score: 0.795166 | Time: 2.432s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[3] Train Loss: 0.119535 | Validation Loss: 0.115019 | Validation F1-score: 0.815862 | Time: 2.281s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[4] Train Loss: 0.109964 | Validation Loss: 0.108146 | Validation F1-score: 0.828753 | Time: 2.686s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[5] Train Loss: 0.103390 | Validation Loss: 0.104043 | Validation F1-score: 0.848528 | Time: 3.526s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[6] Train Loss: 0.098314 | Validation Loss: 0.100249 | Validation F1-score: 0.858229 | Time: 4.089s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[7] Train Loss: 0.094824 | Validation Loss: 0.098361 | Validation F1-score: 0.855157 | Time: 3.743s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[8] Train Loss: 0.091860 | Validation Loss: 0.096786 | Validation F1-score: 0.857125 | Time: 3.032s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[9] Train Loss: 0.089241 | Validation Loss: 0.095744 | Validation F1-score: 0.853532 | Time: 4.548s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[10] Train Loss: 0.086679 | Validation Loss: 0.094612 | Validation F1-score: 0.859441 | Time: 3.168s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[11] Train Loss: 0.085496 | Validation Loss: 0.092405 | Validation F1-score: 0.864113 | Time: 2.657s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[12] Train Loss: 0.083414 | Validation Loss: 0.092942 | Validation F1-score: 0.863457 | Time: 2.288s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[13] Train Loss: 0.081452 | Validation Loss: 0.092676 | Validation F1-score: 0.859803 | Time: 3.052s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[14] Train Loss: 0.080166 | Validation Loss: 0.091846 | Validation F1-score: 0.865380 | Time: 3.304s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[15] Train Loss: 0.078406 | Validation Loss: 0.090663 | Validation F1-score: 0.873913 | Time: 2.823s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[16] Train Loss: 0.077107 | Validation Loss: 0.090668 | Validation F1-score: 0.864660 | Time: 2.948s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[17] Train Loss: 0.075975 | Validation Loss: 0.095160 | Validation F1-score: 0.872102 | Time: 3.316s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[18] Train Loss: 0.075171 | Validation Loss: 0.092186 | Validation F1-score: 0.863992 | Time: 4.191s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[19] Train Loss: 0.073381 | Validation Loss: 0.089642 | Validation F1-score: 0.866353 | Time: 2.316s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[20] Train Loss: 0.072909 | Validation Loss: 0.089223 | Validation F1-score: 0.873517 | Time: 2.301s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[21] Train Loss: 0.068792 | Validation Loss: 0.088176 | Validation F1-score: 0.870026 | Time: 2.482s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[22] Train Loss: 0.068198 | Validation Loss: 0.087715 | Validation F1-score: 0.870298 | Time: 2.271s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[23] Train Loss: 0.068053 | Validation Loss: 0.087556 | Validation F1-score: 0.873292 | Time: 2.809s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[24] Train Loss: 0.067886 | Validation Loss: 0.087699 | Validation F1-score: 0.873116 | Time: 3.383s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[25] Train Loss: 0.067754 | Validation Loss: 0.087691 | Validation F1-score: 0.871444 | Time: 2.929s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[26] Train Loss: 0.067650 | Validation Loss: 0.087638 | Validation F1-score: 0.870951 | Time: 3.969s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[27] Train Loss: 0.067615 | Validation Loss: 0.087594 | Validation F1-score: 0.876722 | Time: 4.945s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[28] Train Loss: 0.067468 | Validation Loss: 0.087508 | Validation F1-score: 0.872930 | Time: 2.775s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[29] Train Loss: 0.067417 | Validation Loss: 0.087472 | Validation F1-score: 0.872593 | Time: 2.429s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[30] Train Loss: 0.067135 | Validation Loss: 0.087636 | Validation F1-score: 0.873398 | Time: 2.262s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[31] Train Loss: 0.067207 | Validation Loss: 0.087490 | Validation F1-score: 0.872760 | Time: 2.281s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[32] Train Loss: 0.066964 | Validation Loss: 0.087341 | Validation F1-score: 0.874701 | Time: 2.376s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[33] Train Loss: 0.066891 | Validation Loss: 0.087582 | Validation F1-score: 0.871622 | Time: 3.426s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[34] Train Loss: 0.066766 | Validation Loss: 0.087326 | Validation F1-score: 0.874898 | Time: 2.353s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[35] Train Loss: 0.066638 | Validation Loss: 0.087210 | Validation F1-score: 0.874841 | Time: 2.286s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[36] Train Loss: 0.066490 | Validation Loss: 0.087149 | Validation F1-score: 0.875621 | Time: 2.282s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[37] Train Loss: 0.066420 | Validation Loss: 0.087224 | Validation F1-score: 0.872509 | Time: 2.465s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[38] Train Loss: 0.066266 | Validation Loss: 0.087532 | Validation F1-score: 0.874081 | Time: 2.828s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[39] Train Loss: 0.066196 | Validation Loss: 0.087090 | Validation F1-score: 0.874403 | Time: 2.901s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[40] Train Loss: 0.066064 | Validation Loss: 0.087288 | Validation F1-score: 0.874649 | Time: 2.305s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[41] Train Loss: 0.065660 | Validation Loss: 0.087044 | Validation F1-score: 0.873413 | Time: 2.453s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[42] Train Loss: 0.065504 | Validation Loss: 0.087049 | Validation F1-score: 0.873849 | Time: 2.262s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[43] Train Loss: 0.065559 | Validation Loss: 0.087065 | Validation F1-score: 0.871920 | Time: 2.297s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[44] Train Loss: 0.065479 | Validation Loss: 0.087036 | Validation F1-score: 0.874234 | Time: 3.059s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[45] Train Loss: 0.065552 | Validation Loss: 0.087032 | Validation F1-score: 0.872700 | Time: 2.894s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[46] Train Loss: 0.065542 | Validation Loss: 0.087017 | Validation F1-score: 0.873248 | Time: 2.236s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[47] Train Loss: 0.065572 | Validation Loss: 0.087021 | Validation F1-score: 0.874444 | Time: 2.248s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[48] Train Loss: 0.065444 | Validation Loss: 0.087021 | Validation F1-score: 0.874553 | Time: 2.259s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[49] Train Loss: 0.065544 | Validation Loss: 0.087014 | Validation F1-score: 0.874490 | Time: 2.712s\n",
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "EPOCH[50] Train Loss: 0.065500 | Validation Loss: 0.087031 | Validation F1-score: 0.873622 | Time: 3.137s\n",
            "Training completed.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Test the model"
      ],
      "metadata": {
        "id": "yzs4Qe9zzSE9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Tester:\n",
        "    def __init__(self, opt):\n",
        "        self.opt = opt\n",
        "        self.model = SimpleCNN(opt).to(opt.device)\n",
        "        self._load_pretrained_model()\n",
        "        self.test_loader = self._load_data()\n",
        "\n",
        "    # Z-score normalization\n",
        "    def _normalize_ecg(self, ecg_data, axis=1):\n",
        "        mean = np.mean(ecg_data, axis=axis, keepdims=True)\n",
        "        std = np.std(ecg_data, axis=axis, keepdims=True)\n",
        "        return (ecg_data - mean) / (std + 1e-8)  # Prevent division by zero\n",
        "\n",
        "    def _load_data(self):\n",
        "        # Load and preprocess data\n",
        "        test_data = np.load(self.opt.path_val_data)  # Load ECG validation data\n",
        "        test_labels = np.load(self.opt.path_val_labels)  # Load validation labels\n",
        "\n",
        "        test_data = self._normalize_ecg(test_data)\n",
        "        Y_test = np.array([label2index(i) for i in test_labels])\n",
        "\n",
        "        # Prepare DataLoader\n",
        "        X_test = np.expand_dims(test_data, 1)\n",
        "        test_loader = DataLoader(ECGDataloader(X_test, Y_test), batch_size=self.opt.batch_size, shuffle=False, num_workers=0)\n",
        "        return test_loader\n",
        "\n",
        "    def _load_pretrained_model(self):\n",
        "        # Load the pretrained model for evaluation\n",
        "        print('Loading the pretrained model...')\n",
        "        chkpt = torch.load(self.opt.pretrained_model, map_location=self.opt.device)\n",
        "        self.model.load_state_dict(chkpt['model'])\n",
        "\n",
        "    def test(self):\n",
        "        self.model.eval()\n",
        "        pred_labels, true_labels = [], []\n",
        "        total_loss = 0\n",
        "        loss_fn = torch.nn.CrossEntropyLoss()\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for X, Y in Bar(self.test_loader):\n",
        "                X, Y = X.float().to(self.opt.device), Y.long().to(self.opt.device)\n",
        "                outputs = self.model(X)\n",
        "                loss = loss_fn(outputs, Y)\n",
        "                total_loss += loss.item()\n",
        "\n",
        "                # Get predicted classes\n",
        "                pred_classes = torch.argmax(outputs, dim=1)\n",
        "                pred_labels.extend(pred_classes.cpu().numpy())\n",
        "                true_labels.extend(Y.cpu().numpy())\n",
        "\n",
        "        # Calculate accuracy\n",
        "        pred_labels = np.array(pred_labels)\n",
        "        true_labels = np.array(true_labels)\n",
        "        accuracy = np.mean(pred_labels == true_labels)\n",
        "        avg_test_loss = total_loss / len(self.test_loader)\n",
        "        f1 = f1_score(true_labels, pred_labels, average='macro')  # Calculate F1-score\n",
        "\n",
        "        # Print results in a more professional format\n",
        "        print(f'==================== Test Results ====================')\n",
        "        # print(f'| Test Accuracy    : {accuracy * 100:.2f}%')\n",
        "        print(f'| Test F1-score    : {f1 * 100:.2f}%')\n",
        "        print(f'=======================================================')\n",
        "\n",
        "# Parse command-line arguments\n",
        "opt = Options().init(argparse.ArgumentParser(description='ECG Classification')).parse_known_args()\n",
        "print(opt[0])\n",
        "\n",
        "torch.manual_seed(1234)\n",
        "np.random.seed(1234)\n",
        "random.seed(1234)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed_all(1234)\n",
        "\n",
        "tester = Tester(opt[0])\n",
        "tester.test()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ffxLDGHQzQ9u",
        "outputId": "8b50f6d4-d5a5-4fec-96e5-231fc061b0b6"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Namespace(batch_size=256, nepoch=50, lr_initial=0.0001, decay_epoch=20, device='cuda', classes=5, log_name='241111', pretrained=False, pretrained_model='./log/241111/models/ckpt_opt.pt', fs=360, path_train_data='./dataset/train_data.npy', path_train_labels='./dataset/train_labels.npy', path_val_data='./dataset/val_data.npy', path_val_labels='./dataset/val_labels.npy', path_test_data='./dataset/test_data.npy', path_test_labels='./dataset/test_labels.npy')\n",
            "Loading the pretrained model...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-12-37262cba5165>:30: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
            "  chkpt = torch.load(self.opt.pretrained_model, map_location=self.opt.device)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "21005/21005: [=================================================>] - ETA 0.0s\n",
            "==================== Test Results ====================\n",
            "| Test F1-score    : 87.67%\n",
            "=======================================================\n"
          ]
        }
      ]
    }
  ]
}