notebooks/TorchSig-Dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sigmf_pytorch_dataset import SigMFDataset\n",
    "from torchsig.utils.visualize import IQVisualizer, SpectrogramVisualizer, two_channel_to_complex\n",
    "from torchsig.utils.dataset import SignalDataset\n",
    "from torchsig.datasets.sig53 import Sig53\n",
    "from torch.utils.data import DataLoader\n",
    "from matplotlib import pyplot as plt\n",
    "from typing import List\n",
    "from tqdm import tqdm\n",
    "from datetime import datetime\n",
    "import numpy as np\n",
    "import os\n",
    "from pathlib import Path\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = \"./dev_data/torchsig_test/\"\n",
    "num_iq_samples = 2096"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logs_dir = datetime.now().strftime('logs/%H_%M_%S_%m_%d_%Y')\n",
    "logs_dir\n",
    "logs_dir = Path(logs_dir)\n",
    "logs_dir.mkdir(parents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset = SigMFDataset( root=dataset_path, sample_count= num_iq_samples, allowed_filetypes=[\".sigmf-data\"])\n",
    "dataset_class_counts = {class_name:0 for class_name in dataset.class_list}\n",
    "for data,label in dataset:\n",
    "    dataset_class_counts[dataset.class_list[label]] += 1\n",
    "print(f\"{len(dataset)=}\")\n",
    "print(dataset_class_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_loader = DataLoader(\n",
    "    dataset=dataset,\n",
    "    batch_size=100,\n",
    "    shuffle=True,\n",
    ")\n",
    "\n",
    "visualizer = IQVisualizer(\n",
    "    data_loader=data_loader\n",
    ")\n",
    "\n",
    "for figure in iter(visualizer):\n",
    "    figure.set_size_inches(16, 16)\n",
    "    plt.show()\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SigMF based Model Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchsig.models.iq_models.efficientnet.efficientnet import efficientnet_b4\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from torchsig.utils.cm_plotter import plot_confusion_matrix\n",
    "from pytorch_lightning import LightningModule, Trainer\n",
    "from sklearn.metrics import classification_report\n",
    "from torchsig.datasets.sig53 import Sig53\n",
    "from torch.utils.data import DataLoader\n",
    "from matplotlib import pyplot as plt\n",
    "from torch import optim\n",
    "from tqdm import tqdm\n",
    "import torch.nn.functional as F\n",
    "import torchsig.transforms as ST\n",
    "import numpy as np\n",
    "import torchsig\n",
    "import torch\n",
    "import os\n",
    "from sigmf_db_dataset import SigMFDB\n",
    "from sigmf_pytorch_dataset import SigMFDataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchsig.transforms import (\n",
    "    Compose,\n",
    "    IQImbalance,\n",
    "    Normalize,\n",
    "    RandomApply,\n",
    "    RandomFrequencyShift,\n",
    "    RandomPhaseShift,\n",
    "    RandomResample,\n",
    "    RandomTimeShift,\n",
    "    RayleighFadingChannel,\n",
    "    TargetSNR,\n",
    "    ComplexTo2D,\n",
    ")\n",
    "\n",
    "eb_no=False\n",
    "level2 = Compose(\n",
    "    [\n",
    "        RandomApply(RandomPhaseShift((-1, 1)), 0.9),\n",
    "        RandomApply(RandomTimeShift((-32, 32)), 0.9),\n",
    "        RandomApply(RandomFrequencyShift((-0.16, 0.16)), 0.7),\n",
    "        RandomApply(\n",
    "            RayleighFadingChannel((0.05, 0.5), power_delay_profile=(1.0, 0.5, 0.1)),\n",
    "            0.5,\n",
    "        ),\n",
    "        RandomApply(\n",
    "            IQImbalance(\n",
    "                (-3, 3),\n",
    "                (-np.pi * 1.0 / 180.0, np.pi * 1.0 / 180.0),\n",
    "                (-0.1, 0.1),\n",
    "            ),\n",
    "            0.9,\n",
    "        ),\n",
    "        RandomApply(\n",
    "            RandomResample((0.75, 1.5), num_iq_samples=num_iq_samples),\n",
    "            0.5,\n",
    "        ),\n",
    "        #TargetSNR((-2, 30), eb_no=eb_no),\n",
    "        Normalize(norm=np.inf),\n",
    "        ComplexTo2D(),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the SigMF File dataset\n",
    "and generate the class list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# transform = ST.Compose([\n",
    "#     # ST.RandomPhaseShift(phase_offset=(-1, 1)),\n",
    "#     ST.Normalize(norm=np.inf),\n",
    "#     ST.ComplexTo2D(),\n",
    "# ])\n",
    "transform = level2\n",
    "# class_list = ['mini2_video']\n",
    "\n",
    "dataset = SigMFDataset( root=dataset_path,\n",
    "                       sample_count=num_iq_samples,\n",
    "                       transform=transform,\n",
    "                       # class_list=class_list,\n",
    ")\n",
    "# labels = []\n",
    "# for _, label in dataset:\n",
    "#     labels.append(label)\n",
    "# #class_list=np.unique(labels).tolist()\n",
    "# print(\"Unique labels: {}\".format(np.unique(labels)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SigMF DB Data Loaders\n",
    "The number of workers needs to be set to 0 when there is not a GPU. Otherwise the dataloader tries to pickle the environment variables: https://github.com/acids-ircam/RAVE/issues/10#issuecomment-1002708207"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# # Specify Sig53 Options\n",
    "\n",
    "# root = \"sigmf/\"\n",
    "# train = False\n",
    "# impaired = False\n",
    "# class_list = ['signal', 'mini2_video']\n",
    "# print(\"Class List: {}\".format(class_list))\n",
    "# transform = ST.Compose([\n",
    "#     ST.RandomPhaseShift(phase_offset=(-1, 1)),\n",
    "#     ST.Normalize(norm=np.inf),\n",
    "#     ST.ComplexTo2D(),\n",
    "# ])\n",
    "# target_transform = ST.DescToClassIndex(class_list=class_list)\n",
    "\n",
    "# sig53_clean_train = SigMFDB(\n",
    "#     root=root, \n",
    "#     transform=transform,\n",
    "#     target_transform=target_transform,\n",
    "#     use_signal_data=True,\n",
    "# )\n",
    "\n",
    "# # Retrieve a sample and print out information to verify\n",
    "# print(len(sig53_clean_train))\n",
    "# idx = np.random.randint(len(sig53_clean_train))\n",
    "# data, label = sig53_clean_train[idx]\n",
    "# print(\"Dataset length: {}\".format(len(sig53_clean_train)))\n",
    "# print(\"Data shape: {}\".format(data.shape))\n",
    "# print(\"Data Object: {}\".format(data))\n",
    "# print(\"Label Index: {}\".format(label))\n",
    "\n",
    "\n",
    "# train_dataloader = DataLoader(\n",
    "#     sig53_clean_train ,\n",
    "#     batch_size=16,\n",
    "#     num_workers=0,\n",
    "#     shuffle=True,\n",
    "#     drop_last=True,\n",
    "# )\n",
    "# val_dataloader = DataLoader(\n",
    "#     sig53_clean_train ,\n",
    "#     batch_size=16,\n",
    "#     num_workers=0,\n",
    "#     shuffle=False,\n",
    "#     drop_last=True,\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SigMF File Dataloaders\n",
    "This is working"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, val_data = torch.utils.data.random_split(dataset, [0.8, 0.2])\n",
    "\n",
    "# train_class_counts = {class_name:0 for class_name in dataset.class_list}\n",
    "# for data,label in train_data:\n",
    "#     train_class_counts[dataset.class_list[label]] += 1\n",
    "# print(f\"{train_class_counts=}\")\n",
    "\n",
    "# val_class_counts = {class_name:0 for class_name in dataset.class_list}\n",
    "# for data,label in val_data:\n",
    "#     val_class_counts[dataset.class_list[label]] += 1\n",
    "# print(f\"{val_class_counts=}\")\n",
    "\n",
    "# weight = 1. / np.array(list(train_class_counts.values()))\n",
    "# samples_weight = np.array([weight[t] for d,t in train_data])\n",
    "\n",
    "# samples_weight = torch.from_numpy(samples_weight)\n",
    "# sampler = torch.utils.data.WeightedRandomSampler(samples_weight, len(samples_weight))\n",
    "sampler = dataset.get_weighted_sampler(indices=train_data.indices)\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    dataset=train_data, #sig53_clean_train ,\n",
    "    batch_size=180,\n",
    "    num_workers=16,\n",
    "    sampler=sampler,\n",
    "    # shuffle=True,\n",
    "    drop_last=True,\n",
    ")\n",
    "val_dataloader = DataLoader(\n",
    "    dataset=val_data, #sig53_clean_train ,\n",
    "    batch_size=180,\n",
    "    num_workers=16,\n",
    "    shuffle=False,\n",
    "    drop_last=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_class_counts = dataset.get_class_counts(indices=train_data.indices)\n",
    "train_class_counts = {dataset.class_list[k]:v for k,v in train_class_counts.items()}\n",
    "val_class_counts = dataset.get_class_counts(indices=val_data.indices)\n",
    "val_class_counts = {dataset.class_list[k]:v for k,v in val_class_counts.items()}\n",
    "print(f\"{len(train_data)=}, {train_class_counts=}\")\n",
    "print(f\"{len(val_data)=}, {val_class_counts=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = efficientnet_b4(\n",
    "    pretrained=True,\n",
    "    path=\"efficientnet_b4.pt\",\n",
    ")\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExampleNetwork(LightningModule):\n",
    "    def __init__(self, model, data_loader, val_data_loader):\n",
    "        super(ExampleNetwork, self).__init__()\n",
    "        self.mdl = model\n",
    "        self.data_loader = data_loader\n",
    "        self.val_data_loader = val_data_loader\n",
    "\n",
    "        # Hyperparameters\n",
    "        self.lr = 0.001\n",
    "        self.batch_size = data_loader.batch_size\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mdl(x)\n",
    "\n",
    "    def predict(self, x):\n",
    "        with torch.no_grad():\n",
    "            out = self.forward(x)\n",
    "        return out\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return optim.Adam(self.parameters(), lr=self.lr)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return self.data_loader\n",
    "\n",
    "    def training_step(self, batch, batch_nb):\n",
    "        x, y = batch\n",
    "        # print(x.shape)\n",
    "        # print(y.shape)\n",
    "        y = torch.squeeze(y.to(torch.int64))\n",
    "        loss = F.cross_entropy(self(x.float()), y)\n",
    "        return {\"loss\": loss}\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return self.val_data_loader\n",
    "\n",
    "    def validation_step(self, batch, batch_nb):\n",
    "        x, y = batch\n",
    "        y = torch.squeeze(y.to(torch.int64))\n",
    "        val_loss = F.cross_entropy(self(x.float()), y)\n",
    "        self.log(\"val_loss\", val_loss, prog_bar=True)\n",
    "        return {\"val_loss\": val_loss}\n",
    "\n",
    "example_model = ExampleNetwork(model, train_dataloader, val_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Setup checkpoint callbacks\n",
    "checkpoint_filename = \"{}/checkpoints/checkpoint\".format(os.getcwd())\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    filename=checkpoint_filename,\n",
    "    save_top_k=True,\n",
    "    monitor=\"val_loss\",\n",
    "    mode=\"min\",\n",
    ")\n",
    "print(\"Doing stuff\")\n",
    "# Create and fit trainer\n",
    "epochs = 10\n",
    "trainer = Trainer(\n",
    "    max_epochs=epochs, callbacks=checkpoint_callback, accelerator=\"gpu\", devices=1\n",
    ")\n",
    "trainer.fit(example_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_callback.best_model_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate the Trained Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load best checkpoint\n",
    "#checkpoint_filename = checkpoint_callback.best_model_path\n",
    "checkpoint_filename = \"/home/ltindall/rfml-dev/rfml-dev/checkpoints/checkpoint-v9.ckpt\"\n",
    "checkpoint = torch.load(checkpoint_filename, map_location=lambda storage, loc: storage)\n",
    "example_model.load_state_dict(checkpoint[\"state_dict\"], strict=False)\n",
    "example_model = example_model.eval()\n",
    "example_model = example_model.cuda() if torch.cuda.is_available() else example_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Infer results over validation set\n",
    "num_test_examples = len(val_data)\n",
    "# num_classes = 5 #len(list(Sig53._idx_to_name_dict.values()))\n",
    "# y_raw_preds = np.empty((num_test_examples,num_classes))\n",
    "y_preds = np.zeros((num_test_examples,))\n",
    "y_true = np.zeros((num_test_examples,))\n",
    "y_true_list = []\n",
    "y_preds_list = []\n",
    "with torch.no_grad():\n",
    "    example_model.eval()\n",
    "    #for i in tqdm(range(0,num_test_examples)):\n",
    "    for data, label in tqdm(val_dataloader):\n",
    "        # Retrieve data\n",
    "        # idx = i # Use index if evaluating over full dataset\n",
    "        # data, label = val_data[idx]\n",
    "        # Infer\n",
    "        data = data.float()\n",
    "        #data = torch.from_numpy(data).float()\n",
    "        #data = torch.from_numpy(np.expand_dims(data,0)).float()\n",
    "        data = data.cuda() if torch.cuda.is_available() else data\n",
    "        pred_tmp = example_model.predict(data)\n",
    "        pred_tmp = pred_tmp.cpu().numpy() if torch.cuda.is_available() else pred_tmp\n",
    "\n",
    "\n",
    "        y_preds_list.extend(np.argmax(pred_tmp, axis=1).tolist())\n",
    "        y_true_list.extend(label.tolist())\n",
    "        # # Argmax\n",
    "        # y_preds[i] = np.argmax(pred_tmp)\n",
    "        # # Store label\n",
    "        # y_true[i] = label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_preds = y_preds_list\n",
    "y_true = y_true_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc = np.sum(np.asarray(y_preds)==np.asarray(y_true))/len(y_true)\n",
    "plot_confusion_matrix(\n",
    "    y_true, \n",
    "    y_preds, \n",
    "    classes=dataset.class_list,\n",
    "    normalize=True,\n",
    "    title=\"Example Modulations Confusion Matrix\\nTotal Accuracy: {:.2f}%\".format(acc*100),\n",
    "    text=True,\n",
    "    rotate_x_text=90,\n",
    "    figsize=(16,9),\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "print(f\"{len(train_data)=}, {train_class_counts=}\")\n",
    "print(f\"{len(val_data)=}, {val_class_counts=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.class_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}