Upload folder using huggingface_hub

.gitattributes

@@ -36,3 +36,68 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 accuracy_vs_snr.png filter=lfs diff=lfs merge=lfs -text
 accuracy_vs_all_parameters.png filter=lfs diff=lfs merge=lfs -text
 accuracy_vs_dm.png filter=lfs diff=lfs merge=lfs -text
+models/.ipynb_checkpoints/recover_crab-checkpoint.ipynb filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/.ipynb_checkpoints/hit_100000000_1739230556_9-checkpoint.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/.ipynb_checkpoints/hit_100000000_1739231399_1-checkpoint.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_100000000_1739230556_9.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_100000000_1739231399_1.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_100000000_1739231802_11.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_100000000_1739234628_13.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_100000000_1739234628_14.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_100000000_1739235333_29.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_100000000_1739235841_12.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_50233055_1739232802_29.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_52111435_1739229641_28.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_52550001_1739233595_4.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_57096732_1739234611_11.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_57521253_1739232651_11.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_58032264_1739232672_22.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_58165746_1739230560_10.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_62177701_1739230031_23.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_64237575_1739233604_2.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_67249737_1739231769_23.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_71882680_1739230648_29.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_72677566_1739232113_26.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_74160848_1739234611_5.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_75109552_1739231790_10.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_79640130_1739231950_5.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_81910572_1739231764_29.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_83296520_1739233906_31.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_84171229_1739231886_5.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_84411238_1739233784_3.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_87957837_1739232059_29.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_88699241_1739235027_30.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_90233018_1739235740_0.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_93808281_1739233104_14.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_93821122_1739232374_4.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_94705507_1739231215_6.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_95400645_1739230724_21.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_95544329_1739233784_17.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_96119644_1739232285_1.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_96369222_1739233852_3.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_96470689_1739233843_18.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_96471079_1739232857_0.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_96497133_1739233692_27.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_98139322_1739231500_5.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_98582385_1739232894_12.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_98697207_1739229930_16.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99172221_1739232365_29.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99314646_1739233667_1.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99756914_1739230207_8.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99939211_1739233705_24.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99972041_1739234066_9.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99977277_1739231773_10.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99986237_1739234058_8.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99998032_1739232348_29.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99999287_1739233700_3.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99999351_1739235476_0.png filter=lfs diff=lfs merge=lfs -text
+models/HITS-FEB-10/hit_99999979_1739232399_15.png filter=lfs diff=lfs merge=lfs -text
+models/combined_frb_detections.pdf filter=lfs diff=lfs merge=lfs -text
+models/combined_frb_detections.png filter=lfs diff=lfs merge=lfs -text
+models/hits.png filter=lfs diff=lfs merge=lfs -text
+models/hits_crab.pdf filter=lfs diff=lfs merge=lfs -text
+models/models_mask/accuracy_vs_all_parameters.png filter=lfs diff=lfs merge=lfs -text
+models/models_mask/accuracy_vs_dm.png filter=lfs diff=lfs merge=lfs -text
+models/models_mask/accuracy_vs_snr.png filter=lfs diff=lfs merge=lfs -text
+models/recover_crab.ipynb filter=lfs diff=lfs merge=lfs -text
+models/recover_new_crab-debug.ipynb filter=lfs diff=lfs merge=lfs -text

models/.ipynb_checkpoints/Untitled-checkpoint.ipynb

@@ -0,0 +1,766 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "5577ffee-a5c9-4648-8849-95c2c7ebcebe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from utils_batched_preproc import transform_batched, preproc_flip\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "import torch\n",
+    "from functorch import vmap"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f1180d60-83e7-47ca-aa09-58d26af3c706",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# def renorm_batched(data):\n",
+    "#     mean = torch.mean(data, dim=tuple(range(1, data.ndim)), keepdim=True)\n",
+    "#     std = torch.std(data, dim=tuple(range(1, data.ndim)), keepdim=True)\n",
+    "#     standardized_data = (data - mean) / std\n",
+    "#     return standardized_data\n",
+    "\n",
+    "# def transform_batched(data):\n",
+    "#     copy_data = data.detach().clone()\n",
+    "#     rms = torch.std(data, dim=tuple(range(1, data.ndim)), keepdim=True)  # Batch-wise std\n",
+    "#     mean = torch.mean(data, dim=tuple(range(1, data.ndim)), keepdim=True)  # Batch-wise mean\n",
+    "#     masks_rms = [-1, 5]\n",
+    "    \n",
+    "#     # Prepare the new_data tensor\n",
+    "#     num_masks = len(masks_rms) + 1\n",
+    "#     new_data = torch.zeros((num_masks, *data.shape), device=data.device)  # Shape: (num_masks, batch_size, ..., ...)\n",
+    "\n",
+    "#     # First layer: Apply renorm(log10(copy_data + epsilon))\n",
+    "#     new_data[0] = renorm_batched(torch.log10(copy_data + 1e-10))\n",
+    "#     for i, scale in enumerate(masks_rms, start=1):\n",
+    "#         copy_data = data.detach().clone()\n",
+    "        \n",
+    "#         # Apply masking based on the scale\n",
+    "#         if scale < 0:\n",
+    "#             ind = copy_data < abs(scale) * rms + mean\n",
+    "#         else:\n",
+    "#             ind = copy_data > scale * rms + mean\n",
+    "#         copy_data[ind] = 0\n",
+    "        \n",
+    "#         # Renormalize and log10 transform\n",
+    "#         new_data[i] = renorm_batched(torch.log10(copy_data + 1e-10))\n",
+    "    \n",
+    "#     # Convert to float32\n",
+    "#     new_data = new_data.type(torch.float32)\n",
+    "\n",
+    "#     # Chunk along the last dimension and stack\n",
+    "#     slices = torch.chunk(new_data, 8, dim=-1)  # Adjust for batch-wise slicing\n",
+    "#     new_data = torch.stack(slices, dim=2)  # Insert a new axis at dim=1\n",
+    "#     new_data = torch.swapaxes(new_data, 0,1)\n",
+    "#     # Reshape into final format\n",
+    "#     new_data = new_data.reshape( new_data.size(0), 24,  new_data.size(3), new_data.size(4))  # Flatten batch and mask dimensions\n",
+    "#     return new_data\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "81cc81a8-ecef-43ef-a5cd-c35765384812",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "num params encoder  50840\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_19147/1680389579.py:7: 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",
+      "  model.load_state_dict(torch.load(model_path))\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "DataParallel(\n",
+       "  (module): ResNet(\n",
+       "    (relu): ReLU()\n",
+       "    (conv1): Sequential(\n",
+       "      (0): Conv2d(24, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3))\n",
+       "      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (2): ReLU()\n",
+       "    )\n",
+       "    (maxpool): MaxPool2d(kernel_size=(2, 2), stride=2, padding=1, dilation=1, ceil_mode=False)\n",
+       "    (layer0): Sequential(\n",
+       "      (0): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (1): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (2): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (layer1): Sequential(\n",
+       "      (0): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (downsample): Sequential(\n",
+       "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (1): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (2): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (3): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (layer2): Sequential(\n",
+       "      (0): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (downsample): Sequential(\n",
+       "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (1): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (2): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (3): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (4): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (5): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (layer3): Sequential(\n",
+       "      (0): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (downsample): Sequential(\n",
+       "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n",
+       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (1): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (2): ResidualBlock(\n",
+       "        (conv1): Sequential(\n",
+       "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "          (2): ReLU()\n",
+       "        )\n",
+       "        (conv2): Sequential(\n",
+       "          (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "        )\n",
+       "        (relu): ReLU()\n",
+       "        (dropout1): Dropout(p=0.5, inplace=False)\n",
+       "        (batchnorm_mod): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (avgpool): AvgPool2d(kernel_size=7, stride=1, padding=0)\n",
+       "    (fc): Linear(in_features=39424, out_features=2, bias=True)\n",
+       "    (dropout1): Dropout(p=0.3, inplace=False)\n",
+       "    (encoder): Sequential(\n",
+       "      (0): Conv2d(24, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "      (1): ReLU(inplace=True)\n",
+       "      (2): Dropout(p=0.3, inplace=False)\n",
+       "      (3): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "      (4): ReLU(inplace=True)\n",
+       "      (5): Dropout(p=0.3, inplace=False)\n",
+       "      (6): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "      (7): ReLU(inplace=True)\n",
+       "      (8): Dropout(p=0.3, inplace=False)\n",
+       "      (9): Conv2d(32, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
+       "      (10): Sigmoid()\n",
+       "    )\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_path = 'models/model-47-99.125.pt'\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=2).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "model.load_state_dict(torch.load(model_path))\n",
+    "model.eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "58b8c338-df2f-4ef0-92cf-409c9f034cab",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[ 4.1780, -4.1750],\n",
+      "        [ 4.6414, -4.6303],\n",
+      "        [ 5.0103, -5.0162],\n",
+      "        [ 4.8273, -4.8311],\n",
+      "        [ 4.8523, -4.8661],\n",
+      "        [ 4.8855, -4.9074],\n",
+      "        [ 4.4973, -4.5213],\n",
+      "        [ 5.5996, -5.6192],\n",
+      "        [ 4.7929, -4.8116],\n",
+      "        [ 5.5999, -5.5925],\n",
+      "        [ 4.7918, -4.7998],\n",
+      "        [ 4.0914, -4.0766],\n",
+      "        [ 0.7072, -0.6955],\n",
+      "        [ 4.7136, -4.7234],\n",
+      "        [ 5.3918, -5.4307],\n",
+      "        [ 4.5491, -4.5524],\n",
+      "        [ 4.5412, -4.5391],\n",
+      "        [ 4.6264, -4.6137],\n",
+      "        [ 3.9378, -3.9300],\n",
+      "        [ 5.0673, -5.0792],\n",
+      "        [ 5.7389, -5.7330],\n",
+      "        [ 5.2259, -5.2326],\n",
+      "        [ 5.3856, -5.4036],\n",
+      "        [ 5.0781, -5.1232],\n",
+      "        [ 5.2432, -5.2584],\n",
+      "        [ 5.8163, -5.8209],\n",
+      "        [ 4.7730, -4.7823],\n",
+      "        [ 5.1320, -5.1657],\n",
+      "        [ 5.6486, -5.6485],\n",
+      "        [ 3.7626, -3.7674],\n",
+      "        [ 4.1834, -4.1797],\n",
+      "        [ 4.4452, -4.4566]], device='cuda:0', grad_fn=<GatherBackward>)\n"
+     ]
+    }
+   ],
+   "source": [
+    "test_in = abs(torch.randn(32, 192, 2048).to(device))\n",
+    "results = []\n",
+    "for i in range(32):\n",
+    "    results.append(transform(test_in[i,:,:]))\n",
+    "intermediate = torch.stack(results).cuda()\n",
+    "out = model(intermediate)\n",
+    "test_in.cpu().detach().numpy().tofile(\"input.bin\")\n",
+    "intermediate.cpu().detach().numpy().tofile(\"intermediate.bin\")\n",
+    "out.cpu().detach().numpy().tofile(\"output.bin\")\n",
+    "print(out)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "ad56299a-44e4-4d6b-afcc-18a5f4cf0138",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'preproc_flip' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[13], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m preproc_model \u001b[38;5;241m=\u001b[39m preproc_flip()\n\u001b[1;32m      2\u001b[0m Convert_ONNX(preproc_model,\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmodels_mask/preproc_flip.onnx\u001b[39m\u001b[38;5;124m'\u001b[39m, input_data_mock\u001b[38;5;241m=\u001b[39mtest_in\u001b[38;5;241m.\u001b[39mto(device))\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'preproc_flip' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "preproc_model = preproc_flip()\n",
+    "Convert_ONNX(preproc_model,f'models_mask/preproc_flip.onnx', input_data_mock=test_in.to(device))\n",
+    "# Convert_ONNX(model.module,f'models_mask/model_test.onnx', input_data_mock=intermediate.to(device))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "30e84a9b-0d4f-4cb2-a92b-2e3f0b2ccb20",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([32, 192, 2048])"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "test_in.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "1bb26727-7914-470e-bb48-43d7ee81cb50",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "        ...,\n",
+       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "        [0., 0., 0.,  ..., 0., 0., 0.]], device='cuda:0')"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "torch.flip(test_in[0,:,:], dims = (0,)) - torch.flipud(test_in[0,:,:])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "aeaaab90-6a2a-4851-a1ca-28c54a446573",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor(float)\n",
+      "torch.float32\n",
+      "Input Name: modelInput\n",
+      "Output Name: modelOutput\n",
+      "[array([[ 4.3262615, -4.3409047],\n",
+      "       [ 4.9648395, -4.968621 ],\n",
+      "       [ 5.5126643, -5.522872 ],\n",
+      "       [ 4.7735534, -4.8004475],\n",
+      "       [ 4.0924144, -4.112945 ],\n",
+      "       [ 4.588802 , -4.6043544],\n",
+      "       [ 4.6231914, -4.617625 ],\n",
+      "       [ 5.229881 , -5.2555394],\n",
+      "       [ 4.877381 , -4.882144 ],\n",
+      "       [ 5.2514744, -5.2786503],\n",
+      "       [ 4.2948875, -4.3169603],\n",
+      "       [ 4.5997186, -4.6177607],\n",
+      "       [ 4.9509926, -4.9685597],\n",
+      "       [ 4.933158 , -4.9568825],\n",
+      "       [ 4.747336 , -4.7639017],\n",
+      "       [ 5.020595 , -5.0202913],\n",
+      "       [ 4.914437 , -4.9206715],\n",
+      "       [ 5.193108 , -5.1925435],\n",
+      "       [ 4.5233765, -4.512763 ],\n",
+      "       [ 4.7573333, -4.762632 ],\n",
+      "       [ 5.268702 , -5.2838397],\n",
+      "       [ 4.857734 , -4.8605857],\n",
+      "       [ 5.1886744, -5.2047734],\n",
+      "       [ 5.512568 , -5.5503583],\n",
+      "       [ 5.320961 , -5.344709 ],\n",
+      "       [ 4.1023226, -4.1073256],\n",
+      "       [ 5.17857  , -5.185736 ],\n",
+      "       [ 4.997028 , -4.9933476],\n",
+      "       [ 4.771303 , -4.767269 ],\n",
+      "       [ 5.312805 , -5.3265243],\n",
+      "       [ 5.0030336, -5.0492   ],\n",
+      "       [ 5.429731 , -5.4249325]], dtype=float32)]\n"
+     ]
+    }
+   ],
+   "source": [
+    "import onnxruntime as ort\n",
+    "import onnx\n",
+    "\n",
+    "# Path to your ONNX model\n",
+    "model_path = \"models/model-47-99.125.onnx\"\n",
+    "\n",
+    "# Load the ONNX model\n",
+    "session = ort.InferenceSession(model_path)\n",
+    "\n",
+    "# Get input and output details\n",
+    "input_name = session.get_inputs()[0].name\n",
+    "output_name = session.get_outputs()[0].name\n",
+    "\n",
+    "print(session.get_inputs()[0].type)\n",
+    "print(test_in.dtype)\n",
+    "\n",
+    "print(f\"Input Name: {input_name}\")\n",
+    "print(f\"Output Name: {output_name}\")\n",
+    "\n",
+    "# Example Input Data (Replace with your actual input data)\n",
+    "import numpy as np\n",
+    "\n",
+    "# Perform inference\n",
+    "outputs = session.run([output_name], {input_name: intermediate.cpu().numpy()})\n",
+    "print(outputs)\n",
+    "\n",
+    "onnx_model = onnx.load(model_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "f250739d-4c8a-4752-964a-d0b929c396f4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import onnxruntime as ort\n",
+    "# import onnx\n",
+    "\n",
+    "# # Path to your ONNX model\n",
+    "# model_path = \"models_mask/preproc_test.onnx\"\n",
+    "\n",
+    "# # Load the ONNX model\n",
+    "# session = ort.InferenceSession(model_path)\n",
+    "\n",
+    "# # Get input and output details\n",
+    "# input_name = session.get_inputs()[0].name\n",
+    "# output_name = session.get_outputs()[0].name\n",
+    "\n",
+    "# print(session.get_inputs()[0].type)\n",
+    "# print(test_in.dtype)\n",
+    "\n",
+    "# print(f\"Input Name: {input_name}\")\n",
+    "# print(f\"Output Name: {output_name}\")\n",
+    "\n",
+    "# # Example Input Data (Replace with your actual input data)\n",
+    "# import numpy as np\n",
+    "\n",
+    "# # Perform inference\n",
+    "# outputs = session.run([output_name], {input_name: test_in.cpu().numpy()})\n",
+    "# print(\"Model Output:\", outputs)\n",
+    "\n",
+    "# onnx_model = onnx.load(model_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "24fed4e7-4838-44cc-9c3a-0862bdbe173a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAicAAAGdCAYAAADJ6dNTAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8fJSN1AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAey0lEQVR4nO3df0xd9f3H8deFymXMcpURL8WCzG124o/L5Jd0dpblboQ6snZZxvaHItm6ZcFFc6NL+w+4rJMsMUiynIXtmyDZr8gaIy5zqanXH/gDQwvFVfEXjhkWvZc26r3luoBezvePxau0UHvhlvs5nOcjuX/ccw/nvDkhl2fuPedej23btgAAAAyRk+0BAAAAPok4AQAARiFOAACAUYgTAABgFOIEAAAYhTgBAABGIU4AAIBRiBMAAGCUTdkeIF2Li4t66623tHnzZnk8nmyPAwAAzoFt2zp16pRKS0uVk3P210YcFydvvfWWysrKsj0GAABYhZmZGW3duvWs6zgmTizLkmVZ+vDDDyX975crLCzM8lQAAOBcxONxlZWVafPmzZ+6rsdp360Tj8fl8/kUi8WIEwAAHCKd/9+cEAsAAIxCnAAAAKMQJwAAwCjECQAAMApxAgAAjEKcAAAAoxAnAADAKMQJAAAwCnECAACMQpwAAACjECcAAMAoxAkAADBKVuJkenpajY2Nqqys1DXXXKNEIpGNMQAAgIE2ZWOnt956qw4cOKAdO3bonXfekdfrzcYYy7vbd9r9WHbmAADApdY9Tl566SVdcMEF2rFjhySpqKhovUcAAAAGS/ttneHhYbW0tKi0tFQej0dDQ0NnrGNZlioqKpSfn6/6+nqNjo6mHnv99dd14YUXqqWlRdddd53uueeeNf0CAABgY0k7ThKJhAKBgCzLWvbxwcFBhUIhdXV1aXx8XIFAQE1NTZqdnZUkffjhh3r66af129/+ViMjIzp8+LAOHz68tt8CAABsGGnHSXNzsw4cOKA9e/Ys+3hPT4/27t2r9vZ2VVZWqq+vTwUFBerv75ckXXrppaqpqVFZWZm8Xq927dqliYmJFfc3Pz+veDy+5AYAADaujF6ts7CwoLGxMQWDwY93kJOjYDCokZERSVJtba1mZ2f17rvvanFxUcPDw7ryyitX3GZ3d7d8Pl/qVlZWlsmRAQCAYTIaJydPnlQymZTf71+y3O/3KxKJSJI2bdqke+65R1/72td07bXX6ktf+pK+9a1vrbjN/fv3KxaLpW4zMzOZHBkAABgmK5cSNzc3q7m5+ZzW9Xq9Zl1qDAAAzquMvnJSXFys3NxcRaPRJcuj0ahKSkrWtG3LslRZWana2to1bQcAAJgto3GSl5en6upqhcPh1LLFxUWFw2E1NDSsadsdHR2anJzUkSNH1jomAAAwWNpv68zNzWlqaip1f3p6WhMTEyoqKlJ5eblCoZDa2tpUU1Ojuro69fb2KpFIqL29PaODAwCAjSntODl69KgaGxtT90OhkCSpra1NAwMDam1t1YkTJ9TZ2alIJKKqqiodOnTojJNkAQAAluOxbdvO9hDnwrIsWZalZDKp1157TbFYTIWFhZnfEd+tAwBAxsXjcfl8vnP6/52VbyVeDc45AQDAHRwTJwAAwB0cEydcSgwAgDs4Jk54WwcAAHdwTJwAAAB3IE4AAIBRiBMAAGAUx8QJJ8QCAOAOjokTTogFAMAdHBMnAADAHYgTAABgFOIEAAAYxTFxwgmxAAC4g2PihBNiAQBwB8fECQAAcAfiBAAAGIU4AQAARiFOAACAURwTJ1ytAwCAOzgmTrhaBwAAd3BMnAAAAHcgTgAAgFGIEwAAYBTiBAAAGIU4AQAARiFOAACAUYgTAABgFMfECR/CBgCAOzgmTvgQNgAA3MExcQIAANyBOAEAAEYhTgAAgFGIEwAAYBTiBAAAGIU4AQAARiFOAACAUYgTAABgFOIEAAAYxTFxwsfXAwDgDo6JEz6+HgAAd3BMnAAAAHcgTgAAgFGIEwAAYBTiBAAAGIU4AQAARiFOAACAUYgTAABgFOIEAAAYhTgBAABGIU4AAIBRiBMAAGAU4gQAABhlUzZ2WlFRocLCQuXk5Ojiiy/WE088kY0xAACAgbISJ5L03HPP6cILL8zW7gEAgKF4WwcAABgl7TgZHh5WS0uLSktL5fF4NDQ0dMY6lmWpoqJC+fn5qq+v1+jo6JLHPR6PbrzxRtXW1urPf/7zqocHAAAbT9pxkkgkFAgEZFnWso8PDg4qFAqpq6tL4+PjCgQCampq0uzsbGqdZ555RmNjY/rb3/6me+65R//85z9X/xsAAIANJe04aW5u1oEDB7Rnz55lH+/p6dHevXvV3t6uyspK9fX1qaCgQP39/al1Lr30UknSli1btGvXLo2Pj6+4v/n5ecXj8SU3AACwcWX0nJOFhQWNjY0pGAx+vIOcHAWDQY2MjEj63ysvp06dkiTNzc3p8ccf11VXXbXiNru7u+Xz+VK3srKyTI4MAAAMk9E4OXnypJLJpPx+/5Llfr9fkUhEkhSNRnXDDTcoEAjo+uuv1y233KLa2toVt7l//37FYrHUbWZmJpMjAwAAw6z7pcSXX365XnjhhXNe3+v1yuv1nseJAACASTL6yklxcbFyc3MVjUaXLI9GoyopKVnTti3LUmVl5VlfZQEAAM6X0TjJy8tTdXW1wuFwatni4qLC4bAaGhrWtO2Ojg5NTk7qyJEjax0TAAAYLO23debm5jQ1NZW6Pz09rYmJCRUVFam8vFyhUEhtbW2qqalRXV2dent7lUgk1N7entHBAQDAxpR2nBw9elSNjY2p+6FQSJLU1tamgYEBtba26sSJE+rs7FQkElFVVZUOHTp0xkmy6bIsS5ZlKZlMrmk7AADAbB7btu1sD5GOeDwun8+nWCymwsLCzO/gbt9p92OZ3wcAAC6Tzv9vvlsHAAAYhTgBAABGcUyccCkxAADu4Jg44VJiAADcwTFxAgAA3IE4AQAARnFMnHDOCQAA7uCYOOGcEwAA3MExcQIAANyBOAEAAEYhTgAAgFGIEwAAYBTHxAlX6wAA4A6OiROu1gEAwB0cEycAAMAdiBMAAGAU4gQAABiFOAEAAEZxTJxwtQ4AAO7gmDjhah0AANzBMXECAADcgTgBAABGIU4AAIBRiBMAAGAU4gQAABiFOAEAAEZxTJzwOScAALiDY+KEzzkBAMAdHBMnAADAHYgTAABgFOIEAAAYhTgBAABGIU4AAIBRiBMAAGAU4gQAABiFOAEAAEYhTgAAgFGIEwAAYBTHxAnfrQMAgDs4Jk74bh0AANzBMXECAADcgTgBAABGIU4AAIBRiBMAAGAU4gQAABiFOAEAAEYhTgAAgFGIEwAAYBTiBAAAGIU4AQAARiFOAACAUYgTAABglKzFyfvvv6/LLrtMd955Z7ZGAAAABspanPzqV7/S9ddfn63dAwAAQ2UlTl5//XW98soram5uzsbuAQCAwdKOk+HhYbW0tKi0tFQej0dDQ0NnrGNZlioqKpSfn6/6+nqNjo4uefzOO+9Ud3f3qocGAAAbV9pxkkgkFAgEZFnWso8PDg4qFAqpq6tL4+PjCgQCampq0uzsrCTp4Ycf1hVXXKErrrhibZMDAIANaVO6P9Dc3HzWt2N6enq0d+9etbe3S5L6+vr0yCOPqL+/X/v27dPzzz+vBx54QAcPHtTc3Jw++OADFRYWqrOzc9ntzc/Pa35+PnU/Ho+nOzIAAHCQjJ5zsrCwoLGxMQWDwY93kJOjYDCokZERSVJ3d7dmZmb073//W/fee6/27t27Yph8tL7P50vdysrKMjkyAAAwTEbj5OTJk0omk/L7/UuW+/1+RSKRVW1z//79isViqdvMzEwmRgUAAIZK+22dTLr11ls/dR2v1yuv13v+hwEAAEbI6CsnxcXFys3NVTQaXbI8Go2qpKRkTdu2LEuVlZWqra1d03YAAIDZMhoneXl5qq6uVjgcTi1bXFxUOBxWQ0PDmrbd0dGhyclJHTlyZK1jAgAAg6X9ts7c3JympqZS96enpzUxMaGioiKVl5crFAqpra1NNTU1qqurU29vrxKJROrqHQAAgLNJO06OHj2qxsbG1P1QKCRJamtr08DAgFpbW3XixAl1dnYqEomoqqpKhw4dOuMk2XRZliXLspRMJte0HQAAYDaPbdt2todIRzwel8/nUywWU2FhYeZ3cLfvtPuxzO8DAACXSef/d9a++A8AAGA5xAkAADCKY+KES4kBAHAHx8QJlxIDAOAOjokTAADgDsQJAAAwimPihHNOAABwB8fECeecAADgDo6JEwAA4A7ECQAAMApxAgAAjOKYOOGEWAAA3MExccIJsQAAuINj4gQAALgDcQIAAIxCnAAAAKMQJwAAwCjECQAAMIpj4oRLiQEAcAfHxAmXEgMA4A6OiRMAAOAOxAkAADAKcQIAAIxCnAAAAKMQJwAAwCjECQAAMIpj4oTPOQEAwB0cEyd8zgkAAO7gmDgBAADuQJwAAACjECcAAMAoxAkAADAKcQIAAIxCnAAAAKMQJwAAwCjECQAAMApxAgAAjOKYOOHj6wEAcAfHxAkfXw8AgDs4Jk4AAIA7ECcAAMAoxAkAADAKcQIAAIxCnAAAAKMQJwAAwCjECQAAMApxAgAAjEKcAAAAoxAnAADAKMQJAAAwCnECAACMsu5x8t5776mmpkZVVVW6+uqr9X//93/rPQIAADDYpvXe4ebNmzU8PKyCggIlEgldffXV+s53vqPPfe5z6z0KAAAw0Lq/cpKbm6uCggJJ0vz8vGzblm3b6z0GAAAwVNpxMjw8rJaWFpWWlsrj8WhoaOiMdSzLUkVFhfLz81VfX6/R0dElj7/33nsKBALaunWr7rrrLhUXF6/6FwAAABtL2nGSSCQUCARkWdayjw8ODioUCqmrq0vj4+MKBAJqamrS7Oxsap2LLrpIL7zwgqanp/WXv/xF0Wh09b8BAADYUNKOk+bmZh04cEB79uxZ9vGenh7t3btX7e3tqqysVF9fnwoKCtTf33/Gun6/X4FAQE8//fSK+5ufn1c8Hl9yAwAAG1dGzzlZWFjQ2NiYgsHgxzvIyVEwGNTIyIgkKRqN6tSpU5KkWCym4eFhbdu2bcVtdnd3y+fzpW5lZWWZHBkAABgmo3Fy8uRJJZNJ+f3+Jcv9fr8ikYgk6c0339SOHTsUCAS0Y8cO/exnP9M111yz4jb379+vWCyWus3MzGRyZAAAYJh1v5S4rq5OExMT57y+1+uV1+s9fwMBAACjZPSVk+LiYuXm5p5xgms0GlVJScmatm1ZliorK1VbW7um7QAAALNlNE7y8vJUXV2tcDicWra4uKhwOKyGhoY1bbujo0OTk5M6cuTIWscEAAAGS/ttnbm5OU1NTaXuT09Pa2JiQkVFRSovL1coFFJbW5tqampUV1en3t5eJRIJtbe3Z3RwAACwMaUdJ0ePHlVjY2PqfigUkiS1tbVpYGBAra2tOnHihDo7OxWJRFRVVaVDhw6dcZIsAADAcjy2Qz473rIsWZalZDKp1157TbFYTIWFhZnf0d2+0+7HMr8PAABcJh6Py+fzndP/73X/bp3V4pwTAADcwTFxAgAA3MExccKlxAAAuINj4oS3dQAAcAfHxAkAAHAH4gQAABiFOAEAAEZxTJxwQiwAAO7gmDjhhFgAANzBMXECAADcgTgBAABGIU4AAIBRHBMnnBALAIA7OCZOOCEWAAB3cEycAAAAdyBOAACAUYgTAABgFOIEAAAYxTFxwtU6AAC4g2PihKt1AABwB8fECQAAcAfiBAAAGIU4AQAARiFOAACAUYgTAABgFOIEAAAYxTFxwuecAADgDo6JEz7nBAAAd3BMnAAAAHcgTgAAgFGIEwAAYBTiBAAAGIU4AQAARiFOAACAUYgTAABgFOIEAAAYhTgBAABGIU4AAIBRHBMnfLcOAADu4Jg44bt1AABwB8fECQAAcAfiBAAAGIU4AQAARiFOAACAUYgTAABgFOIEAAAYhTgBAABGIU4AAIBRiBMAAGAU4gQAABiFOAEAAEYhTgAAgFHWPU5mZma0c+dOVVZW6tprr9XBgwfXewQAAGCwTeu+w02b1Nvbq6qqKkUiEVVXV2vXrl367Gc/u96jAAAAA617nGzZskVbtmyRJJWUlKi4uFjvvPMOcQIAACSt4m2d4eFhtbS0qLS0VB6PR0NDQ2esY1mWKioqlJ+fr/r6eo2Oji67rbGxMSWTSZWVlaU9OAAA2JjSjpNEIqFAICDLspZ9fHBwUKFQSF1dXRofH1cgEFBTU5NmZ2eXrPfOO+/olltu0e9///vVTQ4AADaktN/WaW5uVnNz84qP9/T0aO/evWpvb5ck9fX16ZFHHlF/f7/27dsnSZqfn9fu3bu1b98+bd++/az7m5+f1/z8fOp+PB5Pd2QAAOAgGb1aZ2FhQWNjYwoGgx/vICdHwWBQIyMjkiTbtnXrrbfq61//um6++eZP3WZ3d7d8Pl/qxltAAABsbBmNk5MnTyqZTMrv9y9Z7vf7FYlEJEnPPvusBgcHNTQ0pKqqKlVVVen48eMrbnP//v2KxWKp28zMTCZHBgAAhln3q3VuuOEGLS4unvP6Xq9XXq/3PE4EAABMktFXToqLi5Wbm6toNLpkeTQaVUlJyZq2bVmWKisrVVtbu6btAAAAs2U0TvLy8lRdXa1wOJxatri4qHA4rIaGhjVtu6OjQ5OTkzpy5MhaxwQAAAZL+22dubk5TU1Npe5PT09rYmJCRUVFKi8vVygUUltbm2pqalRXV6fe3l4lEonU1TsAAABnk3acHD16VI2Njan7oVBIktTW1qaBgQG1trbqxIkT6uzsVCQSUVVVlQ4dOnTGSbLpsixLlmUpmUyuaTsAAMBsHtu27WwPkY54PC6fz6dYLKbCwsLM7+Bu32n3Y5nfBwAALpPO/+91/1ZiAACAsyFOAACAURwTJ1xKDACAOzgmTriUGAAAd3BMnAAAAHcgTgAAgFEcEyeccwIAgDs4Jk445wQAAHdwTJwAAAB3IE4AAIBRiBMAAGAU4gQAABjFMXHC1ToAALiDY+KEq3UAAHAHx8QJAABwB+IEAAAYhTgBAABGIU4AAIBRHBMnXK0DAIA7OCZOuFoHAAB3cEycAAAAdyBOAACAUYgTAABgFOIEAAAYhTgBAABGIU4AAIBRHBMnfM4JAADu4Jg44XNOAABwB8fECQAAcAfiBAAAGIU4AQAARiFOAACAUYgTAABgFOIEAAAYhTgBAABGIU4AAIBRiBMAAGAU4gQAABjFMXHCd+sAAOAOjokTvlsHAAB3cEycAAAAdyBOAACAUYgTAABgFOIEAAAYhTgBAABGIU4AAIBRiBMAAGAU4gQAABiFOAEAAEYhTgAAgFGIEwAAYBTiBAAAGCUrcbJnzx5dfPHF+u53v5uN3QMAAINlJU5uv/12/eEPf8jGrgEAgOGyEic7d+7U5s2bs7FrAABguLTjZHh4WC0tLSotLZXH49HQ0NAZ61iWpYqKCuXn56u+vl6jo6OZmBUAALhA2nGSSCQUCARkWdayjw8ODioUCqmrq0vj4+MKBAJqamrS7Ozsqgacn59XPB5fcgMAABtX2nHS3NysAwcOaM+ePcs+3tPTo71796q9vV2VlZXq6+tTQUGB+vv7VzVgd3e3fD5f6lZWVraq7QAAAGfI6DknCwsLGhsbUzAY/HgHOTkKBoMaGRlZ1Tb379+vWCyWus3MzGRqXAAAYKBNmdzYyZMnlUwm5ff7lyz3+/165ZVXUveDwaBeeOEFJRIJbd26VQcPHlRDQ8Oy2/R6vfJ6vZkcEwAAGCyjcXKuHnvssbR/xrIsWZalZDJ5HiYCAMCl7vYtsyy2/nN8Qkbf1ikuLlZubq6i0eiS5dFoVCUlJWvadkdHhyYnJ3XkyJE1bQcAAJgto3GSl5en6upqhcPh1LLFxUWFw+EV37YBAAD4pLTf1pmbm9PU1FTq/vT0tCYmJlRUVKTy8nKFQiG1tbWppqZGdXV16u3tVSKRUHt7e0YHBwAAG1PacXL06FE1Njam7odCIUlSW1ubBgYG1NraqhMnTqizs1ORSERVVVU6dOjQGSfJpotzTgAAcAePbdt2todIRzwel8/nUywWU2FhYeZ3cPqJQVk+KQgAgPNqnU6ITef/d1a+WwcAAGAlxAkAADCKY+LEsixVVlaqtrY226MAAIDzyDFxwuecAADgDo6JEwAA4A7ECQAAMIpj4oRzTgAAcAfHxAnnnAAA4A6OiRMAAOAOxAkAADBK2t+tk20ffdp+PB4/PzuYP+3T/M/XfgAAMMHp//ek8/K/76P/2+fyrTmO+W6dj774b2FhQW+88Ua2xwEAAKswMzOjrVu3nnUdx8TJRxYXF/XWW29p8+bN8ng8Gd12PB5XWVmZZmZmzs+XCroIxzJzOJaZw7HMHI5l5rjlWNq2rVOnTqm0tFQ5OWc/q8Rxb+vk5OR8anGtVWFh4Yb+A1lPHMvM4VhmDscycziWmeOGY+nzLfMNyMvghFgAAGAU4gQAABiFOPkEr9errq4ueb3ebI/ieBzLzOFYZg7HMnM4lpnDsTyT406IBQAAGxuvnAAAAKMQJwAAwCjECQAAMApxAgAAjOK6OLEsSxUVFcrPz1d9fb1GR0fPuv7Bgwf15S9/Wfn5+brmmmv0j3/8Y50mNV86x3JgYEAej2fJLT8/fx2nNdPw8LBaWlpUWloqj8ejoaGhT/2ZJ598Utddd528Xq+++MUvamBg4LzP6QTpHssnn3zyjL9Jj8ejSCSyPgMbrLu7W7W1tdq8ebMuueQS7d69W6+++uqn/hzPl2dazbHk+dJlcTI4OKhQKKSuri6Nj48rEAioqalJs7Ozy67/3HPP6Qc/+IF++MMf6tixY9q9e7d2796tF198cZ0nN0+6x1L636cfvv3226nbm2++uY4TmymRSCgQCMiyrHNaf3p6WjfddJMaGxs1MTGhO+64Qz/60Y/06KOPnudJzZfusfzIq6++uuTv8pJLLjlPEzrHU089pY6ODj3//PM6fPiwPvjgA33zm99UIpFY8Wd4vlzeao6lxPOlbBepq6uzOzo6UveTyaRdWlpqd3d3L7v+9773Pfumm25asqy+vt7+yU9+cl7ndIJ0j+X9999v+3y+dZrOmSTZDz300FnX+fnPf25fddVVS5a1trbaTU1N53Ey5zmXY/nEE0/Ykux33313XWZystnZWVuS/dRTT624Ds+X5+ZcjiXPl7btmldOFhYWNDY2pmAwmFqWk5OjYDCokZGRZX9mZGRkyfqS1NTUtOL6brGaYylJc3Nzuuyyy1RWVqZvf/vbeumll9Zj3A2Fv8nMq6qq0pYtW/SNb3xDzz77bLbHMVIsFpMkFRUVrbgOf5vn5lyOpcTzpWvi5OTJk0omk/L7/UuW+/3+Fd9jjkQiaa3vFqs5ltu2bVN/f78efvhh/elPf9Li4qK2b9+u//znP+sx8oax0t9kPB7Xf//73yxN5UxbtmxRX1+fHnzwQT344IMqKyvTzp07NT4+nu3RjLK4uKg77rhDX/3qV3X11VevuB7Pl5/uXI8lz5cO/FZiOFNDQ4MaGhpS97dv364rr7xSv/vd7/TLX/4yi5PBrbZt26Zt27al7m/fvl1vvPGG7rvvPv3xj3/M4mRm6ejo0Isvvqhnnnkm26M43rkeS54vXfTKSXFxsXJzcxWNRpcsj0ajKikpWfZnSkpK0lrfLVZzLE93wQUX6Ctf+YqmpqbOx4gb1kp/k4WFhfrMZz6Tpak2jrq6Ov4mP+G2227T3//+dz3xxBPaunXrWdfl+fLs0jmWp3Pj86Vr4iQvL0/V1dUKh8OpZYuLiwqHw0sK9ZMaGhqWrC9Jhw8fXnF9t1jNsTxdMpnU8ePHtWXLlvM15obE3+T5NTExwd+kJNu2ddttt+mhhx7S448/rs9//vOf+jP8bS5vNcfydK58vsz2Gbnr6YEHHrC9Xq89MDBgT05O2j/+8Y/tiy66yI5EIrZt2/bNN99s79u3L7X+s88+a2/atMm+99577Zdfftnu6uqyL7jgAvv48ePZ+hWMke6x/MUvfmE/+uij9htvvGGPjY3Z3//+9+38/Hz7pZdeytavYIRTp07Zx44ds48dO2ZLsnt6euxjx47Zb775pm3btr1v3z775ptvTq3/r3/9yy4oKLDvuusu++WXX7Yty7Jzc3PtQ4cOZetXMEa6x/K+++6zh4aG7Ndff90+fvy4ffvtt9s5OTn2Y489lq1fwRg//elPbZ/PZz/55JP222+/nbq9//77qXV4vjw3qzmWPF/atqvixLZt+ze/+Y1dXl5u5+Xl2XV1dfbzzz+feuzGG2+029ralqz/17/+1b7iiivsvLw8+6qrrrIfeeSRdZ7YXOkcyzvuuCO1rt/vt3ft2mWPj49nYWqzfHQ56+m3j45dW1ubfeONN57xM1VVVXZeXp59+eWX2/fff/+6z22idI/lr3/9a/sLX/iCnZ+fbxcVFdk7d+60H3/88ewMb5jljqOkJX9rPF+em9UcS54vbdtj27a9fq/TAAAAnJ1rzjkBAADOQJwAAACjECcAAMAoxAkAADAKcQIAAIxCnAAAAKMQJwAAwCjECQAAMApxAgAAjEKcAAAAoxAnAADAKMQJAAAwyv8D6KAeY7AISbEAAAAASUVORK5CYII=",
+      "text/plain": [
+       "<Figure size 640x480 with 1 Axes>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "%matplotlib inline\n",
+    "plt.hist(abs(intermediate-outputs[0]).ravel(), bins = 100)\n",
+    "plt.yscale('log')\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "71cb219e-b91a-4629-99f6-00db786903c7",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([1.1902e-03, 2.7140e+00, 2.7140e+00, 2.7140e+00, 2.7140e+00, 2.7140e+00,\n",
+       "        2.7140e+00, 2.7140e+00, 2.7140e+00, 2.7140e+00])"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "torch.sort(abs(intermediate-outputs[0]).ravel())[0][-10:]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "92ba5920-5451-4bb1-af0e-5ea987841ab1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import onnxruntime as ort\n",
+    "\n",
+    "session_options = ort.SessionOptions()\n",
+    "session_options.log_severity_level = 0  # Verbose logging\n",
+    "session = ort.InferenceSession(\"models_mask/preproc_test.onnx\", sess_options=session_options)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3277b343-245d-4ac8-a91c-373061dcbf53",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "%matplotlib inline\n",
+    "plt.imshow(outputs[0][0,8,:,:])\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "67e99ef8-e49a-4037-a818-244555b0bdc5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import onnx\n",
+    "\n",
+    "# Path to your ONNX model\n",
+    "model_path = \"models/model-47-99.125.onnx\"\n",
+    "\n",
+    "# Load the ONNX model\n",
+    "onnx_model = onnx.load(model_path)\n",
+    "\n",
+    "# Check the model for validity\n",
+    "onnx.checker.check_model(onnx_model)\n",
+    "\n",
+    "# Print model graph structure (optional)\n",
+    "print(onnx.helper.printable_graph(onnx_model.graph))\n"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/eval_basic-checkpoint.ipynb

@@ -0,0 +1,305 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "98165c88-8ead-4fae-9ea8-6b2e82996fc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "  0%|                                                    | 0/48 [00:22<?, ?it/s]\n"
+     ]
+    },
+    {
+     "ename": "AttributeError",
+     "evalue": "Caught AttributeError in replica 0 on device 0.\nOriginal Traceback (most recent call last):\n  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py\", line 83, in _worker\n    output = module(*input, **kwargs)\n             ^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/modules/module.py\", line 1553, in _wrapped_call_impl\n    return self._call_impl(*args, **kwargs)\n           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/modules/module.py\", line 1562, in _call_impl\n    return forward_call(*args, **kwargs)\n           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/home/pma/projects/frbnn_narrow/CNN/resnet_model.py\", line 106, in forward\n    return x, self.mask, self.value\n              ^^^^^^^^^\n  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/modules/module.py\", line 1729, in __getattr__\n    raise AttributeError(f\"'{type(self).__name__}' object has no attribute '{name}'\")\nAttributeError: 'ResNet' object has no attribute 'mask'\n",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[1], line 50\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m images, labels \u001b[38;5;129;01min\u001b[39;00m tqdm(testloader):\n\u001b[1;32m     49\u001b[0m     inputs, labels \u001b[38;5;241m=\u001b[39m images\u001b[38;5;241m.\u001b[39mto(device), labels\n\u001b[0;32m---> 50\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m model(inputs, return_mask \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m     51\u001b[0m     _, predicted \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mmax(outputs, \u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     52\u001b[0m     results[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124moutput\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mextend(outputs\u001b[38;5;241m.\u001b[39mcpu()\u001b[38;5;241m.\u001b[39mnumpy()\u001b[38;5;241m.\u001b[39mtolist())\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/modules/module.py:1562\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1558\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1559\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1560\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1561\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1565\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py:186\u001b[0m, in \u001b[0;36mDataParallel.forward\u001b[0;34m(self, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m    184\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodule(\u001b[38;5;241m*\u001b[39minputs[\u001b[38;5;241m0\u001b[39m], \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mmodule_kwargs[\u001b[38;5;241m0\u001b[39m])\n\u001b[1;32m    185\u001b[0m replicas \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mreplicate(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodule, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice_ids[:\u001b[38;5;28mlen\u001b[39m(inputs)])\n\u001b[0;32m--> 186\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mparallel_apply(replicas, inputs, module_kwargs)\n\u001b[1;32m    187\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgather(outputs, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_device)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/data_parallel.py:201\u001b[0m, in \u001b[0;36mDataParallel.parallel_apply\u001b[0;34m(self, replicas, inputs, kwargs)\u001b[0m\n\u001b[1;32m    200\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mparallel_apply\u001b[39m(\u001b[38;5;28mself\u001b[39m, replicas: Sequence[T], inputs: Sequence[Any], kwargs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m List[Any]:\n\u001b[0;32m--> 201\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m parallel_apply(replicas, inputs, kwargs, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice_ids[:\u001b[38;5;28mlen\u001b[39m(replicas)])\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:108\u001b[0m, in \u001b[0;36mparallel_apply\u001b[0;34m(modules, inputs, kwargs_tup, devices)\u001b[0m\n\u001b[1;32m    106\u001b[0m     output \u001b[38;5;241m=\u001b[39m results[i]\n\u001b[1;32m    107\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(output, ExceptionWrapper):\n\u001b[0;32m--> 108\u001b[0m         output\u001b[38;5;241m.\u001b[39mreraise()\n\u001b[1;32m    109\u001b[0m     outputs\u001b[38;5;241m.\u001b[39mappend(output)\n\u001b[1;32m    110\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m outputs\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/_utils.py:706\u001b[0m, in \u001b[0;36mExceptionWrapper.reraise\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    702\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m:\n\u001b[1;32m    703\u001b[0m     \u001b[38;5;66;03m# If the exception takes multiple arguments, don't try to\u001b[39;00m\n\u001b[1;32m    704\u001b[0m     \u001b[38;5;66;03m# instantiate since we don't know how to\u001b[39;00m\n\u001b[1;32m    705\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(msg) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 706\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exception\n",
+      "\u001b[0;31mAttributeError\u001b[0m: Caught AttributeError in replica 0 on device 0.\nOriginal Traceback (most recent call last):\n  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py\", line 83, in _worker\n    output = module(*input, **kwargs)\n             ^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/modules/module.py\", line 1553, in _wrapped_call_impl\n    return self._call_impl(*args, **kwargs)\n           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/modules/module.py\", line 1562, in _call_impl\n    return forward_call(*args, **kwargs)\n           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/home/pma/projects/frbnn_narrow/CNN/resnet_model.py\", line 106, in forward\n    return x, self.mask, self.value\n              ^^^^^^^^^\n  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/modules/module.py\", line 1729, in __getattr__\n    raise AttributeError(f\"'{type(self).__name__}' object has no attribute '{name}'\")\nAttributeError: 'ResNet' object has no attribute 'mask'\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "model_1 = 'models/model-23-99.045.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models/test_42.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "64733667-75c3-4fd3-ab9f-62b85c5e27e3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models/model-14-98.005.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models/test_1.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fe74ada8-43e4-4c73-b772-0ef18983345d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models/model-28-98.955.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models/test_7109.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/eval_basic-extend-checkpoint.ipynb

@@ -0,0 +1,485 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "98165c88-8ead-4fae-9ea8-6b2e82996fc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:39<00:00,  1.21it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  98.82\n",
+      "===========================\n",
+      "False Positive Rate: 0.010\n",
+      "Precision: 0.990\n",
+      "Recall: 0.986\n",
+      "F1 Score: 0.988\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "model_1 = 'models/model-23-99.045.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models/test_42.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "from sklearn.metrics import confusion_matrix\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "64733667-75c3-4fd3-ab9f-62b85c5e27e3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:51<00:00,  1.07s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  97.185\n",
+      "===========================\n",
+      "False Positive Rate: 0.038\n",
+      "Precision: 0.963\n",
+      "Recall: 0.981\n",
+      "F1 Score: 0.972\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models/model-14-98.005.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models/test_1.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "from sklearn.metrics import confusion_matrix\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "fe74ada8-43e4-4c73-b772-0ef18983345d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:53<00:00,  1.11s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  98.455\n",
+      "===========================\n",
+      "False Positive Rate: 0.010\n",
+      "Precision: 0.990\n",
+      "Recall: 0.979\n",
+      "F1 Score: 0.984\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models/model-28-98.955.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models/test_7109.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "from sklearn.metrics import confusion_matrix\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "ad4ef08f-6a9b-495f-bb93-0a13251f0adb",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "98.15333333333332 0.7007416705811665\n",
+      "0.9803333333333333 0.0009428090415820641\n",
+      "0.9813333333333333 0.006798692684790386\n",
+      "0.019333333333333334 0.013199326582148887\n"
+     ]
+    }
+   ],
+   "source": [
+    "# acc\n",
+    "print(np.mean([98.82,97.185,98.455]), np.std([98.82,97.185,98.455]))\n",
+    "# recall\n",
+    "print(np.mean([0.981,0.981, 0.979]), np.std([0.981,0.981, 0.979]))\n",
+    "# f1\n",
+    "print(np.mean([0.988,0.972,0.984]),np.std([0.988,0.972,0.984]))\n",
+    "# fp\n",
+    "print(np.mean([0.010,0.038,0.010]),np.std([0.010,0.038,0.010]))"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/eval_mask-8-checkpoint.ipynb

@@ -0,0 +1,372 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "98165c88-8ead-4fae-9ea8-6b2e82996fc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [01:43<00:00,  2.16s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  98.94\n",
+      "===========================\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, bit8 =True, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "model_1 = 'models_8/model-25-99.31_7109.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_8/test_7109.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "64733667-75c3-4fd3-ab9f-62b85c5e27e3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████████████��████████████████████████| 48/48 [00:54<00:00,  1.14s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.17\n",
+      "===========================\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, bit8 =True, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_8/model-44-99.445_42.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_8/test_42.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "fe74ada8-43e4-4c73-b772-0ef18983345d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|███████████████████████████████████████████| 48/48 [00:54<00:00,  1.14s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.035\n",
+      "===========================\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir,  bit8 =True, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_8/model-43-99.355_1.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "        \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_8/test_1.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/eval_mask-8-extend-checkpoint.ipynb

@@ -0,0 +1,483 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "98165c88-8ead-4fae-9ea8-6b2e82996fc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:44<00:00,  1.09it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  98.94\n",
+      "===========================\n",
+      "False Positive Rate: 0.004\n",
+      "Precision: 0.996\n",
+      "Recall: 0.983\n",
+      "F1 Score: 0.989\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, bit8 =True, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "model_1 = 'models_8/model-25-99.31_7109.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_8/test_7109.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "from sklearn.metrics import confusion_matrix\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "64733667-75c3-4fd3-ab9f-62b85c5e27e3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:41<00:00,  1.15it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.17\n",
+      "===========================\n",
+      "False Positive Rate: 0.004\n",
+      "Precision: 0.995\n",
+      "Recall: 0.988\n",
+      "F1 Score: 0.992\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, bit8 =True, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_8/model-44-99.445_42.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_8/test_42.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "fe74ada8-43e4-4c73-b772-0ef18983345d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:53<00:00,  1.12s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.035\n",
+      "===========================\n",
+      "False Positive Rate: 0.010\n",
+      "Precision: 0.990\n",
+      "Recall: 0.990\n",
+      "F1 Score: 0.990\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir,  bit8 =True, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_8/model-43-99.355_1.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "        \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_8/test_1.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "8444ced5-686c-4c6e-bb02-8c4a45ff8af9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "99.04833333333333 0.09436925111261553\n",
+      "0.9936666666666666 0.002624669291337273\n",
+      "0.9903333333333334 0.0012472191289246482\n",
+      "0.006000000000000001 0.00282842712474619\n"
+     ]
+    }
+   ],
+   "source": [
+    "# acc\n",
+    "print(np.mean([98.94,99.17, 99.035 ]), np.std([98.94,99.17, 99.035 ]))\n",
+    "# recall\n",
+    "print(np.mean([0.996,0.988, 0.990]), np.std([0.996,0.995, 0.990]))\n",
+    "# f1\n",
+    "print(np.mean([0.989,0.992,0.990 ]),np.std([0.989,0.992,0.990 ]))\n",
+    "# fp\n",
+    "print(np.mean([0.004,0.004,0.010]),np.std([0.004,0.004,0.010]))\n"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/eval_mask-checkpoint.ipynb

@@ -0,0 +1,323 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "98165c88-8ead-4fae-9ea8-6b2e82996fc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "  8%|███▋                                        | 4/48 [02:22<25:11, 34.35s/it]Exception ignored in: <bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7efcbc3f67d0>>\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/ipykernel/ipkernel.py\", line 775, in _clean_thread_parent_frames\n",
+      "    def _clean_thread_parent_frames(\n",
+      "\n",
+      "KeyboardInterrupt: \n",
+      " 10%|████▌                                       | 5/48 [02:53<24:56, 34.79s/it]\n"
+     ]
+    },
+    {
+     "ename": "RuntimeError",
+     "evalue": "DataLoader worker (pid(s) 4158742, 4158790, 4158838, 4158886, 4158934, 4158982, 4159030, 4159078, 4159126, 4159174, 4159222, 4159270, 4159318, 4159366, 4159414, 4159462, 4159510, 4159558, 4159606, 4159654, 4159702, 4159750, 4159798, 4159846, 4159894, 4159942, 4159990, 4160038, 4160086, 4160134, 4160182, 4160230) exited unexpectedly",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/utils/data/dataloader.py:1131\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m   1130\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1131\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_data_queue\u001b[38;5;241m.\u001b[39mget(timeout\u001b[38;5;241m=\u001b[39mtimeout)\n\u001b[1;32m   1132\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m (\u001b[38;5;28;01mTrue\u001b[39;00m, data)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/multiprocessing/queues.py:122\u001b[0m, in \u001b[0;36mQueue.get\u001b[0;34m(self, block, timeout)\u001b[0m\n\u001b[1;32m    121\u001b[0m \u001b[38;5;66;03m# unserialize the data after having released the lock\u001b[39;00m\n\u001b[0;32m--> 122\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _ForkingPickler\u001b[38;5;241m.\u001b[39mloads(res)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/multiprocessing/reductions.py:496\u001b[0m, in \u001b[0;36mrebuild_storage_fd\u001b[0;34m(cls, df, size)\u001b[0m\n\u001b[1;32m    495\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrebuild_storage_fd\u001b[39m(\u001b[38;5;28mcls\u001b[39m, df, size):\n\u001b[0;32m--> 496\u001b[0m     fd \u001b[38;5;241m=\u001b[39m df\u001b[38;5;241m.\u001b[39mdetach()\n\u001b[1;32m    497\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/multiprocessing/resource_sharer.py:57\u001b[0m, in \u001b[0;36mDupFd.detach\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m'''Get the fd.  This should only be called once.'''\u001b[39;00m\n\u001b[0;32m---> 57\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m _resource_sharer\u001b[38;5;241m.\u001b[39mget_connection(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_id) \u001b[38;5;28;01mas\u001b[39;00m conn:\n\u001b[1;32m     58\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m reduction\u001b[38;5;241m.\u001b[39mrecv_handle(conn)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/multiprocessing/resource_sharer.py:86\u001b[0m, in \u001b[0;36m_ResourceSharer.get_connection\u001b[0;34m(ident)\u001b[0m\n\u001b[1;32m     85\u001b[0m address, key \u001b[38;5;241m=\u001b[39m ident\n\u001b[0;32m---> 86\u001b[0m c \u001b[38;5;241m=\u001b[39m Client(address, authkey\u001b[38;5;241m=\u001b[39mprocess\u001b[38;5;241m.\u001b[39mcurrent_process()\u001b[38;5;241m.\u001b[39mauthkey)\n\u001b[1;32m     87\u001b[0m c\u001b[38;5;241m.\u001b[39msend((key, os\u001b[38;5;241m.\u001b[39mgetpid()))\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/multiprocessing/connection.py:519\u001b[0m, in \u001b[0;36mClient\u001b[0;34m(address, family, authkey)\u001b[0m\n\u001b[1;32m    518\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 519\u001b[0m     c \u001b[38;5;241m=\u001b[39m SocketClient(address)\n\u001b[1;32m    521\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m authkey \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(authkey, \u001b[38;5;28mbytes\u001b[39m):\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/multiprocessing/connection.py:647\u001b[0m, in \u001b[0;36mSocketClient\u001b[0;34m(address)\u001b[0m\n\u001b[1;32m    646\u001b[0m s\u001b[38;5;241m.\u001b[39msetblocking(\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m--> 647\u001b[0m s\u001b[38;5;241m.\u001b[39mconnect(address)\n\u001b[1;32m    648\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m Connection(s\u001b[38;5;241m.\u001b[39mdetach())\n",
+      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory",
+      "\nThe above exception was the direct cause of the following exception:\n",
+      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[1], line 48\u001b[0m\n\u001b[1;32m     46\u001b[0m model\u001b[38;5;241m.\u001b[39meval()\n\u001b[1;32m     47\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[0;32m---> 48\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m images, labels \u001b[38;5;129;01min\u001b[39;00m tqdm(testloader):\n\u001b[1;32m     49\u001b[0m         inputs, labels \u001b[38;5;241m=\u001b[39m images\u001b[38;5;241m.\u001b[39mto(device), labels\n\u001b[1;32m     50\u001b[0m         outputs \u001b[38;5;241m=\u001b[39m model(inputs, return_mask \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/tqdm/std.py:1181\u001b[0m, in \u001b[0;36mtqdm.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1178\u001b[0m time \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_time\n\u001b[1;32m   1180\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1181\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m obj \u001b[38;5;129;01min\u001b[39;00m iterable:\n\u001b[1;32m   1182\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m obj\n\u001b[1;32m   1183\u001b[0m         \u001b[38;5;66;03m# Update and possibly print the progressbar.\u001b[39;00m\n\u001b[1;32m   1184\u001b[0m         \u001b[38;5;66;03m# Note: does not call self.update(1) for speed optimisation.\u001b[39;00m\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/utils/data/dataloader.py:630\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    627\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    628\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[1;32m    629\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[0;32m--> 630\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_next_data()\n\u001b[1;32m    631\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    632\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    633\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    634\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/utils/data/dataloader.py:1327\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1324\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_process_data(data)\n\u001b[1;32m   1326\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_shutdown \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tasks_outstanding \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[0;32m-> 1327\u001b[0m idx, data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_data()\n\u001b[1;32m   1328\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_tasks_outstanding \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m   1329\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable:\n\u001b[1;32m   1330\u001b[0m     \u001b[38;5;66;03m# Check for _IterableDatasetStopIteration\u001b[39;00m\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/utils/data/dataloader.py:1293\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._get_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1289\u001b[0m     \u001b[38;5;66;03m# In this case, `self._data_queue` is a `queue.Queue`,. But we don't\u001b[39;00m\n\u001b[1;32m   1290\u001b[0m     \u001b[38;5;66;03m# need to call `.task_done()` because we don't use `.join()`.\u001b[39;00m\n\u001b[1;32m   1291\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1292\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m-> 1293\u001b[0m         success, data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_try_get_data()\n\u001b[1;32m   1294\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m success:\n\u001b[1;32m   1295\u001b[0m             \u001b[38;5;28;01mreturn\u001b[39;00m data\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/torch/utils/data/dataloader.py:1144\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m   1142\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(failed_workers) \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m   1143\u001b[0m     pids_str \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;28mstr\u001b[39m(w\u001b[38;5;241m.\u001b[39mpid) \u001b[38;5;28;01mfor\u001b[39;00m w \u001b[38;5;129;01min\u001b[39;00m failed_workers)\n\u001b[0;32m-> 1144\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mDataLoader worker (pid(s) \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mpids_str\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m) exited unexpectedly\u001b[39m\u001b[38;5;124m'\u001b[39m) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m   1145\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(e, queue\u001b[38;5;241m.\u001b[39mEmpty):\n\u001b[1;32m   1146\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m (\u001b[38;5;28;01mFalse\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m)\n",
+      "\u001b[0;31mRuntimeError\u001b[0m: DataLoader worker (pid(s) 4158742, 4158790, 4158838, 4158886, 4158934, 4158982, 4159030, 4159078, 4159126, 4159174, 4159222, 4159270, 4159318, 4159366, 4159414, 4159462, 4159510, 4159558, 4159606, 4159654, 4159702, 4159750, 4159798, 4159846, 4159894, 4159942, 4159990, 4160038, 4160086, 4160134, 4160182, 4160230) exited unexpectedly"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "model_1 = 'models_mask/model-43-99.235_42.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_42.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "64733667-75c3-4fd3-ab9f-62b85c5e27e3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_mask/model-36-99.11999999999999_1.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_1.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fe74ada8-43e4-4c73-b772-0ef18983345d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_mask/model-26-99.13_7109.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_7109.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/eval_mask-extend-checkpoint.ipynb

@@ -0,0 +1,500 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "98165c88-8ead-4fae-9ea8-6b2e82996fc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:45<00:00,  1.07it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.125\n",
+      "===========================\n",
+      "Precision: 0.992\n",
+      "Recall: 0.991\n",
+      "F1 Score: 0.991\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "model_1 = 'models_mask/model-43-99.235_42.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_42.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "64733667-75c3-4fd3-ab9f-62b85c5e27e3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:43<00:00,  1.11it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  98.77\n",
+      "===========================\n",
+      "Precision: 0.982\n",
+      "Recall: 0.993\n",
+      "F1 Score: 0.988\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_mask/model-36-99.11999999999999_1.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_1.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "fe74ada8-43e4-4c73-b772-0ef18983345d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                    | 0/48 [00:00<?, ?it/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|███████████████████████████████████████████| 48/48 [00:43<00:00,  1.11it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.03\n",
+      "===========================\n",
+      "Precision: 0.990\n",
+      "Recall: 0.990\n",
+      "F1 Score: 0.990\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_mask/model-26-99.13_7109.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in tqdm(testloader):\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = model(inputs, return_mask = True)\n",
+    "        _, predicted = torch.max(outputs, 1)\n",
+    "        results['output'].extend(outputs.cpu().numpy().tolist())\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_7109.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "974e62d6-5088-4cd8-9721-6702717eadee",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "98.97499999999998 0.1500555452713003\n",
+      "0.9913333333333333 0.0012472191289246482\n",
+      "0.9896666666666666 0.0012472191289246482\n"
+     ]
+    }
+   ],
+   "source": [
+    "# acc\n",
+    "print(np.mean([99.125,98.77, 99.03 ]), np.std([99.125,98.77, 99.03 ]))\n",
+    "# precision\n",
+    "print(np.mean([0.991,0.990, 0.993]), np.std([0.991,0.990, 0.993]))\n",
+    "# f1\n",
+    "print(np.mean([0.990,0.988,0.991 ]),np.std([0.990,0.988,0.991 ]))\n",
+    "# recall"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "3eee97ad-114f-4090-b54a-6ec0cc7150f5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "False Positive Rate: 0.200\n"
+     ]
+    }
+   ],
+   "source": [
+    "from sklearn.metrics import confusion_matrix\n",
+    "\n",
+    "# Ground truth and predictions\n",
+    "true = [1, 0, 1, 1, 0, 1, 0, 0, 1, 0]\n",
+    "pred = [1, 0, 1, 0, 0, 1, 1, 0, 1, 0]\n",
+    "\n"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/eval_mask_threshold-extend-checkpoint.ipynb

@@ -0,0 +1,460 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "98165c88-8ead-4fae-9ea8-6b2e82996fc5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.195\n",
+      "===========================\n",
+      "False Positive Rate: 0.005\n",
+      "Precision: 0.995\n",
+      "Recall: 0.989\n",
+      "F1 Score: 0.992\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "threshold = 0.992\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "model_1 = 'models_mask/model-43-99.235_42.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in testloader:\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = nn.Softmax(dim = 1)(model(inputs))\n",
+    "        selection = outputs[:, 1] > threshold\n",
+    "        predicted = selection.int()\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_42.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "from sklearn.metrics import confusion_matrix\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "64733667-75c3-4fd3-ab9f-62b85c5e27e3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.195\n",
+      "===========================\n",
+      "False Positive Rate: 0.007\n",
+      "Precision: 0.993\n",
+      "Recall: 0.991\n",
+      "F1 Score: 0.992\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_mask/model-36-99.11999999999999_1.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in testloader:\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = nn.Softmax(dim = 1)(model(inputs))\n",
+    "        selection = outputs[:, 1] > threshold\n",
+    "        predicted = selection.int()\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_1.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "fe74ada8-43e4-4c73-b772-0ef18983345d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "===========================\n",
+      "accuracy:  99.035\n",
+      "===========================\n",
+      "False Positive Rate: 0.007\n",
+      "Precision: 0.993\n",
+      "Recall: 0.987\n",
+      "F1 Score: 0.990\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from utils import CustomDataset, TestingDataset, transform\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "test_data_dir = '/mnt/buf1/pma/frbnn/test_ready'\n",
+    "test_dataset = TestingDataset(test_data_dir, transform=transform)\n",
+    "\n",
+    "num_classes = 2\n",
+    "testloader = DataLoader(test_dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n",
+    "\n",
+    "\n",
+    "model_1 = 'models_mask/model-26-99.13_7109.pt'\n",
+    "# model_1 ='models/model-47-99.125.pt'\n",
+    "model.load_state_dict(torch.load(model_1, weights_only=True))\n",
+    "model = model.eval()\n",
+    "\n",
+    "# eval\n",
+    "val_loss = 0.0\n",
+    "correct_valid = 0\n",
+    "total = 0\n",
+    "results = {'output': [],'pred': [], 'true':[], 'freq':[], 'snr':[], 'dm':[], 'boxcar':[]}\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for images, labels in testloader:\n",
+    "        inputs, labels = images.to(device), labels\n",
+    "        outputs = nn.Softmax(dim = 1)(model(inputs))\n",
+    "        selection = outputs[:, 1] > threshold\n",
+    "        predicted = selection.int()\n",
+    "        results['pred'].extend(predicted.cpu().numpy().tolist())\n",
+    "        results['true'].extend(labels[0].cpu().numpy().tolist())\n",
+    "        results['freq'].extend(labels[2].cpu().numpy().tolist())\n",
+    "        results['dm'].extend(labels[1].cpu().numpy().tolist())\n",
+    "        results['snr'].extend(labels[3].cpu().numpy().tolist())\n",
+    "        results['boxcar'].extend(labels[4].cpu().numpy().tolist())\n",
+    "        total += labels[0].size(0)\n",
+    "        correct_valid += (predicted.cpu() == labels[0].cpu()).sum().item()\n",
+    "    \n",
+    "    \n",
+    "# Calculate training accuracy after each epoch\n",
+    "val_accuracy = correct_valid / total * 100.0\n",
+    "print(\"===========================\")\n",
+    "print('accuracy: ',  val_accuracy)\n",
+    "print(\"===========================\")\n",
+    "\n",
+    "import pickle\n",
+    "\n",
+    "# Pickle the dictionary to a file\n",
+    "with open('models_mask/test_7109.pkl', 'wb') as f:\n",
+    "    pickle.dump(results, f)\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "\n",
+    "# Example binary labels\n",
+    "true = results['true']  # ground truth\n",
+    "pred = results['pred']  # predicted\n",
+    "\n",
+    "# Compute metrics\n",
+    "precision = precision_score(true, pred)\n",
+    "recall = recall_score(true, pred)\n",
+    "f1 = f1_score(true, pred)\n",
+    "# Get confusion matrix: TN, FP, FN, TP\n",
+    "tn, fp, fn, tp = confusion_matrix(true, pred).ravel()\n",
+    "\n",
+    "# Compute FPR\n",
+    "fpr = fp / (fp + tn)\n",
+    "\n",
+    "print(f\"False Positive Rate: {fpr:.3f}\")\n",
+    "\n",
+    "print(f\"Precision: {precision:.3f}\")\n",
+    "print(f\"Recall: {recall:.3f}\")\n",
+    "print(f\"F1 Score: {f1:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "974e62d6-5088-4cd8-9721-6702717eadee",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "99.14166666666665 0.07542472332656346\n",
+      "0.9913333333333333 0.0012472191289246482\n",
+      "0.9913333333333334 0.0012472191289246482\n",
+      "0.006333333333333333 0.0009428090415820634\n"
+     ]
+    }
+   ],
+   "source": [
+    "# acc\n",
+    "print(np.mean([99.195,99.195, 99.035 ]), np.std([99.195,99.195, 99.035]))\n",
+    "# recall\n",
+    "print(np.mean([0.991,0.991, 0.987]), np.std([0.991,0.990, 0.993]))\n",
+    "# f1\n",
+    "print(np.mean([0.992,0.992,0.990 ]),np.std([0.990,0.988,0.991 ]))\n",
+    "# fp\n",
+    "print(np.mean([0.005,0.007,0.007 ]),np.std([0.005,0.007,0.007]))\n"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/practice_cnn_train-checkpoint.ipynb

@@ -0,0 +1,326 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "851f001a-3882-42cf-8e45-1bb7c4193d20",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, preproc, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "# Create custom dataset instance\n",
+    "# Create custom dataset instance\n",
+    "data_dir = '/mnt/buf0/pma/frbnn/train_ready'\n",
+    "dataset = CustomDataset(data_dir, transform=transform)\n",
+    "valid_data_dir = '/mnt/buf0/pma/frbnn/valid_ready'\n",
+    "valid_dataset = CustomDataset(valid_data_dir, transform=transform)\n",
+    "\n",
+    "\n",
+    "num_classes = 2\n",
+    "trainloader = DataLoader(dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "676a6ffa-5bed-403d-ba03-627f14b36de2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                | 0/477 [00:00<?, ?batch/s]/home/pma/.conda/envs/frbnn/lib/python3.11/site-packages/torch/nn/parallel/parallel_apply.py:79: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled):\n",
+      "100%|██████████████████████████████████████| 477/477 [08:57<00:00,  1.13s/batch]\n"
+     ]
+    },
+    {
+     "ename": "NameError",
+     "evalue": "name 'validloader' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[2], line 29\u001b[0m\n\u001b[1;32m     27\u001b[0m model\u001b[38;5;241m.\u001b[39meval()\n\u001b[1;32m     28\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[0;32m---> 29\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m images, labels \u001b[38;5;129;01min\u001b[39;00m validloader:\n\u001b[1;32m     30\u001b[0m         inputs, labels \u001b[38;5;241m=\u001b[39m images\u001b[38;5;241m.\u001b[39mto(device), labels\u001b[38;5;241m.\u001b[39mto(device)\u001b[38;5;241m.\u001b[39mfloat()\n\u001b[1;32m     31\u001b[0m         optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'validloader' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "criterion = nn.CrossEntropyLoss(weight = torch.tensor([1,1]).to(device))\n",
+    "optimizer = optim.Adam(model.parameters(), lr=0.0001)\n",
+    "scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10)\n",
+    "\n",
+    "for epoch in range(5):\n",
+    "    running_loss = 0.0\n",
+    "    correct_train = 0\n",
+    "    total_train = 0\n",
+    "    with tqdm(trainloader, unit=\"batch\") as tepoch:\n",
+    "        model.train()\n",
+    "        for i, (images, labels) in enumerate(tepoch):\n",
+    "            inputs, labels = images.to(device), labels.to(device).float()\n",
+    "            optimizer.zero_grad()\n",
+    "            outputs = model(inputs, return_mask=False).to(device)\n",
+    "            new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32).to(device)\n",
+    "            loss = criterion(outputs, new_label)\n",
+    "            loss.backward()\n",
+    "            optimizer.step()\n",
+    "            running_loss += loss.item()\n",
+    "            # Calculate training accuracy\n",
+    "            _, predicted = torch.max(outputs.data, 1)\n",
+    "            total_train += labels.size(0)\n",
+    "            correct_train += (predicted == labels).sum().item()        \n",
+    "        val_loss = 0.0\n",
+    "        correct_valid = 0\n",
+    "        total = 0\n",
+    "        model.eval()\n",
+    "        with torch.no_grad():\n",
+    "            for images, labels in validloader:\n",
+    "                inputs, labels = images.to(device), labels.to(device).float()\n",
+    "                optimizer.zero_grad()\n",
+    "                outputs = model(inputs, return_mask=False)\n",
+    "                new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32)\n",
+    "                loss = criterion(outputs, new_label)\n",
+    "                val_loss += loss.item()\n",
+    "                _, predicted = torch.max(outputs, 1)\n",
+    "                total += labels.size(0)\n",
+    "                correct_valid += (predicted == labels).sum().item()\n",
+    "    scheduler.step(val_loss)\n",
+    "    # Calculate training accuracy after each epoch\n",
+    "    train_accuracy = 100 * correct_train / total_train\n",
+    "    val_accuracy = correct_valid / total * 100.0\n",
+    "\n",
+    "\n",
+    "    print(\"===========================\")\n",
+    "    print('accuracy: ', epoch, train_accuracy,  val_accuracy)\n",
+    "    print('learning rate: ', scheduler.get_last_lr())\n",
+    "    print(\"===========================\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3faa4a11-89fb-4556-ae87-3645a47fa00d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_accuracy = 100 * correct_train / total_train\n",
+    "print('accuracy: ', epoch, train_accuracy)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e586c4d2-a7f4-4f14-81fc-4f84ffac52b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sigpyproc.readers as r\n",
+    "import cv2\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from scipy.special import softmax\n",
+    "%matplotlib inline\n",
+    "path = '/mnt/primary/ata/projects/p051/fil_60565_59210_9756774_J0534+2200_0001/LoB.C0928/fil_60565_59210_9756774_J0534+2200_0001-beam0000.fil'\n",
+    "# path = '/mnt/primary/ata/projects/p051/fil_60564_62428_4679748_J0332+5434_0001/LoB.C0928/fil_60564_62428_4679748_J0332+5434_0001-beam0000.fil'\n",
+    "\n",
+    "# Get some metadata\n",
+    "\n",
+    "# Open the filterbank file\n",
+    "fil = r.FilReader(path)\n",
+    "header = fil.header\n",
+    "print(\"Header:\", header)\n",
+    "n=100\n",
+    "li = [ 7257608,  7324207, 10393163, 10641071, 11130537, 11085081,\n",
+    "       11419145, 11964112, 12329364, 13047181]\n",
+    "for el in li:\n",
+    "    data = torch.tensor(fil.read_block(el-1024, 2048)).cuda()\n",
+    "    print(data.shape)\n",
+    "    out = model(transform(torch.tensor(data).cuda())[None])\n",
+    "    print(softmax(out.detach().cpu().numpy(), axis=1))\n",
+    "    plt.figure(figsize=(10,10))\n",
+    "    plt.imshow(data.cpu().numpy(), aspect = 10)\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "609e5564-f14f-4bd1-b604-68e7e7d42834",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "triggers = []\n",
+    "counter = 0\n",
+    "with torch.no_grad():\n",
+    "    for i in range(2048,10201921, 2048 ):\n",
+    "        data = torch.tensor(fil.read_block(i-1024, 2048)).cuda()\n",
+    "        # Shuffle the tensor using the random indices\n",
+    "        out = model(transform(torch.tensor(data).cuda())[None])\n",
+    "        triggers.append(softmax(out.detach().cpu().numpy(), axis=1))\n",
+    "        counter += 1\n",
+    "        if counter > 1000:\n",
+    "            break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "08ee6dcf-cb30-4490-8624-4e52552fdf39",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(triggers[0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8c56c6f5-5a0b-4854-8a94-066a9baf4cfc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "stack = np.stack(triggers)\n",
+    "positives = stack[:,0,1]\n",
+    "num_pos = np.where(positives > 0.5)[0].shape[0]\n",
+    "print(num_pos)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eb1d1591-8855-4989-bf12-c8a9cdbf2a4d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "\n",
+    "# Path to your pickle file\n",
+    "file_path = \"../dataset_generator/dir.pkl\"\n",
+    "\n",
+    "# Open and load the pickle file\n",
+    "with open(file_path, \"rb\") as file:  # Use \"rb\" mode for reading binary files\n",
+    "    data = pickle.load(file)\n",
+    "\n",
+    "# Print the contents of the file\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "46f61d7e-55fa-44fe-be94-d4ddb3c576f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sigpyproc.readers as r\n",
+    "import cv2\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from scipy.special import softmax\n",
+    "%matplotlib inline\n",
+    "path = data[0]\n",
+    "model.eval()\n",
+    "\n",
+    "fil = r.FilReader(path)\n",
+    "header = fil.header\n",
+    "print(\"Header:\", header)\n",
+    "n=100\n",
+    "\n",
+    "\n",
+    "triggers = []\n",
+    "counter = 0\n",
+    "for i in range(2048,10201921, 2048):\n",
+    "    data = torch.tensor(fil.read_block(i-1024, 2048)).cuda()\n",
+    "    # Shuffle the tensor using the random indices\n",
+    "    out = model(transform(torch.tensor(data).cuda())[None])\n",
+    "    triggers.append(softmax(out.detach().cpu().numpy(), axis=1))\n",
+    "    counter += 1\n",
+    "    if counter > 1000:\n",
+    "        break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "413d402e-2ce3-49fc-bbd4-a3cf1cc92388",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(triggers[0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5c039dee-1b9b-4664-b42a-a79d780f37f1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "stack = np.stack(triggers)\n",
+    "positives = stack[:,0,1]\n",
+    "num_pos = np.where(positives > 0.5)[0].shape[0]\n",
+    "print(num_pos)"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/recover_crab-checkpoint.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b063cb5ad71d38a551a5a4beb9ae21399e5e633af128c2608645398509854239
+size 16982029

models/.ipynb_checkpoints/recover_new_crab-debug-checkpoint.ipynb

@@ -0,0 +1,273 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "851f001a-3882-42cf-8e45-1bb7c4193d20",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "6\n",
+      "num params encoder  50840\n",
+      "num params  21496282\n"
+     ]
+    }
+   ],
+   "source": [
+    "from utils import CustomDataset, transform, Convert_ONNX\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from resnet_model_mask import  ResidualBlock, ResNet\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from tqdm import tqdm \n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
+    "import pickle\n",
+    "\n",
+    "torch.manual_seed(1)\n",
+    "# torch.manual_seed(42)\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "num_gpus = torch.cuda.device_count()\n",
+    "print(num_gpus)\n",
+    "\n",
+    "# Create custom dataset instance\n",
+    "data_dir = '/mnt/buf0/pma/frbnn/train_ready'\n",
+    "dataset = CustomDataset(data_dir, transform=transform)\n",
+    "valid_data_dir = '/mnt/buf0/pma/frbnn/valid_ready'\n",
+    "valid_dataset = CustomDataset(valid_data_dir, transform=transform)\n",
+    "\n",
+    "\n",
+    "num_classes = 2\n",
+    "trainloader = DataLoader(dataset, batch_size=420, shuffle=True, num_workers=32)\n",
+    "\n",
+    "model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
+    "model = nn.DataParallel(model)\n",
+    "model = model.to(device)\n",
+    "params = sum(p.numel() for p in model.parameters())\n",
+    "print(\"num params \",params)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "676a6ffa-5bed-403d-ba03-627f14b36de2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model_path = 'models/model-62-98.78.pt'\n",
+    "# model_path = 'models/model-47-99.125.pt'\n",
+    "model_path = 'models_mask/model-37-99.175_42.pt'\n",
+    "\n",
+    "# model_path = 'models_mask/model-10-97.055_1.pt'\n",
+    "model.load_state_dict(torch.load(model_path, weights_only=True))\n",
+    "model = model.eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "89d108de-7eae-4bbd-837c-8e657082a1e6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Header(filename='/mnt/primary/ata/projects/p031/fil_60692_02772_253151611_crab_0001/LoA.C0736/fil_60692_02772_253151611_crab_0001-beam0000.fil', data_type='raw data', nchans=192, foff=-0.5, fch1=1187.5, nbits=32, tsamp=6.4e-05, tstart=60692.03208333333, nsamples=28125184, nifs=1, coord=<SkyCoord (ICRS): (ra, dec) in deg\n",
+      "    (83.63322, 22.01446)>, azimuth=<Angle 80.54659271 deg>, zenith=<Angle 66.41192055 deg>, telescope='Effelsberg LOFAR', backend='FAKE', source='crab', frame='topocentric', ibeam=0, nbeams=2, dm=0, period=0, accel=0, signed=False, rawdatafile='', stream_info=StreamInfo(entries=[FileInfo(filename='/mnt/primary/ata/projects/p031/fil_60692_02772_253151611_crab_0001/LoA.C0736/fil_60692_02772_253151611_crab_0001-beam0000.fil', hdrlen=338, datalen=21600141312, nsamples=28125184, tstart=60692.03208333333, tsamp=6.4e-05)]))\n"
+     ]
+    }
+   ],
+   "source": [
+    "import sigpyproc.readers as r\n",
+    "import cv2\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "fil = r.FilReader('/mnt/primary/ata/projects/p031/fil_60692_02772_253151611_crab_0001/LoA.C0736/fil_60692_02772_253151611_crab_0001-beam0000.fil')\n",
+    "# fil = r.FilReader('/mnt/primary/ata/projects/p031/fil_60692_02772_253151611_crab_0001/LoB.C0736/fil_60692_02772_253151611_crab_0001-beam0000.fil')\n",
+    "header = fil.header\n",
+    "print(header)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "0b276e6e-d6c8-41da-808d-542ee22133d1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                 | 0/13732 [00:00<?, ?it/s]/tmp/ipykernel_19961/1777549771.py:15: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
+      "  out = model(transform(torch.tensor(data).cuda())[None])\n",
+      "  3%|▉                                      | 351/13732 [00:13<08:27, 26.38it/s]\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[11], line 13\u001b[0m\n\u001b[1;32m     11\u001b[0m counter \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m tqdm(\u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m2048\u001b[39m,header\u001b[38;5;241m.\u001b[39mnsamples, \u001b[38;5;241m2048\u001b[39m)):\n\u001b[0;32m---> 13\u001b[0m     data \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor(fil\u001b[38;5;241m.\u001b[39mread_block(i\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1024\u001b[39m, \u001b[38;5;241m2048\u001b[39m))\u001b[38;5;241m.\u001b[39mcuda()\n\u001b[1;32m     14\u001b[0m     \u001b[38;5;66;03m# Shuffle the tensor using the random indices\u001b[39;00m\n\u001b[1;32m     15\u001b[0m     out \u001b[38;5;241m=\u001b[39m model(transform(torch\u001b[38;5;241m.\u001b[39mtensor(data)\u001b[38;5;241m.\u001b[39mcuda())[\u001b[38;5;28;01mNone\u001b[39;00m])\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Error in callback <function flush_figures at 0x7f6c8689ae80> (for post_execute), with arguments args (),kwargs {}:\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib_inline/backend_inline.py:126\u001b[0m, in \u001b[0;36mflush_figures\u001b[0;34m()\u001b[0m\n\u001b[1;32m    123\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m InlineBackend\u001b[38;5;241m.\u001b[39minstance()\u001b[38;5;241m.\u001b[39mclose_figures:\n\u001b[1;32m    124\u001b[0m     \u001b[38;5;66;03m# ignore the tracking, just draw and close all figures\u001b[39;00m\n\u001b[1;32m    125\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 126\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m show(\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    127\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    128\u001b[0m         \u001b[38;5;66;03m# safely show traceback if in IPython, else raise\u001b[39;00m\n\u001b[1;32m    129\u001b[0m         ip \u001b[38;5;241m=\u001b[39m get_ipython()\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib_inline/backend_inline.py:90\u001b[0m, in \u001b[0;36mshow\u001b[0;34m(close, block)\u001b[0m\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m     89\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m figure_manager \u001b[38;5;129;01min\u001b[39;00m Gcf\u001b[38;5;241m.\u001b[39mget_all_fig_managers():\n\u001b[0;32m---> 90\u001b[0m         display(\n\u001b[1;32m     91\u001b[0m             figure_manager\u001b[38;5;241m.\u001b[39mcanvas\u001b[38;5;241m.\u001b[39mfigure,\n\u001b[1;32m     92\u001b[0m             metadata\u001b[38;5;241m=\u001b[39m_fetch_figure_metadata(figure_manager\u001b[38;5;241m.\u001b[39mcanvas\u001b[38;5;241m.\u001b[39mfigure)\n\u001b[1;32m     93\u001b[0m         )\n\u001b[1;32m     94\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m     95\u001b[0m     show\u001b[38;5;241m.\u001b[39m_to_draw \u001b[38;5;241m=\u001b[39m []\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/IPython/core/display_functions.py:298\u001b[0m, in \u001b[0;36mdisplay\u001b[0;34m(include, exclude, metadata, transient, display_id, raw, clear, *objs, **kwargs)\u001b[0m\n\u001b[1;32m    296\u001b[0m     publish_display_data(data\u001b[38;5;241m=\u001b[39mobj, metadata\u001b[38;5;241m=\u001b[39mmetadata, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    297\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 298\u001b[0m     format_dict, md_dict \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mformat\u001b[39m(obj, include\u001b[38;5;241m=\u001b[39minclude, exclude\u001b[38;5;241m=\u001b[39mexclude)\n\u001b[1;32m    299\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m format_dict:\n\u001b[1;32m    300\u001b[0m         \u001b[38;5;66;03m# nothing to display (e.g. _ipython_display_ took over)\u001b[39;00m\n\u001b[1;32m    301\u001b[0m         \u001b[38;5;28;01mcontinue\u001b[39;00m\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/IPython/core/formatters.py:179\u001b[0m, in \u001b[0;36mDisplayFormatter.format\u001b[0;34m(self, obj, include, exclude)\u001b[0m\n\u001b[1;32m    177\u001b[0m md \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    178\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 179\u001b[0m     data \u001b[38;5;241m=\u001b[39m formatter(obj)\n\u001b[1;32m    180\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    181\u001b[0m     \u001b[38;5;66;03m# FIXME: log the exception\u001b[39;00m\n\u001b[1;32m    182\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/decorator.py:232\u001b[0m, in \u001b[0;36mdecorate.<locals>.fun\u001b[0;34m(*args, **kw)\u001b[0m\n\u001b[1;32m    230\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m kwsyntax:\n\u001b[1;32m    231\u001b[0m     args, kw \u001b[38;5;241m=\u001b[39m fix(args, kw, sig)\n\u001b[0;32m--> 232\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m caller(func, \u001b[38;5;241m*\u001b[39m(extras \u001b[38;5;241m+\u001b[39m args), \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkw)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/IPython/core/formatters.py:223\u001b[0m, in \u001b[0;36mcatch_format_error\u001b[0;34m(method, self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    221\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"show traceback on failed format call\"\"\"\u001b[39;00m\n\u001b[1;32m    222\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 223\u001b[0m     r \u001b[38;5;241m=\u001b[39m method(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    224\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m:\n\u001b[1;32m    225\u001b[0m     \u001b[38;5;66;03m# don't warn on NotImplementedErrors\u001b[39;00m\n\u001b[1;32m    226\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_return(\u001b[38;5;28;01mNone\u001b[39;00m, args[\u001b[38;5;241m0\u001b[39m])\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/IPython/core/formatters.py:340\u001b[0m, in \u001b[0;36mBaseFormatter.__call__\u001b[0;34m(self, obj)\u001b[0m\n\u001b[1;32m    338\u001b[0m     \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[1;32m    339\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 340\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m printer(obj)\n\u001b[1;32m    341\u001b[0m \u001b[38;5;66;03m# Finally look for special method names\u001b[39;00m\n\u001b[1;32m    342\u001b[0m method \u001b[38;5;241m=\u001b[39m get_real_method(obj, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprint_method)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/IPython/core/pylabtools.py:152\u001b[0m, in \u001b[0;36mprint_figure\u001b[0;34m(fig, fmt, bbox_inches, base64, **kwargs)\u001b[0m\n\u001b[1;32m    149\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackend_bases\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FigureCanvasBase\n\u001b[1;32m    150\u001b[0m     FigureCanvasBase(fig)\n\u001b[0;32m--> 152\u001b[0m fig\u001b[38;5;241m.\u001b[39mcanvas\u001b[38;5;241m.\u001b[39mprint_figure(bytes_io, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkw)\n\u001b[1;32m    153\u001b[0m data \u001b[38;5;241m=\u001b[39m bytes_io\u001b[38;5;241m.\u001b[39mgetvalue()\n\u001b[1;32m    154\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m fmt \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124msvg\u001b[39m\u001b[38;5;124m'\u001b[39m:\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/backend_bases.py:2164\u001b[0m, in \u001b[0;36mFigureCanvasBase.print_figure\u001b[0;34m(self, filename, dpi, facecolor, edgecolor, orientation, format, bbox_inches, pad_inches, bbox_extra_artists, backend, **kwargs)\u001b[0m\n\u001b[1;32m   2161\u001b[0m     \u001b[38;5;66;03m# we do this instead of `self.figure.draw_without_rendering`\u001b[39;00m\n\u001b[1;32m   2162\u001b[0m     \u001b[38;5;66;03m# so that we can inject the orientation\u001b[39;00m\n\u001b[1;32m   2163\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mgetattr\u001b[39m(renderer, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_draw_disabled\u001b[39m\u001b[38;5;124m\"\u001b[39m, nullcontext)():\n\u001b[0;32m-> 2164\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfigure\u001b[38;5;241m.\u001b[39mdraw(renderer)\n\u001b[1;32m   2165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m bbox_inches:\n\u001b[1;32m   2166\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m bbox_inches \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtight\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/artist.py:95\u001b[0m, in \u001b[0;36m_finalize_rasterization.<locals>.draw_wrapper\u001b[0;34m(artist, renderer, *args, **kwargs)\u001b[0m\n\u001b[1;32m     93\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(draw)\n\u001b[1;32m     94\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdraw_wrapper\u001b[39m(artist, renderer, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m---> 95\u001b[0m     result \u001b[38;5;241m=\u001b[39m draw(artist, renderer, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m     96\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m renderer\u001b[38;5;241m.\u001b[39m_rasterizing:\n\u001b[1;32m     97\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mstop_rasterizing()\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/artist.py:72\u001b[0m, in \u001b[0;36mallow_rasterization.<locals>.draw_wrapper\u001b[0;34m(artist, renderer)\u001b[0m\n\u001b[1;32m     69\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     70\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mstart_filter()\n\u001b[0;32m---> 72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m draw(artist, renderer)\n\u001b[1;32m     73\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m     74\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/figure.py:3154\u001b[0m, in \u001b[0;36mFigure.draw\u001b[0;34m(self, renderer)\u001b[0m\n\u001b[1;32m   3151\u001b[0m         \u001b[38;5;66;03m# ValueError can occur when resizing a window.\u001b[39;00m\n\u001b[1;32m   3153\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpatch\u001b[38;5;241m.\u001b[39mdraw(renderer)\n\u001b[0;32m-> 3154\u001b[0m mimage\u001b[38;5;241m.\u001b[39m_draw_list_compositing_images(\n\u001b[1;32m   3155\u001b[0m     renderer, \u001b[38;5;28mself\u001b[39m, artists, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msuppressComposite)\n\u001b[1;32m   3157\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m sfig \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msubfigs:\n\u001b[1;32m   3158\u001b[0m     sfig\u001b[38;5;241m.\u001b[39mdraw(renderer)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/image.py:132\u001b[0m, in \u001b[0;36m_draw_list_compositing_images\u001b[0;34m(renderer, parent, artists, suppress_composite)\u001b[0m\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m not_composite \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m has_images:\n\u001b[1;32m    131\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m artists:\n\u001b[0;32m--> 132\u001b[0m         a\u001b[38;5;241m.\u001b[39mdraw(renderer)\n\u001b[1;32m    133\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    134\u001b[0m     \u001b[38;5;66;03m# Composite any adjacent images together\u001b[39;00m\n\u001b[1;32m    135\u001b[0m     image_group \u001b[38;5;241m=\u001b[39m []\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/artist.py:72\u001b[0m, in \u001b[0;36mallow_rasterization.<locals>.draw_wrapper\u001b[0;34m(artist, renderer)\u001b[0m\n\u001b[1;32m     69\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     70\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mstart_filter()\n\u001b[0;32m---> 72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m draw(artist, renderer)\n\u001b[1;32m     73\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m     74\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/axes/_base.py:3070\u001b[0m, in \u001b[0;36m_AxesBase.draw\u001b[0;34m(self, renderer)\u001b[0m\n\u001b[1;32m   3067\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m artists_rasterized:\n\u001b[1;32m   3068\u001b[0m     _draw_rasterized(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfigure, artists_rasterized, renderer)\n\u001b[0;32m-> 3070\u001b[0m mimage\u001b[38;5;241m.\u001b[39m_draw_list_compositing_images(\n\u001b[1;32m   3071\u001b[0m     renderer, \u001b[38;5;28mself\u001b[39m, artists, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfigure\u001b[38;5;241m.\u001b[39msuppressComposite)\n\u001b[1;32m   3073\u001b[0m renderer\u001b[38;5;241m.\u001b[39mclose_group(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124maxes\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m   3074\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstale \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/image.py:132\u001b[0m, in \u001b[0;36m_draw_list_compositing_images\u001b[0;34m(renderer, parent, artists, suppress_composite)\u001b[0m\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m not_composite \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m has_images:\n\u001b[1;32m    131\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m artists:\n\u001b[0;32m--> 132\u001b[0m         a\u001b[38;5;241m.\u001b[39mdraw(renderer)\n\u001b[1;32m    133\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    134\u001b[0m     \u001b[38;5;66;03m# Composite any adjacent images together\u001b[39;00m\n\u001b[1;32m    135\u001b[0m     image_group \u001b[38;5;241m=\u001b[39m []\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/artist.py:72\u001b[0m, in \u001b[0;36mallow_rasterization.<locals>.draw_wrapper\u001b[0;34m(artist, renderer)\u001b[0m\n\u001b[1;32m     69\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     70\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mstart_filter()\n\u001b[0;32m---> 72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m draw(artist, renderer)\n\u001b[1;32m     73\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m     74\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/image.py:649\u001b[0m, in \u001b[0;36m_ImageBase.draw\u001b[0;34m(self, renderer, *args, **kwargs)\u001b[0m\n\u001b[1;32m    647\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mdraw_image(gc, l, b, im, trans)\n\u001b[1;32m    648\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 649\u001b[0m     im, l, b, trans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmake_image(\n\u001b[1;32m    650\u001b[0m         renderer, renderer\u001b[38;5;241m.\u001b[39mget_image_magnification())\n\u001b[1;32m    651\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m im \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    652\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mdraw_image(gc, l, b, im)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/image.py:939\u001b[0m, in \u001b[0;36mAxesImage.make_image\u001b[0;34m(self, renderer, magnification, unsampled)\u001b[0m\n\u001b[1;32m    936\u001b[0m transformed_bbox \u001b[38;5;241m=\u001b[39m TransformedBbox(bbox, trans)\n\u001b[1;32m    937\u001b[0m clip \u001b[38;5;241m=\u001b[39m ((\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_clip_box() \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maxes\u001b[38;5;241m.\u001b[39mbbox) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_clip_on()\n\u001b[1;32m    938\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfigure\u001b[38;5;241m.\u001b[39mbbox)\n\u001b[0;32m--> 939\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_make_image(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_A, bbox, transformed_bbox, clip,\n\u001b[1;32m    940\u001b[0m                         magnification, unsampled\u001b[38;5;241m=\u001b[39munsampled)\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/image.py:526\u001b[0m, in \u001b[0;36m_ImageBase._make_image\u001b[0;34m(self, A, in_bbox, out_bbox, clip_bbox, magnification, unsampled, round_to_pixel_border)\u001b[0m\n\u001b[1;32m    521\u001b[0m mask \u001b[38;5;241m=\u001b[39m (np\u001b[38;5;241m.\u001b[39mwhere(A\u001b[38;5;241m.\u001b[39mmask, np\u001b[38;5;241m.\u001b[39mfloat32(np\u001b[38;5;241m.\u001b[39mnan), np\u001b[38;5;241m.\u001b[39mfloat32(\u001b[38;5;241m1\u001b[39m))\n\u001b[1;32m    522\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m A\u001b[38;5;241m.\u001b[39mmask\u001b[38;5;241m.\u001b[39mshape \u001b[38;5;241m==\u001b[39m A\u001b[38;5;241m.\u001b[39mshape  \u001b[38;5;66;03m# nontrivial mask\u001b[39;00m\n\u001b[1;32m    523\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m np\u001b[38;5;241m.\u001b[39mones_like(A, np\u001b[38;5;241m.\u001b[39mfloat32))\n\u001b[1;32m    524\u001b[0m \u001b[38;5;66;03m# we always have to interpolate the mask to account for\u001b[39;00m\n\u001b[1;32m    525\u001b[0m \u001b[38;5;66;03m# non-affine transformations\u001b[39;00m\n\u001b[0;32m--> 526\u001b[0m out_alpha \u001b[38;5;241m=\u001b[39m _resample(\u001b[38;5;28mself\u001b[39m, mask, out_shape, t, resample\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    527\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m mask  \u001b[38;5;66;03m# Make sure we don't use mask anymore!\u001b[39;00m\n\u001b[1;32m    528\u001b[0m \u001b[38;5;66;03m# Agg updates out_alpha in place.  If the pixel has no image\u001b[39;00m\n\u001b[1;32m    529\u001b[0m \u001b[38;5;66;03m# data it will not be updated (and still be 0 as we initialized\u001b[39;00m\n\u001b[1;32m    530\u001b[0m \u001b[38;5;66;03m# it), if input data that would go into that output pixel than\u001b[39;00m\n\u001b[1;32m    531\u001b[0m \u001b[38;5;66;03m# it will be `nan`, if all the input data for a pixel is good\u001b[39;00m\n\u001b[1;32m    532\u001b[0m \u001b[38;5;66;03m# it will be 1, and if there is _some_ good data in that output\u001b[39;00m\n\u001b[1;32m    533\u001b[0m \u001b[38;5;66;03m# pixel it will be between [0, 1] (such as a rotated image).\u001b[39;00m\n",
+      "File \u001b[0;32m~/.conda/envs/frbnn/lib/python3.11/site-packages/matplotlib/image.py:208\u001b[0m, in \u001b[0;36m_resample\u001b[0;34m(image_obj, data, out_shape, transform, resample, alpha)\u001b[0m\n\u001b[1;32m    206\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m resample \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    207\u001b[0m     resample \u001b[38;5;241m=\u001b[39m image_obj\u001b[38;5;241m.\u001b[39mget_resample()\n\u001b[0;32m--> 208\u001b[0m _image\u001b[38;5;241m.\u001b[39mresample(data, out, transform,\n\u001b[1;32m    209\u001b[0m                 _interpd_[interpolation],\n\u001b[1;32m    210\u001b[0m                 resample,\n\u001b[1;32m    211\u001b[0m                 alpha,\n\u001b[1;32m    212\u001b[0m                 image_obj\u001b[38;5;241m.\u001b[39mget_filternorm(),\n\u001b[1;32m    213\u001b[0m                 image_obj\u001b[38;5;241m.\u001b[39mget_filterrad())\n\u001b[1;32m    214\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "import sigpyproc.readers as r\n",
+    "import cv2\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "from scipy.special import softmax\n",
+    "from tqdm import tqdm\n",
+    "%matplotlib inline\n",
+    "\n",
+    "header = fil.header\n",
+    "triggers = []\n",
+    "counter = 0\n",
+    "for i in tqdm(range(2048,header.nsamples, 2048)):\n",
+    "    data = torch.tensor(fil.read_block(i-1024, 2048)).cuda()\n",
+    "    # Shuffle the tensor using the random indices\n",
+    "    out = model(transform(torch.tensor(data).cuda())[None])\n",
+    "    out = softmax(out.detach().cpu().numpy(), axis=1)\n",
+    "    triggers.append(out)\n",
+    "    counter += 1\n",
+    "    # if counter > 1000:\n",
+    "    #     break\n",
+    "    # if out[0, 1]>0.999:\n",
+    "    #     key = data.cpu().numpy()\n",
+    "    #     plt.figure(figsize=(10,10))\n",
+    "    #     plt.imshow(data.cpu().numpy(), aspect = 10, vmax = 54557.824)\n",
+    "stack = np.stack(triggers)\n",
+    "positives = stack[:,0,1]\n",
+    "num_pos = np.where(positives > 0.999)[0].shape[0]\n",
+    "print(num_pos)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "64df934d-f4a2-49f0-857d-2661b1d78b21",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.flipud()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1eafb2c1-857e-48be-aa8b-18669c0e0f8c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10,10))\n",
+    "# plt.imshow(key, aspect = 10, vmax = 54557.824)\n",
+    "plt.imshow(key, aspect = 10)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed3783c3-8ed1-46d6-91e4-e906dfa44913",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "key.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8b56a356-a582-4f5d-a8e2-20f725a48fb3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "total_data =[]\n",
+    "for i in range(32):\n",
+    "    total_data.append(key)\n",
+    "total_data = torch.tensor(np.array(total_data))\n",
+    "total_data.cpu().detach().numpy().tofile(\"crab_in.bin\")\n",
+    "print(total_data.shape)\n",
+    "outputs_data = []\n",
+    "for i in range(32):\n",
+    "    temp = model(transform(total_data.cuda()[i,:,:])[None])\n",
+    "    print(temp)\n",
+    "    # outputs_data.append(softmax(temp.detach().cpu().numpy(), axis=1))\n",
+    "    outputs_data.append(temp.detach().cpu().numpy())\n",
+    "outputs_data = torch.tensor(outputs_data)\n",
+    "outputs_data.cpu().detach().numpy().tofile(\"crab_out.bin\")"
+   ]
+  }
+ ],
+ "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/.ipynb_checkpoints/resnet_model-checkpoint.py

@@ -0,0 +1,160 @@
+import torch
+import torch.nn as  nn
+import torch.nn.functional as F
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride = 1, downsample = None):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Sequential(
+                        nn.Conv2d(in_channels, out_channels, kernel_size = 3, stride = stride, padding = 1),
+                        nn.BatchNorm2d(out_channels),
+                        nn.ReLU())
+        self.conv2 = nn.Sequential(
+                        nn.Conv2d(out_channels, out_channels, kernel_size = 3, stride = 1, padding = 1),
+                        nn.BatchNorm2d(out_channels))
+        self.downsample = downsample
+        self.relu = nn.ReLU()
+        self.out_channels = out_channels
+        self.dropout_percentage = 0.5
+        self.dropout1 = nn.Dropout(p=self.dropout_percentage)
+        self.batchnorm_mod = nn.BatchNorm2d(out_channels)
+        
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.dropout1(out)
+        # out = self.batchnorm_mod(out)
+        out = self.conv2(out)
+        out = self.dropout1(out)
+        # out = self.batchnorm_mod(out)
+        if self.downsample:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, inchan, block, layers, num_classes = 10):
+        super(ResNet, self).__init__()
+        self.inplanes = 64
+        self.eps = 1e-5
+        self.relu = nn.ReLU()
+        self.conv1 = nn.Sequential(
+                        nn.Conv2d(inchan, 64, kernel_size = 7, stride = 2, padding = 3),
+                        nn.BatchNorm2d(64),
+                        nn.ReLU())
+        self.maxpool = nn.MaxPool2d(kernel_size = (2, 2), stride = 2, padding = 1)
+        self.layer0 = self._make_layer(block, 64, layers[0], stride = 1)
+        self.layer1 = self._make_layer(block, 128, layers[1], stride = 2)
+        self.layer2 = self._make_layer(block, 256, layers[2], stride = 2)
+        self.layer3 = self._make_layer(block, 512, layers[3], stride = 1)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc = nn.Linear(39424, num_classes)
+        self.dropout_percentage = 0.3
+        self.dropout1 = nn.Dropout(p=self.dropout_percentage)
+
+        # Encoder
+        self.encoder = nn.Sequential(
+            nn.Conv2d(24, 32,  kernel_size = 3, stride =1, padding = 1),
+            nn.ReLU(True),nn.Dropout(p=self.dropout_percentage),
+            nn.Conv2d(32, 64,  kernel_size = 3, stride =1, padding = 1),
+            nn.ReLU(True),nn.Dropout(p=self.dropout_percentage),
+            nn.Conv2d(64, 32,  kernel_size = 3, stride = 1, padding = 1),
+            nn.ReLU(True),nn.Dropout(p=self.dropout_percentage),
+            nn.Conv2d(32, 24, kernel_size = 3, stride = 1, padding = 1),
+            nn.Sigmoid()
+        )
+        params = sum(p.numel() for p in self.encoder.parameters())
+        print("num params encoder ",params)
+
+    def norm(self, x):
+        shifted = x-x.min()
+        maxes = torch.amax(abs(shifted), dim=(-2, -1))
+        repeated_maxes = maxes.unsqueeze(2).unsqueeze(3).repeat(1, 1, x.shape[-2],x.shape[-1])
+        x = shifted/repeated_maxes
+        return x
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes, kernel_size=1, stride=stride),
+                nn.BatchNorm2d(planes),
+            )
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+        return nn.Sequential(*layers)
+    
+    def forward(self, x, return_mask=False):
+        # x = self.norm(x)
+        x = self.conv1(x)
+        x = self.maxpool(x)
+        x = self.layer0(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.dropout1(x)
+        x = self.fc(x)
+        # return x
+        if return_mask:
+            return x, self.mask, self.value
+        else:
+            return x
+
+
+class ConvAutoencoder(nn.Module):
+    def __init__(self):
+        super(ConvAutoencoder, self).__init__()
+        
+        # Encoder
+        self.encoder = nn.Sequential(
+            nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1),  # (16, 96, 128)
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1), # (32, 48, 64)
+            nn.ReLU(),
+            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # (64, 24, 32)
+            nn.ReLU(),
+            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),# (128, 12, 16)
+            nn.ReLU()
+        )
+        
+        # Fully connected latent space
+        self.fc1 = nn.Linear(128 * 12 * 16, 8)
+        self.fc2 = nn.Linear(8, 128 * 12 * 16)
+        
+        # Decoder
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1), # (64, 24, 32)
+            nn.ReLU(),
+            nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),  # (32, 48, 64)
+            nn.ReLU(),
+            nn.ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, output_padding=1),  # (16, 96, 128)
+            nn.ReLU(),
+            nn.ConvTranspose2d(16, 3, kernel_size=3, stride=2, padding=1, output_padding=1),   # (3, 192, 256)
+            nn.Sigmoid()  # Using Sigmoid for the final activation to get output in range [0, 1]
+        )
+        
+    def forward(self, x):
+        # Encode
+        x = self.encoder(x)
+        
+        # Flatten the encoded output
+        x = x.view(x.size(0), -1)
+        
+        # Fully connected latent space
+        x = self.fc1(x)
+        x = self.fc2(x)
+        
+        # Reshape the output to the shape suitable for the decoder
+        x = x.view(x.size(0), 128, 12, 16)
+        
+        # Decode
+        x = self.decoder(x)
+        
+        return x

models/.ipynb_checkpoints/resnet_model_mask-checkpoint.py

@@ -0,0 +1,166 @@
+import torch
+import torch.nn as  nn
+import torch.nn.functional as F
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride = 1, downsample = None):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Sequential(
+                        nn.Conv2d(in_channels, out_channels, kernel_size = 3, stride = stride, padding = 1),
+                        nn.BatchNorm2d(out_channels),
+                        nn.ReLU())
+        self.conv2 = nn.Sequential(
+                        nn.Conv2d(out_channels, out_channels, kernel_size = 3, stride = 1, padding = 1),
+                        nn.BatchNorm2d(out_channels))
+        self.downsample = downsample
+        self.relu = nn.ReLU()
+        self.out_channels = out_channels
+        self.dropout_percentage = 0.5
+        self.dropout1 = nn.Dropout(p=self.dropout_percentage)
+        self.batchnorm_mod = nn.BatchNorm2d(out_channels)
+        
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.dropout1(out)
+        # out = self.batchnorm_mod(out)
+        out = self.conv2(out)
+        out = self.dropout1(out)
+        # out = self.batchnorm_mod(out)
+        if self.downsample:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, inchan, block, layers, num_classes = 10):
+        super(ResNet, self).__init__()
+        self.inplanes = 64
+        self.eps = 1e-5
+        self.relu = nn.ReLU()
+        self.conv1 = nn.Sequential(
+                        nn.Conv2d(inchan, 64, kernel_size = 7, stride = 2, padding = 3),
+                        nn.BatchNorm2d(64),
+                        nn.ReLU())
+        self.maxpool = nn.MaxPool2d(kernel_size = (2, 2), stride = 2, padding = 1)
+        self.layer0 = self._make_layer(block, 64, layers[0], stride = 1)
+        self.layer1 = self._make_layer(block, 128, layers[1], stride = 2)
+        self.layer2 = self._make_layer(block, 256, layers[2], stride = 2)
+        self.layer3 = self._make_layer(block, 512, layers[3], stride = 1)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc = nn.Linear(39424, num_classes)
+        self.dropout_percentage = 0.3
+        self.dropout1 = nn.Dropout(p=self.dropout_percentage)
+
+        # Encoder
+        self.encoder = nn.Sequential(
+            nn.Conv2d(24, 32,  kernel_size = 3, stride =1, padding = 1),
+            nn.ReLU(True),nn.Dropout(p=self.dropout_percentage),
+            nn.Conv2d(32, 64,  kernel_size = 3, stride =1, padding = 1),
+            nn.ReLU(True),nn.Dropout(p=self.dropout_percentage),
+            nn.Conv2d(64, 32,  kernel_size = 3, stride = 1, padding = 1),
+            nn.ReLU(True),nn.Dropout(p=self.dropout_percentage),
+            nn.Conv2d(32, 24, kernel_size = 3, stride = 1, padding = 1),
+            nn.Sigmoid()
+        )
+        params = sum(p.numel() for p in self.encoder.parameters())
+        print("num params encoder ",params)
+
+    def norm(self, x):
+        shifted = x-x.min()
+        maxes = torch.amax(abs(shifted), dim=(-2, -1))
+        repeated_maxes = maxes.unsqueeze(2).unsqueeze(3).repeat(1, 1, x.shape[-2],x.shape[-1])
+        x = shifted/repeated_maxes
+        return x
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes, kernel_size=1, stride=stride),
+                nn.BatchNorm2d(planes),
+            )
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+        return nn.Sequential(*layers)
+    
+    def forward(self, x, return_mask=False):
+        # # m = self.encoder(x).unsqueeze(-1).repeat(1, 1, 1, x.shape[-1])
+        m = self.encoder(x)
+        self.mask = m
+        self.value = x
+        # # m = nn.Sigmoid()(self.encoder(x))
+        x = x * m
+        # x = self.norm(x)
+        x = self.conv1(x)
+        x = self.maxpool(x)
+        x = self.layer0(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.dropout1(x)
+        x = self.fc(x)
+        return x
+        # if return_mask:
+        #     return x, self.mask, self.value
+        # else:
+        #     return x
+
+
+class ConvAutoencoder(nn.Module):
+    def __init__(self):
+        super(ConvAutoencoder, self).__init__()
+        
+        # Encoder
+        self.encoder = nn.Sequential(
+            nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1),  # (16, 96, 128)
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1), # (32, 48, 64)
+            nn.ReLU(),
+            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # (64, 24, 32)
+            nn.ReLU(),
+            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),# (128, 12, 16)
+            nn.ReLU()
+        )
+        
+        # Fully connected latent space
+        self.fc1 = nn.Linear(128 * 12 * 16, 8)
+        self.fc2 = nn.Linear(8, 128 * 12 * 16)
+        
+        # Decoder
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1), # (64, 24, 32)
+            nn.ReLU(),
+            nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),  # (32, 48, 64)
+            nn.ReLU(),
+            nn.ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, output_padding=1),  # (16, 96, 128)
+            nn.ReLU(),
+            nn.ConvTranspose2d(16, 3, kernel_size=3, stride=2, padding=1, output_padding=1),   # (3, 192, 256)
+            nn.Sigmoid()  # Using Sigmoid for the final activation to get output in range [0, 1]
+        )
+        
+    def forward(self, x):
+        # Encode
+        x = self.encoder(x)
+        
+        # Flatten the encoded output
+        x = x.view(x.size(0), -1)
+        
+        # Fully connected latent space
+        x = self.fc1(x)
+        x = self.fc2(x)
+        
+        # Reshape the output to the shape suitable for the decoder
+        x = x.view(x.size(0), 128, 12, 16)
+        
+        # Decode
+        x = self.decoder(x)
+        
+        return x

models/.ipynb_checkpoints/train-checkpoint.py

@@ -0,0 +1,105 @@
+from utils import CustomDataset, transform, preproc, Convert_ONNX
+from torch.utils.data import Dataset, DataLoader
+import torch
+import numpy as np
+from resnet_model import  ResidualBlock, ResNet
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import tqdm 
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import pickle
+import sys 
+
+ind = int(sys.argv[1])
+seeds = [1,42,7109,2002,32]
+seed = seeds[ind]
+print("using seed: ",seed)
+torch.manual_seed(seed)
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+num_gpus = torch.cuda.device_count()
+print(num_gpus)
+
+# Create custom dataset instance
+data_dir = '/mnt/buf1/pma/frbnn/train_ready'
+dataset = CustomDataset(data_dir, transform=transform)
+valid_data_dir = '/mnt/buf1/pma/frbnn/valid_ready'
+valid_dataset = CustomDataset(valid_data_dir, transform=transform)
+
+
+num_classes = 2
+trainloader = DataLoader(dataset, batch_size=420, shuffle=True, num_workers=32)
+validloader = DataLoader(valid_dataset, batch_size=512, shuffle=True, num_workers=32)
+
+model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)
+model = nn.DataParallel(model)
+model = model.to(device)
+params = sum(p.numel() for p in model.parameters())
+print("num params ",params)
+torch.save(model.state_dict(), 'models/test.pt')
+model.load_state_dict(torch.load('models/test.pt'))
+
+preproc_model = preproc()
+Convert_ONNX(model.module,'models/test.onnx', input_data_mock=torch.randn(1, 24, 192, 256).to(device))
+Convert_ONNX(preproc_model,'models/preproc.onnx', input_data_mock=torch.randn(32, 192, 2048).to(device))
+
+# Define optimizer and loss function
+
+criterion = nn.CrossEntropyLoss(weight = torch.tensor([1,1]).to(device))
+optimizer = optim.Adam(model.parameters(), lr=0.0001)
+scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10)
+
+
+from tqdm import tqdm
+# Training loop
+epochs = 10000
+for epoch in range(epochs):
+    running_loss = 0.0
+    correct_train = 0
+    total_train = 0
+    with tqdm(trainloader, unit="batch") as tepoch:
+        model.train()
+        for i, (images, labels) in enumerate(tepoch):
+            inputs, labels = images.to(device), labels.to(device).float()
+            optimizer.zero_grad()
+            outputs = model(inputs, return_mask=False).to(device)
+            new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32).to(device)
+            loss = criterion(outputs, new_label)
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item()
+            # Calculate training accuracy
+            _, predicted = torch.max(outputs.data, 1)
+            total_train += labels.size(0)
+            correct_train += (predicted == labels).sum().item()        
+        val_loss = 0.0
+        correct_valid = 0
+        total = 0
+        model.eval()
+        with torch.no_grad():
+            for images, labels in validloader:
+                inputs, labels = images.to(device), labels.to(device).float()
+                optimizer.zero_grad()
+                outputs = model(inputs, return_mask=False)
+                new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32)
+                loss = criterion(outputs, new_label)
+                val_loss += loss.item()
+                _, predicted = torch.max(outputs, 1)
+                total += labels.size(0)
+                correct_valid += (predicted == labels).sum().item()
+    scheduler.step(val_loss)
+    # Calculate training accuracy after each epoch
+    train_accuracy = 100 * correct_train / total_train
+    val_accuracy = correct_valid / total * 100.0
+    torch.save(model.state_dict(), 'models/model-'+str(epoch)+'-'+str(val_accuracy)+'.pt')
+    Convert_ONNX(model.module,'models/model-'+str(epoch)+'-'+str(val_accuracy)+'.onnx', input_data_mock=inputs)
+
+    print("===========================")
+    print('accuracy: ', epoch, train_accuracy,  val_accuracy)
+    print('learning rate: ', scheduler.get_last_lr())
+    print("===========================")
+    if scheduler.get_last_lr()[0] <1e-6:
+        break

models/.ipynb_checkpoints/train-mask-8-checkpoint.py

@@ -0,0 +1,103 @@
+from utils import CustomDataset, transform, preproc, Convert_ONNX
+from torch.utils.data import Dataset, DataLoader
+import torch
+import numpy as np
+from resnet_model_mask import  ResidualBlock, ResNet
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import tqdm 
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import pickle
+import sys
+# [1,42,7109,2002,32]
+ind = int(sys.argv[1])
+seeds = [1,42,7109,2002,32]
+seed = seeds[ind]
+torch.manual_seed(seed)
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+num_gpus = torch.cuda.device_count()
+print(num_gpus)
+
+# Create custom dataset instance
+data_dir = '/mnt/buf1/pma/frbnn/train_ready'
+dataset = CustomDataset(data_dir,  bit8 = True, transform=transform)
+valid_data_dir = '/mnt/buf1/pma/frbnn/valid_ready'
+valid_dataset = CustomDataset(valid_data_dir,  bit8 = True, transform=transform)
+
+
+num_classes = 2
+trainloader = DataLoader(dataset, batch_size=420,  shuffle=True, num_workers=32)
+validloader = DataLoader(valid_dataset, batch_size=512, shuffle=True, num_workers=32)
+
+model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)
+model = nn.DataParallel(model)
+model = model.to(device)
+params = sum(p.numel() for p in model.parameters())
+print("num params ",params)
+torch.save(model.state_dict(), f'models_8/test_{seed}.pt')
+model.load_state_dict(torch.load(f'models_8/test_{seed}.pt'))
+
+preproc_model = preproc()
+Convert_ONNX(model.module,f'models_8/test_{seed}.onnx', input_data_mock=torch.randn(1, 24, 192, 256).to(device))
+Convert_ONNX(preproc_model,f'models_8/preproc_{seed}.onnx', input_data_mock=torch.randn(32, 192, 2048).to(device))
+
+# Define optimizer and loss function
+
+criterion = nn.CrossEntropyLoss(weight = torch.tensor([1,1]).to(device))
+optimizer = optim.Adam(model.parameters(), lr=0.0001)
+scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10)
+
+
+from tqdm import tqdm
+# Training loop
+epochs = 10000
+for epoch in range(epochs):
+    running_loss = 0.0
+    correct_train = 0
+    total_train = 0
+    with tqdm(trainloader, unit="batch") as tepoch:
+        model.train()
+        for i, (images, labels) in enumerate(tepoch):
+            inputs, labels = images.to(device), labels.to(device).float()
+            optimizer.zero_grad()
+            outputs = model(inputs, return_mask=False).to(device)
+            new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32).to(device)
+            loss = criterion(outputs, new_label)
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item()
+            # Calculate training accuracy
+            _, predicted = torch.max(outputs.data, 1)
+            total_train += labels.size(0)
+            correct_train += (predicted == labels).sum().item()        
+        val_loss = 0.0
+        correct_valid = 0
+        total = 0
+        model.eval()
+        with torch.no_grad():
+            for images, labels in validloader:
+                inputs, labels = images.to(device), labels.to(device).float()
+                optimizer.zero_grad()
+                outputs = model(inputs, return_mask=False)
+                new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32)
+                loss = criterion(outputs, new_label)
+                val_loss += loss.item()
+                _, predicted = torch.max(outputs, 1)
+                total += labels.size(0)
+                correct_valid += (predicted == labels).sum().item()
+    scheduler.step(val_loss)
+    # Calculate training accuracy after each epoch
+    train_accuracy = 100 * correct_train / total_train
+    val_accuracy = correct_valid / total * 100.0
+    torch.save(model.state_dict(), 'models_8/model-'+str(epoch)+'-'+str(val_accuracy)+f'_{seed}.pt')
+    Convert_ONNX(model.module,'models_8/model-'+str(epoch)+'-'+str(val_accuracy)+f'_{seed}.onnx', input_data_mock=inputs)
+
+    print("===========================")
+    print('accuracy: ', epoch, train_accuracy,  val_accuracy)
+    print('learning rate: ', scheduler.get_last_lr())
+    print("===========================")
+    if scheduler.get_last_lr()[0] <1e-6:
+        break

models/.ipynb_checkpoints/train-mask-checkpoint.py

@@ -0,0 +1,104 @@
+from utils import CustomDataset, transform, preproc, Convert_ONNX
+from torch.utils.data import Dataset, DataLoader
+import torch
+import numpy as np
+from resnet_model_mask import  ResidualBlock, ResNet
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import tqdm 
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import pickle
+import sys
+
+ind = int(sys.argv[1])
+seeds = [1,42,7109,2002,32]
+seed = seeds[ind]
+print("using seed: ",seed)
+torch.manual_seed(seed)
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+num_gpus = torch.cuda.device_count()
+print(num_gpus)
+
+# Create custom dataset instance
+data_dir = '/mnt/buf1/pma/frbnn/train_ready'
+dataset = CustomDataset(data_dir, transform=transform)
+valid_data_dir = '/mnt/buf1/pma/frbnn/valid_ready'
+valid_dataset = CustomDataset(valid_data_dir, transform=transform)
+
+
+num_classes = 2
+trainloader = DataLoader(dataset, batch_size=420, shuffle=True, num_workers=32)
+validloader = DataLoader(valid_dataset, batch_size=512, shuffle=True, num_workers=32)
+
+model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device)
+model = nn.DataParallel(model)
+model = model.to(device)
+params = sum(p.numel() for p in model.parameters())
+print("num params ",params)
+torch.save(model.state_dict(), f'models_mask/test_{seed}.pt')
+model.load_state_dict(torch.load(f'models_mask/test_{seed}.pt'))
+
+preproc_model = preproc()
+Convert_ONNX(model.module,f'models_mask/test_{seed}.onnx', input_data_mock=torch.randn(1, 24, 192, 256).to(device))
+Convert_ONNX(preproc_model,f'models_mask/preproc_{seed}.onnx', input_data_mock=torch.randn(32, 192, 2048).to(device))
+
+# Define optimizer and loss function
+
+criterion = nn.CrossEntropyLoss(weight = torch.tensor([1,1]).to(device))
+optimizer = optim.Adam(model.parameters(), lr=0.0001)
+scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10)
+
+
+from tqdm import tqdm
+# Training loop
+epochs = 10000
+for epoch in range(epochs):
+    running_loss = 0.0
+    correct_train = 0
+    total_train = 0
+    with tqdm(trainloader, unit="batch") as tepoch:
+        model.train()
+        for i, (images, labels) in enumerate(tepoch):
+            inputs, labels = images.to(device), labels.to(device).float()
+            optimizer.zero_grad()
+            outputs = model(inputs, return_mask=False).to(device)
+            new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32).to(device)
+            loss = criterion(outputs, new_label)
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item()
+            # Calculate training accuracy
+            _, predicted = torch.max(outputs.data, 1)
+            total_train += labels.size(0)
+            correct_train += (predicted == labels).sum().item()        
+        val_loss = 0.0
+        correct_valid = 0
+        total = 0
+        model.eval()
+        with torch.no_grad():
+            for images, labels in validloader:
+                inputs, labels = images.to(device), labels.to(device).float()
+                optimizer.zero_grad()
+                outputs = model(inputs, return_mask=False)
+                new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32)
+                loss = criterion(outputs, new_label)
+                val_loss += loss.item()
+                _, predicted = torch.max(outputs, 1)
+                total += labels.size(0)
+                correct_valid += (predicted == labels).sum().item()
+    scheduler.step(val_loss)
+    # Calculate training accuracy after each epoch
+    train_accuracy = 100 * correct_train / total_train
+    val_accuracy = correct_valid / total * 100.0
+    torch.save(model.state_dict(), 'models_mask/model-'+str(epoch)+'-'+str(val_accuracy)+f'_{seed}.pt')
+    Convert_ONNX(model.module,'models_mask/model-'+str(epoch)+'-'+str(val_accuracy)+f'_{seed}.onnx', input_data_mock=inputs)
+
+    print("===========================")
+    print('accuracy: ', epoch, train_accuracy,  val_accuracy)
+    print('learning rate: ', scheduler.get_last_lr())
+    print("===========================")
+    if scheduler.get_last_lr()[0] <1e-6:
+        break

models/.ipynb_checkpoints/utils-checkpoint.py

@@ -0,0 +1,393 @@
+import numpy as np
+import pickle
+from torch.utils.data import Dataset, DataLoader
+import os
+import torch
+from copy import deepcopy
+from blimpy import Waterfall
+from tqdm import tqdm
+from copy import deepcopy
+from sigpyproc.readers import FilReader
+from torch import nn
+
+
+def load_pickled_data(file_path):
+    with open(file_path, 'rb') as f:
+        data = pickle.load(f)
+    return data
+
+# Custom dataset class
+class CustomDataset(Dataset):
+    def __init__(self, data_dir, bit8=False, transform=None):
+        self.data_dir = data_dir
+        self.transform = transform
+        self.images = []
+        self.labels = []
+        self.classes = os.listdir(data_dir)
+        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
+        self.bit8 = bit8
+        # Load images and labels
+        for cls in self.classes:
+            class_dir = os.path.join(data_dir, cls)
+            for image_name in os.listdir(class_dir):
+                image_path = os.path.join(class_dir, image_name)
+                self.images.append(image_path)
+                self.labels.append(self.class_to_idx[cls])
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image_path = self.images[idx]
+        label = self.labels[idx]
+        # Load image
+        image = load_pickled_data(image_path)
+        if self.transform is not None:
+            if self.bit8 == True:
+                new_image = self.transform(torch.from_numpy(image['8_data']).type(torch.float32))
+            else:
+                new_image = self.transform(torch.from_numpy(image['data']))
+            # new_image = self.transform(image['data'])
+        return new_image, label
+
+# Custom dataset class
+class CustomDataset_Masked(Dataset):
+    def __init__(self, data_dir, transform=None):
+        self.data_dir = data_dir
+        self.transform = transform
+        self.images = []
+        self.labels = []
+        self.classes = os.listdir(data_dir)
+        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
+
+        # Load images and labels
+        for cls in self.classes:
+            class_dir = os.path.join(data_dir, cls)
+            for image_name in os.listdir(class_dir):
+                image_path = os.path.join(class_dir, image_name)
+                self.images.append(image_path)
+                self.labels.append(self.class_to_idx[cls])
+        
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image_path = self.images[idx]
+        
+        label = self.labels[idx]
+        # Load image
+        image = load_pickled_data(image_path)
+        if self.transform is not None:
+            if image['burst'].max() ==0:
+                new_burst = torch.from_numpy(image['burst'])
+            else:
+                new_burst = torch.from_numpy(image['burst']/image['burst'].max())
+            ind = new_burst > 0.1
+            ind_not = new_burst <= 0.1
+            new_burst[ind] = 1
+            new_burst[ind_not] = 0
+            new_image = self.transform(torch.from_numpy(image['data'].data))
+            new_burst_arr = torch.zeros_like(new_image)
+            new_burst_arr[ 0, :,:] =   new_burst 
+            new_burst_arr[ 1, :,:] =   new_burst 
+            new_burst_arr[ 2, :,:] =   new_burst 
+        return new_image, label, new_burst_arr
+
+# Custom dataset class
+class TestingDataset(Dataset):
+    def __init__(self, data_dir, bit8=False, transform=None):
+        self.data_dir = data_dir
+        self.transform = transform
+        self.images = []
+        self.labels = []
+        self.classes = os.listdir(data_dir)
+        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
+        self.bit8 = bit8
+        # Load images and labels
+        for cls in self.classes:
+            class_dir = os.path.join(data_dir, cls)
+            for image_name in os.listdir(class_dir):
+                image_path = os.path.join(class_dir, image_name)
+                self.images.append(image_path)
+                self.labels.append(self.class_to_idx[cls])
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image_path = self.images[idx]
+        label = self.labels[idx]
+        # Load image
+        image = load_pickled_data(image_path)
+        params = image['params']
+        if self.transform is not None:
+            params = image['params']
+            if self.bit8 == True:
+                new_image = self.transform(torch.from_numpy(image['8_data']).type(torch.float32))
+            else:
+                new_image = self.transform(torch.from_numpy(image['data']))
+        params['labels'] = label
+        return new_image, (label, params['dm'], params['freq_ref'], params['snr'],  params['boxcard'])
+
+# Custom dataset class
+class SearchDataset(Dataset):
+    def __init__(self, data_dir, transform=None, pickle_data=False):
+        self.window_size = 2048
+        
+        if pickle_data:
+            with open(data_dir, 'rb') as f:
+                self.d = pickle.load(f)
+            self.header = self.d['header']
+            self.images = self.crop(self.d['data'][:,0,:], self.window_size)
+        else:
+            self.obs = Waterfall(data_dir, max_load = 50)
+            self.header = self.obs.header
+            self.images = self.crop(self.obs.data[:,0,:], self.window_size)
+        self.transform = transform
+        self.SEC_PER_DAY = 86400
+        
+    def crop(self, data, window_size = 2048):
+        n_samp = data.shape[0]//window_size
+        new_data = np.zeros((n_samp, window_size, 192 ))
+        for i in range(n_samp):
+            new_data[i, :,:] = data[ i*window_size : (i+1)*window_size, :]
+        return new_data
+        
+    def __len__(self):
+        return self.images.shape[0]
+    def __getitem__(self, idx):
+        data = self.images[idx, :, :].T 
+        tindex = idx * self.window_size
+        time = self.header['tsamp'] * tindex / self.SEC_PER_DAY + self.header['tstart']
+        if self.transform is not None:
+            new_image = self.transform(data)
+        return new_image, idx
+
+# Custom dataset class
+class SearchDataset_Sigproc(Dataset):
+    def __init__(self, data_dir, transform=None):
+        self.window_size = 2048
+        fil = FilReader(data_dir)
+        self.header = fil.header
+        # print("check shape ",fil.read_block(0, fil.header.nsamples).shape)
+        read_data = fil.read_block(0, fil.header.nsamples)[:,1024:-1024]
+        read_data = np.swapaxes(read_data, 0,-1)
+        self.images = self.crop(read_data, self.window_size) 
+        self.transform = transform
+        self.SEC_PER_DAY = 86400
+        
+    def crop(self, data, window_size = 2048):
+        n_samp = data.shape[0]//window_size
+        new_data = np.zeros((n_samp, window_size, 192 ))
+        for i in range(n_samp):
+            new_data[i, :,:] = data[ i*window_size : (i+1)*window_size, :]
+        return new_data
+        
+    def __len__(self):
+        return self.images.shape[0]
+        
+    def __getitem__(self, idx):
+        data = self.images[idx, :, :].T 
+        tindex = idx * self.window_size
+        time = self.header.tsamp * tindex / self.SEC_PER_DAY + self.header.tstart
+        if self.transform is not None:
+            new_image = self.transform(torch.from_numpy(data))
+        return new_image, idx
+
+# def renorm(data):
+#     shifted = data - data.min()
+#     shifted = shifted/shifted.max()
+#     return shifted
+
+def renorm(data):
+    mean = torch.mean(data)
+    std = torch.std(data)
+    # Standardize the data
+    standardized_data = (data - mean) / std
+    return standardized_data
+
+def transform(data):
+    copy_data = data.detach().clone()
+    rms = torch.std(data)
+    mean = torch.mean(data)
+    masks_rms = [-1, 5]
+    new_data = torch.zeros((len(masks_rms)+1, data.shape[0], data.shape[1]))
+    new_data[0,:,:] = renorm(torch.log10(copy_data+1e-10))  
+    for i in range(1, len(masks_rms)+1):
+        scale = masks_rms[i-1]
+        copy_data = data.detach().clone() #deepcopy(data)
+        if scale < 0:
+            ind = copy_data < abs(scale) * rms  + mean
+            copy_data[ind] = 0
+        else:
+            ind = copy_data > (scale) * rms + mean
+            copy_data[ind] = 0
+        new_data[i,:,:] = renorm(torch.log10(copy_data+1e-10))
+    new_data = new_data.type(torch.float32)
+    slices = torch.chunk(new_data, 8, dim=-1)  # dim=1 is the height dimension
+    new_data = torch.stack(slices, dim=1)  # New axis is inserted at dim=1
+    new_data = new_data.view(-1, new_data.size(2), new_data.size(3))
+    return new_data
+
+
+def renorm_batched(data):
+    mean = torch.mean(data, dim=tuple(range(1, data.ndim)), keepdim=True)
+    std = torch.std(data, dim=tuple(range(1, data.ndim)), keepdim=True)
+    standardized_data = (data - mean) / std
+    return standardized_data
+
+def transform_batched(data):
+    copy_data = data.detach().clone()
+    rms = torch.std(data, dim=tuple(range(1, data.ndim)), keepdim=True)  # Batch-wise std
+    mean = torch.mean(data, dim=tuple(range(1, data.ndim)), keepdim=True)  # Batch-wise mean
+    masks_rms = [-1, 5]
+    
+    # Prepare the new_data tensor
+    num_masks = len(masks_rms) + 1
+    new_data = torch.zeros((num_masks, *data.shape), device=data.device)  # Shape: (num_masks, batch_size, ..., ...)
+
+    # First layer: Apply renorm(log10(copy_data + epsilon))
+    new_data[0] = renorm_batched(torch.log10(copy_data + 1e-10))
+    for i, scale in enumerate(masks_rms, start=1):
+        copy_data = data.detach().clone()
+        
+        # Apply masking based on the scale
+        if scale < 0:
+            ind = copy_data < abs(scale) * rms + mean
+        else:
+            ind = copy_data > scale * rms + mean
+        copy_data[ind] = 0
+        
+        # Renormalize and log10 transform
+        new_data[i] = renorm_batched(torch.log10(copy_data + 1e-10))
+    
+    # Convert to float32
+    new_data = new_data.type(torch.float32)
+
+    # Chunk along the last dimension and stack
+    slices = torch.chunk(new_data, 8, dim=-1)  # Adjust for batch-wise slicing
+    new_data = torch.stack(slices, dim=2)  # Insert a new axis at dim=1
+    new_data = torch.swapaxes(new_data, 0,1)
+    # Reshape into final format
+    new_data = new_data.reshape( new_data.size(0), 24,  new_data.size(3), new_data.size(4))  # Flatten batch and mask dimensions
+    return new_data
+
+
+
+class preproc(nn.Module):
+    def forward(self, x, flip=True):
+        if flip:
+            transform_batched(torch.flip(x, dims = (-2,)))
+        else:
+            transform_batched(x)
+        return template
+        
+# class preproc_debug(nn.Module):
+#     def forward(self, x):
+#         template = torch.zeros((32, 24, 192, 256))
+#         # for i in torch.arange(x.shape[0]):  # Use a tensor-based range
+#         template[0,:,:,:] = transform_debug(torch.flip(x[0,:,:], dims = (0,)))
+#         template[1,:,:,:] = transform_debug(torch.flip(x[1,:,:], dims = (0,)))
+#         template[2,:,:,:] = transform_debug(torch.flip(x[2,:,:], dims = (0,)))
+#         template[3,:,:,:] = transform_debug(torch.flip(x[3,:,:], dims = (0,)))
+#         template[4,:,:,:] = transform_debug(torch.flip(x[4,:,:], dims = (0,)))
+#         template[5,:,:,:] = transform_debug(torch.flip(x[5,:,:], dims = (0,)))
+#         template[6,:,:,:] = transform_debug(torch.flip(x[6,:,:], dims = (0,)))
+#         template[7,:,:,:] = transform_debug(torch.flip(x[7,:,:], dims = (0,)))
+#         template[8,:,:,:] = transform_debug(torch.flip(x[8,:,:], dims = (0,)))
+#         template[9,:,:,:] = transform_debug(torch.flip(x[9,:,:], dims = (0,)))
+#         template[10,:,:,:] = transform_debug(torch.flip(x[10,:,:], dims = (0,)))
+#         template[11,:,:,:] = transform_debug(torch.flip(x[11,:,:], dims = (0,)))
+#         template[12,:,:,:] = transform_debug(torch.flip(x[12,:,:], dims = (0,)))
+#         template[13,:,:,:] = transform_debug(torch.flip(x[13,:,:], dims = (0,)))
+#         template[14,:,:,:] = transform_debug(torch.flip(x[14,:,:], dims = (0,)))
+#         template[15,:,:,:] = transform_debug(torch.flip(x[15,:,:], dims = (0,)))
+#         template[16,:,:,:] = transform_debug(torch.flip(x[16,:,:], dims = (0,)))
+#         template[17,:,:,:] = transform_debug(torch.flip(x[17,:,:], dims = (0,)))
+#         template[18,:,:,:] = transform_debug(torch.flip(x[18,:,:], dims = (0,)))
+#         template[19,:,:,:] = transform_debug(torch.flip(x[19,:,:], dims = (0,)))
+#         template[20,:,:,:] = transform_debug(torch.flip(x[20,:,:], dims = (0,)))
+#         template[21,:,:,:] = transform_debug(torch.flip(x[21,:,:], dims = (0,)))
+#         template[22,:,:,:] = transform_debug(torch.flip(x[22,:,:], dims = (0,)))
+#         template[23,:,:,:] = transform_debug(torch.flip(x[23,:,:], dims = (0,)))
+#         template[24,:,:,:] = transform_debug(torch.flip(x[24,:,:], dims = (0,)))
+#         template[25,:,:,:] = transform_debug(torch.flip(x[25,:,:], dims = (0,)))
+#         template[26,:,:,:] = transform_debug(torch.flip(x[26,:,:], dims = (0,)))
+#         template[27,:,:,:] = transform_debug(torch.flip(x[27,:,:], dims = (0,)))
+#         template[28,:,:,:] = transform_debug(torch.flip(x[28,:,:], dims = (0,)))
+#         template[29,:,:,:] = transform_debug(torch.flip(x[29,:,:], dims = (0,)))
+#         template[30,:,:,:] = transform_debug(torch.flip(x[30,:,:], dims = (0,)))
+#         template[31,:,:,:] = transform_debug(torch.flip(x[31,:,:], dims = (0,)))
+#         return template
+        
+# def transform_debug(data):
+#     copy_data = data.detach().clone()
+#     rms = torch.std(data)
+#     mean = torch.mean(data)
+#     masks_rms = [-1, 5]
+#     new_data = torch.zeros((len(masks_rms)+1, data.shape[0], data.shape[1]))
+#     new_data[0,:,:] = renorm(torch.log10(copy_data+1e-10))  
+#     for i in range(1, len(masks_rms)+1):
+#         scale = masks_rms[i-1]
+#         copy_data = data.detach().clone()
+#         if scale < 0:
+#             ind = copy_data < abs(scale) * rms  + mean
+#             copy_data[ind] = 0
+#         else:
+#             ind = copy_data > (scale) * rms + mean
+#             copy_data[ind] = 0
+#         new_data[i,:,:] = renorm(torch.log10(copy_data+1e-10))
+#     new_data = new_data.type(torch.float32)
+#     slices = torch.chunk(new_data, 8, dim=-1)  # dim=1 is the height dimension
+#     new_data = torch.stack(slices, dim=1)  # New axis is inserted at dim=1
+#     new_data = new_data.view(-1, new_data.size(2), new_data.size(3))
+#     return new_data
+    
+def renorm_batched(data):
+    mins = torch.amin(data, (-2, -1))
+    mins = mins.unsqueeze(1).unsqueeze(2)
+    mins = mins.expand(data.shape[0], 192, 2048)
+    shifted = data - mins
+    maxs = torch.amax(shifted, (-2, -1))
+    maxs = maxs.unsqueeze(1).unsqueeze(2)
+    maxs = maxs.expand(data.shape[0], 192, 2048)
+    shifted = shifted/maxs
+    return shifted
+    
+
+def transform_mask(data):
+    copy_data = deepcopy(data)
+    shift = copy_data - copy_data.min()
+    normalized_data = shift / shift.max()
+    new_data = np.zeros((3, data.shape[0], data.shape[1]))
+    for i in range(3):
+        new_data[i,:,:] = normalized_data 
+    new_data = new_data.astype(np.float32)
+    return new_data
+
+
+#Function to Convert to ONNX 
+def Convert_ONNX(model, saveloc, input_data_mock): 
+    print("Saving to ONNX")
+    # set the model to inference mode 
+    model.eval()
+
+    # Let's create a dummy input tensor  
+    dummy_input = torch.autograd.Variable(input_data_mock)
+
+    # Export the model   
+    torch.onnx.export(model,         # model being run 
+         dummy_input,       # model input (or a tuple for multiple inputs) 
+         saveloc,       # where to save the model  
+         input_names = ['modelInput'],   # the model's input names 
+         output_names = ['modelOutput'], # the model's output names 
+         dynamic_axes={'modelInput' : {0 : 'batch_size'},    # variable length axes
+                    'modelOutput' : {0 : 'batch_size'}} ) 
+    print(" ") 
+    print('Model has been converted to ONNX')
+
+
+
+
+
+

models/.ipynb_checkpoints/utils_batched_preproc-checkpoint.py

@@ -0,0 +1,65 @@
+import numpy as np
+import pickle
+from torch.utils.data import Dataset, DataLoader
+import os
+import torch
+from copy import deepcopy
+from blimpy import Waterfall
+from tqdm import tqdm
+from copy import deepcopy
+from sigpyproc.readers import FilReader
+from torch import nn
+
+
+def renorm_batched(data):
+    mean = torch.mean(data, dim=tuple(range(1, data.ndim)), keepdim=True)
+    std = torch.std(data, dim=tuple(range(1, data.ndim)), keepdim=True)
+    standardized_data = (data - mean) / std
+    return standardized_data
+
+def transform_batched(data):
+    copy_data = data.detach().clone()
+    rms = torch.std(data, dim=tuple(range(1, data.ndim)), keepdim=True)  # Batch-wise std
+    mean = torch.mean(data, dim=tuple(range(1, data.ndim)), keepdim=True)  # Batch-wise mean
+    masks_rms = [-1, 5]
+    
+    # Prepare the new_data tensor
+    num_masks = len(masks_rms) + 1
+    new_data = torch.zeros((num_masks, *data.shape), device=data.device)  # Shape: (num_masks, batch_size, ..., ...)
+
+    # First layer: Apply renorm(log10(copy_data + epsilon))
+    new_data[0] = renorm_batched(torch.log10(copy_data + 1e-10))
+    for i, scale in enumerate(masks_rms, start=1):
+        copy_data = data.detach().clone()
+        
+        # Apply masking based on the scale
+        if scale < 0:
+            ind = copy_data < abs(scale) * rms + mean
+        else:
+            ind = copy_data > scale * rms + mean
+        copy_data[ind] = 0
+        
+        # Renormalize and log10 transform
+        new_data[i] = renorm_batched(torch.log10(copy_data + 1e-10))
+    
+    # Convert to float32
+    new_data = new_data.type(torch.float32)
+
+    # Chunk along the last dimension and stack
+    slices = torch.chunk(new_data, 8, dim=-1)  # Adjust for batch-wise slicing
+    new_data = torch.stack(slices, dim=2)  # Insert a new axis at dim=1
+    new_data = torch.swapaxes(new_data, 0,1)
+    # Reshape into final format
+    new_data = new_data.reshape( new_data.size(0), 24,  new_data.size(3), new_data.size(4))  # Flatten batch and mask dimensions
+    return new_data
+
+class preproc_flip(nn.Module):
+    def forward(self, x, flip=True):
+        template = transform_batched(torch.flip(x, dims = (-2,)))
+        return template 
+
+class preproc(nn.Module):
+    def forward(self, x, flip=True):
+        template = transform_batched(x)
+        return template
+        

models/HITS-FEB-10.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ce8b602cef03e22c666cdea792741411f623fb5eb0a254ef1ffd9a32864d754
+size 270858960

models/HITS-FEB-10/hit_100000000_1739230556_9.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7cc8774df726c0db3358653dc48894f70f2964cac7604dc5319b44f8f6340b71
+size 1572992

models/HITS-FEB-10/hit_100000000_1739231399_1.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f66e781e7776a11067d0b103b36e9f42b7130d5297ab8936a814af659e125524
+size 1572992

models/HITS-FEB-10/hit_100000000_1739231802_11.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9edde89b6a0526420dcff5e35516979a20f379fac7a2c98d59b9dae897bf4426
+size 1572992

models/HITS-FEB-10/hit_100000000_1739234628_13.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:550f498a863202da1dad26237f7daf0a4f347e7fecde3650e9abee7611994078
+size 1572992

models/HITS-FEB-10/hit_100000000_1739234628_14.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b49681cd9e55c1b3f81be0d2e9c04aacece284581b5b534cf8531ef5464b084
+size 1572992

models/HITS-FEB-10/hit_100000000_1739235333_29.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04f284312e3c2d80e49d617bec875f1b6664bbe21e1418ba9fee8f0cafe70e24
+size 1572992

models/HITS-FEB-10/hit_100000000_1739235841_12.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5b8008a6d13ef7b36fe03fd9dab1e58a6bbb770b433ffb77a3eb11a5033a09f
+size 1572992

models/HITS-FEB-10/hit_50233055_1739232802_29.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2d0a49b0bf3825cfbbdbd91c465764f82b326e172c378685c9de87a44f1296e
+size 1572992

models/HITS-FEB-10/hit_52111435_1739229641_28.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:caa8876c4c3185ad998f1a61a2b4176dec62d5ad38e3ab585aa9ff730848f83f
+size 1572992

models/HITS-FEB-10/hit_52550001_1739233595_4.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a69ae7f971552525b8beac6adee0baa0eb066158fa73e2748265e7a803f8c9f
+size 1572992