From a156a0c2731cb4ad6c7ec507a8fcdf1cf03648a3 Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sun, 23 Aug 2026 16:11:13 +0530 Subject: [PATCH 1/9] chore: setup github actions ci pipeline and 2d tensor shape validation --- .github/workflows/2d_model_ci.yml | 45 ++++++ health_report.md | 48 +++++++ test_2d_shapes.py | 222 ++++++++++++++++++++++++++++++ 3 files changed, 315 insertions(+) create mode 100644 .github/workflows/2d_model_ci.yml create mode 100644 health_report.md create mode 100644 test_2d_shapes.py diff --git a/.github/workflows/2d_model_ci.yml b/.github/workflows/2d_model_ci.yml new file mode 100644 index 0000000..f84c360 --- /dev/null +++ b/.github/workflows/2d_model_ci.yml @@ -0,0 +1,45 @@ +# .github/workflows/2d_model_ci.yml +# ========================================================================= +# CI Pipeline: 2D Model Tensor Shape Validation +# Runs on every push to feature/2d-models to catch tensor shape mismatches +# before they propagate into model training code. +# ========================================================================= + +name: 2D Model CI Pipeline + +on: + push: + branches: [feature/2d-models] + pull_request: + branches: [main, feature/2d-models] + +jobs: + validate-2d-shapes: + name: 2D Tensor Shape Validation + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python 3.10 + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Cache pip dependencies + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-2d-${{ hashFiles('requirements.txt') }} + restore-keys: | + ${{ runner.os }}-pip-2d- + + - name: Install dependencies (CPU-only PyTorch for CI speed) + run: | + python -m pip install --upgrade pip + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install numpy scipy scikit-learn matplotlib + + - name: Run 2D shape validation tests + run: python test_2d_shapes.py diff --git a/health_report.md b/health_report.md new file mode 100644 index 0000000..47169d2 --- /dev/null +++ b/health_report.md @@ -0,0 +1,48 @@ +# Repository Health Report: RF_Signals_Classifier + +## 1. Overview +The **Intelligent RF Signal Classification Pipeline** aims to be an advanced, 4-stage machine learning architecture for Automatic Modulation Classification (AMC). The goal is to use deep learning for robust spectrum sensing. + +However, based on the current codebase, the repository is **only partially implemented**. It currently contains robust data ingestion, synthetic generation, preprocessing, and dataset verification tools. The core Machine Learning models and training pipelines (Stages 3 and 4) are **missing**. + +## 2. Current Status of the 4-Stage Pipeline + +| Stage | Name | Status | Description | +|---|---|---|---| +| **Stage 1** | Ingestion & Dataset Agent | **✅ Implemented** | Both `dataset_agent.py` (for RadioML dataset) and `generate_custom_rf_v4.m` (for synthetic generation) are present and functional. | +| **Stage 2** | Dual-Representation Preprocessing | **✅ Implemented** | Both Time-Domain (I/Q normalization) and Time-Frequency (STFT via `stft_converter.py`) branches exist and are integrated into PyTorch DataLoaders. | +| **Stage 3** | Adaptive CNN-Transformer Core | **❌ Missing** | No PyTorch models, convolutional layers, or Transformer architectures exist in the codebase. | +| **Stage 4** | Spectrum Cop Decision Engine | **❌ Missing** | No Outlier Detection, Continual Learning (EWC loss), or XAI (Grad-CAM) implementations are present. | + +## 3. File-by-File Breakdown + +### Data Preprocessing & Pipeline +* **`dataset_agent.py`**: A robust script that loads the legacy RadioML 2016.10a pickle file, normalizes the I/Q signals (zero-mean, unit-variance), and creates a locked 70/20/10 train/val/test split (saving as `.npy` files). +* **`dataset_loader.py`**: Provides the `RFSignalDataset` class and `get_rf_dataloader` function. Standardizes the PyTorch loading process using the splits generated by `dataset_agent.py`. +* **`verify_splits.py`**: A comprehensive analytics script that verifies data integrity, ensures no data leakage across splits, and generates a detailed visualization dashboard (`verification_results.png`). + +### Custom Dataset Generation +* **`generate_custom_rf_v4.m`**: A MATLAB script that generates custom synthetic RF signals (BPSK, QPSK, 8PSK, 16QAM, 64QAM) with realistic channel impairments (fading, Doppler, CFO). Fallbacks are implemented if the Communications Toolbox is unavailable. +* **`load_mat_dataset.py`**: PyTorch integration script that loads the `.mat` outputs from the MATLAB generator into a PyTorch `TensorDataset`. + +### Feature Engineering +* **`stft_converter.py`**: Handles the Time-Frequency branch of Stage 2. It converts I/Q sequences into Kaiser-windowed 2D power spectrograms for CNN processing. + +### Configuration +* **`requirements.txt`**: Standard Python environment requirements (includes `torch`, `scipy`, `scikit-learn`, `matplotlib`, `numpy`, etc.). +* **`README.md`**: Outlines the architectural framework but contains claims for features that have not yet been implemented in the code. + +## 4. Actionable Next Steps for Development + +To continue the development of this repository with full context, here is what you need to focus on next: + +1. **Generate or Download Data**: + * *Option A*: Download the `RML2016.10a_dict.pkl` dataset and place it in the `./data` directory, then run `python dataset_agent.py`. + * *Option B*: Run `generate_custom_rf_v4.m` in MATLAB to generate synthetic data, then use `load_mat_dataset.py`. +2. **Verify Data Pipeline**: Run `python verify_splits.py` to ensure the ingestion and splitting logic is sound. +3. **Implement Stage 3 (The Core Model)**: + * Create a new file (e.g., `models.py`). + * Implement the 1D/2D CNNs for Spatial Feature Extraction. + * Implement the local-windowed Transformer Encoder for Temporal Modeling. +4. **Implement Training Loop**: Create a `train.py` script to orchestrate the dataloaders (`dataset_loader.py`) and the new PyTorch models. +5. **Implement Stage 4 (Decision Engine)**: Later on, build out the MCD outlier detection, Elastic Weight Consolidation (EWC) for continual learning, and Grad-CAM integration for interpretability. diff --git a/test_2d_shapes.py b/test_2d_shapes.py new file mode 100644 index 0000000..201ad9d --- /dev/null +++ b/test_2d_shapes.py @@ -0,0 +1,222 @@ +# test_2d_shapes.py +# ========================================================================= +# CHUNK 1: DevOps Scaffolding & 2D Data Ingestion +# Validates that Person 3's Kaiser-windowed STFT converter produces +# correctly shaped 2D tensors for downstream CNN models (ResNet-18, STFT-RADN). +# Uses synthetic I/Q data so this runs in CI without real dataset files. +# ========================================================================= + +import sys +import numpy as np +import torch +from scipy.signal import stft as scipy_stft + + +# ────────────────────────────────────────────────────────────── +# Constants: Must match Person 3's stft_converter.py parameters +# ────────────────────────────────────────────────────────────── + +STFT_WINDOW = ("kaiser", 0.85) +STFT_NPERSEG = 64 +STFT_NOVERLAP = 32 +SEQ_LEN = 128 # RadioML 2016.10a I/Q sequence length +BATCH_SIZE = 8 + +# Expected spectrogram dimensions from complex-input two-sided STFT: +# H = nperseg = 64 (frequency bins, two-sided for complex signal) +# W = 5 (time bins: scipy default boundary='zeros', padded=True) +EXPECTED_H = 64 +EXPECTED_W = 5 + + +# ────────────────────────────────────────────────────────────── +# Converter Utilities (mirror Person 3's logic for CI isolation) +# ────────────────────────────────────────────────────────────── + +def compute_spectrogram(signal_np): + """ + Convert a single (2, 128) I/Q numpy array to a 2D power spectrogram. + Mirrors the exact logic in stft_converter.py → compute_stft(). + + Returns: + power_db: np.ndarray of shape (H, W) + """ + I = signal_np[0] + Q = signal_np[1] + complex_signal = I + 1j * Q + + _, _, Zxx = scipy_stft( + complex_signal, + window=STFT_WINDOW, + nperseg=STFT_NPERSEG, + noverlap=STFT_NOVERLAP, + ) + + power = np.abs(Zxx) ** 2 + power_db = 10 * np.log10(power + 1e-10) + + # Center zero-frequency (matches Person 3's fftshift) + power_db = np.fft.fftshift(power_db, axes=0) + + return power_db + + +def batch_to_grayscale_tensor(signals_np): + """ + Convert a batch of (B, 2, 128) I/Q signals to grayscale spectrogram tensors. + + Returns: + torch.Tensor of shape [B, 1, H, W] + """ + spectrograms = [] + for i in range(signals_np.shape[0]): + spec = compute_spectrogram(signals_np[i]) + spectrograms.append(spec) + + specs = np.array(spectrograms, dtype=np.float32) # (B, H, W) + tensor = torch.from_numpy(specs).unsqueeze(1) # (B, 1, H, W) + return tensor + + +def batch_to_hybrid_tensor(signals_np): + """ + Convert a batch of (B, 2, 128) I/Q signals to 3-channel hybrid tensors. + Channels: [Power (dB), Phase, Magnitude] + Compatible with pretrained ResNet-18 that expects 3-channel input. + + Returns: + torch.Tensor of shape [B, 3, H, W] + """ + batch = [] + for i in range(signals_np.shape[0]): + I = signals_np[i, 0] + Q = signals_np[i, 1] + complex_signal = I + 1j * Q + + _, _, Zxx = scipy_stft( + complex_signal, + window=STFT_WINDOW, + nperseg=STFT_NPERSEG, + noverlap=STFT_NOVERLAP, + ) + Zxx = np.fft.fftshift(Zxx, axes=0) + + power_db = 10 * np.log10(np.abs(Zxx) ** 2 + 1e-10) + phase = np.angle(Zxx) + magnitude = np.abs(Zxx) + + hybrid = np.stack([power_db, phase, magnitude], axis=0) # (3, H, W) + batch.append(hybrid) + + batch_np = np.array(batch, dtype=np.float32) # (B, 3, H, W) + return torch.from_numpy(batch_np) + + +# ────────────────────────────────────────────────────────────── +# Test Suite +# ────────────────────────────────────────────────────────────── + +def test_single_spectrogram_shape(): + """Verify a single spectrogram has shape (H, W) = (64, 5).""" + signal = np.random.randn(2, SEQ_LEN).astype(np.float32) + spec = compute_spectrogram(signal) + + assert spec.shape == (EXPECTED_H, EXPECTED_W), \ + f"Expected ({EXPECTED_H}, {EXPECTED_W}), got {spec.shape}" + print(f"[PASS] Single spectrogram shape: {spec.shape}") + + +def test_grayscale_batch_shape(): + """Verify grayscale batch has shape [B, 1, H, W].""" + signals = np.random.randn(BATCH_SIZE, 2, SEQ_LEN).astype(np.float32) + tensor = batch_to_grayscale_tensor(signals) + + expected = (BATCH_SIZE, 1, EXPECTED_H, EXPECTED_W) + assert tensor.shape == torch.Size(expected), \ + f"Expected {expected}, got {tuple(tensor.shape)}" + assert tensor.dtype == torch.float32, \ + f"Expected float32, got {tensor.dtype}" + print(f"[PASS] Grayscale batch shape: {tuple(tensor.shape)}") + + +def test_hybrid_batch_shape(): + """Verify 3-channel hybrid batch has shape [B, 3, H, W].""" + signals = np.random.randn(BATCH_SIZE, 2, SEQ_LEN).astype(np.float32) + tensor = batch_to_hybrid_tensor(signals) + + expected = (BATCH_SIZE, 3, EXPECTED_H, EXPECTED_W) + assert tensor.shape == torch.Size(expected), \ + f"Expected {expected}, got {tuple(tensor.shape)}" + assert tensor.dtype == torch.float32, \ + f"Expected float32, got {tensor.dtype}" + print(f"[PASS] Hybrid batch shape: {tuple(tensor.shape)}") + + +def test_no_nan_or_inf(): + """Verify no NaN or Inf values in converted tensors.""" + signals = np.random.randn(BATCH_SIZE, 2, SEQ_LEN).astype(np.float32) + + gray = batch_to_grayscale_tensor(signals) + assert not torch.isnan(gray).any(), "NaN detected in grayscale tensor" + assert not torch.isinf(gray).any(), "Inf detected in grayscale tensor" + + hybrid = batch_to_hybrid_tensor(signals) + assert not torch.isnan(hybrid).any(), "NaN detected in hybrid tensor" + assert not torch.isinf(hybrid).any(), "Inf detected in hybrid tensor" + + print("[PASS] No NaN or Inf values detected in any tensor") + + +def test_stft_converter_consistency(): + """Verify our logic matches Person 3's stft_converter.compute_stft().""" + from stft_converter import compute_stft + + signal = np.random.randn(2, SEQ_LEN).astype(np.float32) + + _, _, p3_spec = compute_stft(signal) + our_spec = compute_spectrogram(signal) + + assert p3_spec.shape == our_spec.shape, \ + f"Shape mismatch: Person3={p3_spec.shape} vs ours={our_spec.shape}" + assert np.allclose(p3_spec, our_spec, atol=1e-5), \ + "Numerical mismatch with Person 3's stft_converter.compute_stft()" + print(f"[PASS] Consistency with stft_converter.py verified (shape={our_spec.shape})") + + +# ────────────────────────────────────────────────────────────── +# Runner +# ────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + print("=" * 58) + print(" CHUNK 1: 2D TENSOR SHAPE VALIDATION SUITE") + print("=" * 58) + + tests = [ + test_single_spectrogram_shape, + test_grayscale_batch_shape, + test_hybrid_batch_shape, + test_no_nan_or_inf, + test_stft_converter_consistency, + ] + + passed = 0 + failed = 0 + + for test_fn in tests: + try: + test_fn() + passed += 1 + except Exception as e: + print(f"[FAIL] {test_fn.__name__}: {e}") + failed += 1 + + print("\n" + "=" * 58) + print(f" RESULTS: {passed} passed, {failed} failed") + print("=" * 58) + + if failed > 0: + sys.exit(1) + + print("\n[SUCCESS] All 2D tensor shape validations passed!") + sys.exit(0) From fa96663b7ca907b8e7d626810438410d87b5ded1 Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sun, 23 Aug 2026 16:19:18 +0530 Subject: [PATCH 2/9] feat: implement baseline ResNet-18 architecture for 2d spectrograms --- models_2d.py | 189 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 189 insertions(+) create mode 100644 models_2d.py diff --git a/models_2d.py b/models_2d.py new file mode 100644 index 0000000..e46c8b0 --- /dev/null +++ b/models_2d.py @@ -0,0 +1,189 @@ +# models_2d.py +# ========================================================================= +# CHUNK 2: Baseline ResNet-18 for 2D RF Spectrograms +# Modified ResNet-18 architecture optimized for small (64x5) spectrograms +# generated by Person 3's Kaiser-windowed STFT converter. +# No torchvision dependency — built from scratch for full control. +# ========================================================================= + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +# ────────────────────────────────────────────────────────────── +# Building Block +# ────────────────────────────────────────────────────────────── + +class BasicBlock(nn.Module): + """Standard ResNet BasicBlock: two 3x3 conv layers + skip connection.""" + expansion = 1 + + def __init__(self, in_planes, planes, stride=1): + super().__init__() + self.conv1 = nn.Conv2d( + in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False + ) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d( + planes, planes, kernel_size=3, stride=1, padding=1, bias=False + ) + self.bn2 = nn.BatchNorm2d(planes) + + # Shortcut: identity when dims match, learned 1x1 projection otherwise + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != planes * self.expansion: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, planes * self.expansion, + kernel_size=1, stride=stride, bias=False + ), + nn.BatchNorm2d(planes * self.expansion), + ) + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + out += self.shortcut(x) + out = F.relu(out) + return out + + +# ────────────────────────────────────────────────────────────── +# ResNet-18 for 2D Spectrograms +# ────────────────────────────────────────────────────────────── + +class ResNet18_2D(nn.Module): + """ + Modified ResNet-18 for small 2D RF spectrograms. + + Adaptations from standard ImageNet ResNet-18: + 1. conv1: 3x3 kernel (not 7x7), stride=1 — preserves spatial + resolution on small (64×5) spectrograms. + 2. No initial MaxPool — prevents excessive downsampling. + 3. Configurable input channels: 1 (grayscale) or 3 (hybrid). + 4. Dynamic final FC layer sized to num_classes. + + Expected input shape: [B, in_channels, 64, 5] + Output shape: [B, num_classes] + + Args: + num_classes (int): Number of modulation classes. Default 11 (RadioML 2016.10a). + in_channels (int): Number of input channels. 1 for grayscale, 3 for hybrid. + """ + + def __init__(self, num_classes=11, in_channels=1): + super().__init__() + self.in_planes = 64 + + # ── Stem (no MaxPool) ────────────────────────────────── + self.conv1 = nn.Conv2d( + in_channels, 64, kernel_size=3, stride=1, padding=1, bias=False + ) + self.bn1 = nn.BatchNorm2d(64) + + # ── Residual Layers [2, 2, 2, 2] ────────────────────── + self.layer1 = self._make_layer(64, num_blocks=2, stride=1) + self.layer2 = self._make_layer(128, num_blocks=2, stride=2) + self.layer3 = self._make_layer(256, num_blocks=2, stride=2) + self.layer4 = self._make_layer(512, num_blocks=2, stride=2) + + # ── Classification Head ──────────────────────────────── + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(512 * BasicBlock.expansion, num_classes) + + # ── Weight Initialization (Kaiming) ──────────────────── + self._init_weights() + + def _make_layer(self, planes, num_blocks, stride): + """Build a residual layer with `num_blocks` BasicBlocks.""" + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for s in strides: + layers.append(BasicBlock(self.in_planes, planes, s)) + self.in_planes = planes * BasicBlock.expansion + return nn.Sequential(*layers) + + def _init_weights(self): + """Kaiming initialization for stable gradient flow.""" + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + nn.init.constant_(m.bias, 0) + + def forward(self, x): + """ + Forward pass through the modified ResNet-18. + + Feature map sizes for input [B, 1, 64, 5]: + conv1 → [B, 64, 64, 5] + layer1 → [B, 64, 64, 5] + layer2 → [B, 128, 32, 3] + layer3 → [B, 256, 16, 2] + layer4 → [B, 512, 8, 1] + avgpool→ [B, 512, 1, 1] + fc → [B, num_classes] + """ + # Stem + x = F.relu(self.bn1(self.conv1(x))) + + # Residual blocks + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + # Pool → Flatten → Classify + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + + return x + + +# ========================================================================= +# Smoke Test +# ========================================================================= +if __name__ == "__main__": + print("=" * 58) + print(" CHUNK 2: ResNet-18 2D ARCHITECTURE SMOKE TEST") + print("=" * 58) + + # ── Test 1: Grayscale input [B, 1, 64, 5] ───────────────── + model_gray = ResNet18_2D(num_classes=11, in_channels=1) + x_gray = torch.randn(8, 1, 64, 5) + out_gray = model_gray(x_gray) + + assert out_gray.shape == (8, 11), f"Expected (8, 11), got {out_gray.shape}" + print(f"[PASS] Grayscale input {tuple(x_gray.shape)} -> output {tuple(out_gray.shape)}") + + # ── Test 2: Hybrid input [B, 3, 64, 5] ──────────────────── + model_hybrid = ResNet18_2D(num_classes=11, in_channels=3) + x_hybrid = torch.randn(8, 3, 64, 5) + out_hybrid = model_hybrid(x_hybrid) + + assert out_hybrid.shape == (8, 11), f"Expected (8, 11), got {out_hybrid.shape}" + print(f"[PASS] Hybrid input {tuple(x_hybrid.shape)} -> output {tuple(out_hybrid.shape)}") + + # ── Test 3: Custom class count (MATLAB dataset = 5 classes) ─ + model_custom = ResNet18_2D(num_classes=5, in_channels=1) + x_custom = torch.randn(4, 1, 64, 5) + out_custom = model_custom(x_custom) + + assert out_custom.shape == (4, 5), f"Expected (4, 5), got {out_custom.shape}" + print(f"[PASS] Custom input {tuple(x_custom.shape)} -> output {tuple(out_custom.shape)} (5 classes)") + + # ── Parameter Count ──────────────────────────────────────── + total_params = sum(p.numel() for p in model_gray.parameters()) + trainable = sum(p.numel() for p in model_gray.parameters() if p.requires_grad) + print(f"\n[INFO] Total parameters : {total_params:,}") + print(f"[INFO] Trainable parameters: {trainable:,}") + + print("\n" + "=" * 58) + print(" [SUCCESS] ResNet-18 2D architecture verified!") + print("=" * 58) From 547f286ee7f35eb524e1cced2af9560c7e631abc Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sun, 23 Aug 2026 16:30:34 +0530 Subject: [PATCH 3/9] feat: build advanced STFT-RADN model integrating CBAM and residual dense blocks --- models_2d.py | 353 +++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 326 insertions(+), 27 deletions(-) diff --git a/models_2d.py b/models_2d.py index e46c8b0..df73b52 100644 --- a/models_2d.py +++ b/models_2d.py @@ -114,7 +114,8 @@ def _init_weights(self): nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) - nn.init.constant_(m.bias, 0) + if m.bias is not None: + nn.init.constant_(m.bias, 0) def forward(self, x): """ @@ -146,44 +147,342 @@ def forward(self, x): return x +# ══════════════════════════════════════════════════════════════ +# CHUNK 3: STFT-RADN (Residual Attention Dense Network) +# ══════════════════════════════════════════════════════════════ + + +# ────────────────────────────────────────────────────────────── +# CBAM: Convolutional Block Attention Module +# (Woo et al., ECCV 2018) +# ────────────────────────────────────────────────────────────── + +class ChannelAttention(nn.Module): + """ + Channel Attention sub-module of CBAM. + Learns inter-channel feature dependencies via shared MLP + applied to both global average-pooled and max-pooled descriptors. + """ + + def __init__(self, channels, reduction=16): + super().__init__() + mid = max(channels // reduction, 8) + self.mlp = nn.Sequential( + nn.Linear(channels, mid, bias=False), + nn.ReLU(inplace=True), + nn.Linear(mid, channels, bias=False), + ) + + def forward(self, x): + # x: [B, C, H, W] + avg_pool = x.mean(dim=[2, 3]) # [B, C] + max_pool = x.amax(dim=[2, 3]) # [B, C] + + scale = torch.sigmoid(self.mlp(avg_pool) + self.mlp(max_pool)) # [B, C] + return x * scale.unsqueeze(-1).unsqueeze(-1) + + +class SpatialAttention(nn.Module): + """ + Spatial Attention sub-module of CBAM. + Learns intra-spatial feature importance from channel-wise + average and max descriptors. + """ + + def __init__(self, kernel_size=7): + super().__init__() + self.conv = nn.Conv2d( + 2, 1, kernel_size, padding=kernel_size // 2, bias=False + ) + + def forward(self, x): + # x: [B, C, H, W] + avg_pool = x.mean(dim=1, keepdim=True) # [B, 1, H, W] + max_pool = x.amax(dim=1, keepdim=True) # [B, 1, H, W] + + scale = torch.sigmoid( + self.conv(torch.cat([avg_pool, max_pool], dim=1)) + ) # [B, 1, H, W] + return x * scale + + +class CBAM(nn.Module): + """ + Full CBAM block: Channel Attention followed by Spatial Attention. + + Args: + channels (int): Number of input feature channels. + reduction (int): Channel reduction ratio for the shared MLP. + spatial_kernel (int): Kernel size for the spatial attention conv. + """ + + def __init__(self, channels, reduction=16, spatial_kernel=7): + super().__init__() + self.channel_att = ChannelAttention(channels, reduction) + self.spatial_att = SpatialAttention(spatial_kernel) + + def forward(self, x): + x = self.channel_att(x) + x = self.spatial_att(x) + return x + + +# ────────────────────────────────────────────────────────────── +# Residual Dense Block (RDB) +# (Zhang et al., CVPR 2018 — adapted with CBAM) +# ────────────────────────────────────────────────────────────── + +class DenseLayer(nn.Module): + """ + Single dense layer: BN -> ReLU -> 3x3 Conv -> concatenate with input. + Each layer produces `growth_rate` new feature maps that are + concatenated onto the running feature tensor. + """ + + def __init__(self, in_channels, growth_rate): + super().__init__() + self.layer = nn.Sequential( + nn.BatchNorm2d(in_channels), + nn.ReLU(inplace=True), + nn.Conv2d(in_channels, growth_rate, 3, padding=1, bias=False), + ) + + def forward(self, x): + return torch.cat([x, self.layer(x)], dim=1) + + +class ResidualDenseBlock(nn.Module): + """ + Residual Dense Block with CBAM attention. + + Internal structure: + Input (C ch) -> DenseLayer x num_layers -> Local Feature Fusion (1x1) + -> CBAM -> + Input (local residual) + + Dense connections ensure maximum feature reuse; CBAM refines + the fused features before the residual addition. + + Args: + in_channels (int): Number of input channels (and output channels). + growth_rate (int): New feature maps produced per dense layer. + num_layers (int): Number of dense layers inside this block. + """ + + def __init__(self, in_channels=64, growth_rate=32, num_layers=4): + super().__init__() + + # Build dense layers with growing channel counts + layers = [] + ch = in_channels + for _ in range(num_layers): + layers.append(DenseLayer(ch, growth_rate)) + ch += growth_rate + self.dense_layers = nn.Sequential(*layers) + + # Local Feature Fusion: compress concatenated features back to in_channels + self.lff = nn.Sequential( + nn.Conv2d(ch, in_channels, 1, bias=False), + nn.BatchNorm2d(in_channels), + ) + + # Attention refinement + self.cbam = CBAM(in_channels) + + def forward(self, x): + out = self.dense_layers(x) # [B, C + L*G, H, W] + out = self.lff(out) # [B, C, H, W] + out = self.cbam(out) # [B, C, H, W] (attention-refined) + return out + x # Local residual learning + + +# ────────────────────────────────────────────────────────────── +# STFT-RADN: Full Architecture +# ────────────────────────────────────────────────────────────── + +class STFTRADN(nn.Module): + """ + STFT-based Residual Attention Dense Network. + + Integrates CBAM attention and Residual Dense Blocks for rich + spectral feature extraction from 2D STFT spectrograms. + + Architecture: + Shallow Feature Extraction (2-layer stem) + -> RDB+CBAM x num_rdb (dense feature reuse + attention) + -> Global Feature Fusion (1x1 conv on concatenated RDB outputs) + -> Global Residual (+ shallow features) + -> Classification Backbone (strided convs for spatial reduction) + -> FC + + Expected input: [B, in_channels, 64, 5] + Output: [B, num_classes] + + Args: + num_classes (int): Number of modulation classes. Default 11. + in_channels (int): Input channels. 1 (grayscale) or 3 (hybrid). + base_channels (int): Feature width through the RDB blocks. + growth_rate (int): New feature maps per dense layer. + num_dense_layers (int): Dense layers inside each RDB. + num_rdb (int): Number of RDB blocks. + """ + + def __init__(self, num_classes=11, in_channels=1, + base_channels=64, growth_rate=32, + num_dense_layers=4, num_rdb=3): + super().__init__() + self.num_rdb = num_rdb + + # ── Shallow Feature Extraction (2-layer stem) ───────── + self.sfe1 = nn.Conv2d( + in_channels, base_channels, 3, padding=1, bias=False + ) + self.sfe2 = nn.Sequential( + nn.BatchNorm2d(base_channels), + nn.ReLU(inplace=True), + nn.Conv2d(base_channels, base_channels, 3, padding=1, bias=False), + ) + + # ── Residual Dense Blocks with CBAM ─────────────────── + self.rdbs = nn.ModuleList([ + ResidualDenseBlock(base_channels, growth_rate, num_dense_layers) + for _ in range(num_rdb) + ]) + + # ── Global Feature Fusion ───────────────────────────── + self.gff = nn.Sequential( + nn.Conv2d(base_channels * num_rdb, base_channels, 1, bias=False), + nn.BatchNorm2d(base_channels), + nn.ReLU(inplace=True), + ) + + # ── Classification Backbone (downsample + deepen) ───── + self.backbone = nn.Sequential( + nn.Conv2d(base_channels, 128, 3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(128), + nn.ReLU(inplace=True), + nn.Conv2d(128, 256, 3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(256), + nn.ReLU(inplace=True), + nn.AdaptiveAvgPool2d((1, 1)), + ) + + # ── Classification Head ─────────────────────────────── + self.fc = nn.Linear(256, num_classes) + + # ── Weight Initialization ───────────────────────────── + self._init_weights() + + def _init_weights(self): + """Kaiming initialization for stable gradient flow.""" + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.normal_(m.weight, 0, 0.01) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x): + """ + Forward pass through STFT-RADN. + + Feature map sizes for input [B, 1, 64, 5]: + sfe1 -> [B, 64, 64, 5] (shallow features F0) + sfe2 -> [B, 64, 64, 5] (deeper stem F1) + RDB x3 -> [B, 64, 64, 5] each (local residual preserved) + GFF -> [B, 64, 64, 5] (fused from 3 x 64 = 192 ch) + + F0 -> [B, 64, 64, 5] (global residual) + backbone -> [B, 256, 1, 1] (spatial reduction + pooling) + fc -> [B, num_classes] + """ + # Shallow Feature Extraction + f0 = self.sfe1(x) + f1 = self.sfe2(f0) + + # Pass through RDB blocks, collecting each output + rdb_outputs = [] + h = f1 + for rdb in self.rdbs: + h = rdb(h) + rdb_outputs.append(h) + + # Global Feature Fusion + gff_out = self.gff(torch.cat(rdb_outputs, dim=1)) + + # Global Residual Learning + out = gff_out + f0 + + # Classification + out = self.backbone(out) + out = torch.flatten(out, 1) + out = self.fc(out) + + return out + + # ========================================================================= -# Smoke Test +# Smoke Test (both architectures) # ========================================================================= if __name__ == "__main__": + + def count_params(model): + total = sum(p.numel() for p in model.parameters()) + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + return total, trainable + + x_gray = torch.randn(8, 1, 64, 5) + x_hybrid = torch.randn(8, 3, 64, 5) + x_custom = torch.randn(4, 1, 64, 5) + + # ── ResNet-18 Tests ──────────────────────────────────────── print("=" * 58) - print(" CHUNK 2: ResNet-18 2D ARCHITECTURE SMOKE TEST") + print(" ResNet-18 2D SMOKE TEST") print("=" * 58) - # ── Test 1: Grayscale input [B, 1, 64, 5] ───────────────── - model_gray = ResNet18_2D(num_classes=11, in_channels=1) - x_gray = torch.randn(8, 1, 64, 5) - out_gray = model_gray(x_gray) + r1 = ResNet18_2D(num_classes=11, in_channels=1) + out = r1(x_gray) + assert out.shape == (8, 11) + print(f"[PASS] Grayscale {tuple(x_gray.shape)} -> {tuple(out.shape)}") - assert out_gray.shape == (8, 11), f"Expected (8, 11), got {out_gray.shape}" - print(f"[PASS] Grayscale input {tuple(x_gray.shape)} -> output {tuple(out_gray.shape)}") + r2 = ResNet18_2D(num_classes=11, in_channels=3) + out = r2(x_hybrid) + assert out.shape == (8, 11) + print(f"[PASS] Hybrid {tuple(x_hybrid.shape)} -> {tuple(out.shape)}") - # ── Test 2: Hybrid input [B, 3, 64, 5] ──────────────────── - model_hybrid = ResNet18_2D(num_classes=11, in_channels=3) - x_hybrid = torch.randn(8, 3, 64, 5) - out_hybrid = model_hybrid(x_hybrid) + r3 = ResNet18_2D(num_classes=5, in_channels=1) + out = r3(x_custom) + assert out.shape == (4, 5) + print(f"[PASS] Custom {tuple(x_custom.shape)} -> {tuple(out.shape)} (5 cls)") - assert out_hybrid.shape == (8, 11), f"Expected (8, 11), got {out_hybrid.shape}" - print(f"[PASS] Hybrid input {tuple(x_hybrid.shape)} -> output {tuple(out_hybrid.shape)}") + t, tr = count_params(r1) + print(f"[INFO] Params: {t:,} total, {tr:,} trainable") - # ── Test 3: Custom class count (MATLAB dataset = 5 classes) ─ - model_custom = ResNet18_2D(num_classes=5, in_channels=1) - x_custom = torch.randn(4, 1, 64, 5) - out_custom = model_custom(x_custom) + # ── STFT-RADN Tests ─────────────────────────────────────── + print("\n" + "=" * 58) + print(" STFT-RADN SMOKE TEST") + print("=" * 58) + + s1 = STFTRADN(num_classes=11, in_channels=1) + out = s1(x_gray) + assert out.shape == (8, 11) + print(f"[PASS] Grayscale {tuple(x_gray.shape)} -> {tuple(out.shape)}") + + s2 = STFTRADN(num_classes=11, in_channels=3) + out = s2(x_hybrid) + assert out.shape == (8, 11) + print(f"[PASS] Hybrid {tuple(x_hybrid.shape)} -> {tuple(out.shape)}") - assert out_custom.shape == (4, 5), f"Expected (4, 5), got {out_custom.shape}" - print(f"[PASS] Custom input {tuple(x_custom.shape)} -> output {tuple(out_custom.shape)} (5 classes)") + s3 = STFTRADN(num_classes=5, in_channels=1) + out = s3(x_custom) + assert out.shape == (4, 5) + print(f"[PASS] Custom {tuple(x_custom.shape)} -> {tuple(out.shape)} (5 cls)") - # ── Parameter Count ──────────────────────────────────────── - total_params = sum(p.numel() for p in model_gray.parameters()) - trainable = sum(p.numel() for p in model_gray.parameters() if p.requires_grad) - print(f"\n[INFO] Total parameters : {total_params:,}") - print(f"[INFO] Trainable parameters: {trainable:,}") + t, tr = count_params(s1) + print(f"[INFO] Params: {t:,} total, {tr:,} trainable") print("\n" + "=" * 58) - print(" [SUCCESS] ResNet-18 2D architecture verified!") + print(" [SUCCESS] Both architectures verified!") print("=" * 58) From c404bf36b8f8d0aa1ce4a0a0681ff9a4a7568717 Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sun, 23 Aug 2026 17:21:17 +0530 Subject: [PATCH 4/9] feat: implement 2d training loop with adamw and validation early stopping --- train_2d.py | 378 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 378 insertions(+) create mode 100644 train_2d.py diff --git a/train_2d.py b/train_2d.py new file mode 100644 index 0000000..ae7a364 --- /dev/null +++ b/train_2d.py @@ -0,0 +1,378 @@ +# train_2d.py +# ========================================================================= +# CHUNK 4: 2D Training Loop +# Comprehensive training orchestrator for ResNet-18 and STFT-RADN. +# AdamW optimizer, CrossEntropy loss, ReduceLROnPlateau scheduler, +# early stopping on validation loss, and per-SNR accuracy logging +# to match Team A's benchmarking format. +# ========================================================================= + +import os +import sys +import time +import warnings +import argparse +import numpy as np +import torch +import torch.nn as nn +from torch.optim import AdamW +from torch.optim.lr_scheduler import ReduceLROnPlateau +from scipy.signal import stft as scipy_stft + +from dataset_loader import get_rf_dataloader +from models_2d import ResNet18_2D, STFTRADN + +# Suppress scipy's harmless "complex input" warning during STFT +warnings.filterwarnings("ignore", message=".*complex.*return_onesided.*") + + +# ────────────────────────────────────────────────────────────── +# Transform: 1D IQ -> 2D Spectrogram (on-the-fly) +# ────────────────────────────────────────────────────────────── + +class IQToSpectrogram: + """ + On-the-fly transform bridging Person 1's DataLoader and Person 4's models. + Converts a (2, 128) IQ tensor to a 2D spectrogram tensor using + Person 3's exact Kaiser-windowed STFT parameters. + + Args: + mode: "grayscale" -> (1, 64, 5) single power spectrogram + "hybrid" -> (3, 64, 5) [power_dB, phase, magnitude] + """ + STFT_PARAMS = dict(window=("kaiser", 0.85), nperseg=64, noverlap=32) + + def __init__(self, mode="grayscale"): + self.mode = mode + + def __call__(self, signal_tensor): + signal_np = signal_tensor.numpy() + I, Q = signal_np[0], signal_np[1] + complex_signal = I + 1j * Q + + _, _, Zxx = scipy_stft(complex_signal, **self.STFT_PARAMS) + + if self.mode == "grayscale": + power_db = 10 * np.log10(np.abs(Zxx) ** 2 + 1e-10) + spec = np.fft.fftshift(power_db, axes=0) + return torch.tensor(spec, dtype=torch.float32).unsqueeze(0) # (1, H, W) + + else: # hybrid + Zxx_s = np.fft.fftshift(Zxx, axes=0) + power_db = 10 * np.log10(np.abs(Zxx_s) ** 2 + 1e-10) + phase = np.angle(Zxx_s) + magnitude = np.abs(Zxx_s) + return torch.tensor( + np.stack([power_db, phase, magnitude], axis=0), + dtype=torch.float32, + ) # (3, H, W) + + +# ────────────────────────────────────────────────────────────── +# Training & Validation +# ────────────────────────────────────────────────────────────── + +def train_one_epoch(model, loader, criterion, optimizer, device): + """Run one training epoch. Returns (avg_loss, accuracy).""" + model.train() + running_loss = 0.0 + correct = 0 + total = 0 + + for signals, labels, _snrs in loader: + signals = signals.to(device) + labels = labels.to(device) + + optimizer.zero_grad() + outputs = model(signals) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() * signals.size(0) + _, predicted = outputs.max(1) + correct += predicted.eq(labels).sum().item() + total += labels.size(0) + + return running_loss / total, 100.0 * correct / total + + +@torch.no_grad() +def validate(model, loader, criterion, device): + """Run validation pass. Returns (avg_loss, accuracy).""" + model.eval() + running_loss = 0.0 + correct = 0 + total = 0 + + for signals, labels, _snrs in loader: + signals = signals.to(device) + labels = labels.to(device) + + outputs = model(signals) + loss = criterion(outputs, labels) + + running_loss += loss.item() * signals.size(0) + _, predicted = outputs.max(1) + correct += predicted.eq(labels).sum().item() + total += labels.size(0) + + return running_loss / total, 100.0 * correct / total + + +@torch.no_grad() +def evaluate_per_snr(model, loader, device): + """ + Per-SNR accuracy evaluation — matches Team A's benchmarking output format. + Returns: dict {snr_dB: accuracy_%} + """ + model.eval() + snr_correct = {} + snr_total = {} + + for signals, labels, snrs in loader: + signals = signals.to(device) + labels = labels.to(device) + + outputs = model(signals) + _, predicted = outputs.max(1) + + for i in range(labels.size(0)): + snr = int(snrs[i].item()) + snr_correct.setdefault(snr, 0) + snr_total.setdefault(snr, 0) + snr_total[snr] += 1 + if predicted[i] == labels[i]: + snr_correct[snr] += 1 + + return { + snr: 100.0 * snr_correct[snr] / snr_total[snr] + for snr in sorted(snr_correct) + } + + +# ────────────────────────────────────────────────────────────── +# Early Stopping +# ────────────────────────────────────────────────────────────── + +class EarlyStopping: + """ + Stops training when validation loss has not improved for + `patience` consecutive epochs. + """ + + def __init__(self, patience=10, min_delta=1e-4): + self.patience = patience + self.min_delta = min_delta + self.best_loss = float("inf") + self.counter = 0 + self.should_stop = False + + def step(self, val_loss): + if val_loss < self.best_loss - self.min_delta: + self.best_loss = val_loss + self.counter = 0 + else: + self.counter += 1 + if self.counter >= self.patience: + self.should_stop = True + + +# ────────────────────────────────────────────────────────────── +# Main Training Orchestrator +# ────────────────────────────────────────────────────────────── + +def main(args): + # ── Device ───────────────────────────────────────────────── + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + print("=" * 70) + print(" 2D MODEL TRAINING LOOP") + print("=" * 70) + print(f" Device : {device}") + print(f" Model : {args.model}") + print(f" Input mode : {args.input_mode}") + print(f" LR : {args.lr} | Weight decay: {args.weight_decay}") + print(f" Batch size : {args.batch_size} | Max epochs: {args.epochs}") + print(f" Patience : {args.patience}") + print("=" * 70) + + # ── Transform ────────────────────────────────────────────── + transform = IQToSpectrogram(mode=args.input_mode) + in_channels = 1 if args.input_mode == "grayscale" else 3 + + # ── DataLoaders ──────────────────────────────────────────── + print("\n--- Loading dataset splits ---") + train_loader = get_rf_dataloader( + data_dir=args.data_dir, split="train", + batch_size=args.batch_size, transform=transform, + ) + val_loader = get_rf_dataloader( + data_dir=args.data_dir, split="val", + batch_size=args.batch_size, transform=transform, + ) + test_loader = get_rf_dataloader( + data_dir=args.data_dir, split="test", + batch_size=args.batch_size, transform=transform, + ) + print(f" Train : {len(train_loader.dataset):>6} samples ({len(train_loader)} batches)") + print(f" Val : {len(val_loader.dataset):>6} samples ({len(val_loader)} batches)") + print(f" Test : {len(test_loader.dataset):>6} samples ({len(test_loader)} batches)") + + # ── Model ────────────────────────────────────────────────── + num_classes = len(np.unique( + np.load(os.path.join(args.data_dir, "y_all.npy")) + )) + print(f" Classes: {num_classes}") + + if args.model == "resnet18": + model = ResNet18_2D(num_classes=num_classes, in_channels=in_channels) + elif args.model == "stft-radn": + model = STFTRADN(num_classes=num_classes, in_channels=in_channels) + else: + sys.exit(f"[ERROR] Unknown model: {args.model}") + + model = model.to(device) + total_params = sum(p.numel() for p in model.parameters()) + print(f" Parameters: {total_params:,}") + + # ── Optimizer / Loss / Scheduler ─────────────────────────── + optimizer = AdamW( + model.parameters(), lr=args.lr, weight_decay=args.weight_decay + ) + criterion = nn.CrossEntropyLoss() + scheduler = ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=5 + ) + early_stop = EarlyStopping(patience=args.patience) + + # ── Checkpoint dir ───────────────────────────────────────── + os.makedirs(args.save_dir, exist_ok=True) + best_val_acc = 0.0 + best_model_path = os.path.join(args.save_dir, f"best_{args.model}.pt") + + # ── Training Loop ────────────────────────────────────────── + print("\n" + "=" * 74) + print(f"{'Ep':>4} | {'Train Loss':>10} | {'Train Acc':>9} | " + f"{'Val Loss':>9} | {'Val Acc':>8} | {'LR':>10} | {'Time':>6}") + print("-" * 74) + + for epoch in range(1, args.epochs + 1): + t0 = time.time() + + train_loss, train_acc = train_one_epoch( + model, train_loader, criterion, optimizer, device + ) + val_loss, val_acc = validate( + model, val_loader, criterion, device + ) + + current_lr = optimizer.param_groups[0]["lr"] + scheduler.step(val_loss) + elapsed = time.time() - t0 + + marker = " *" if val_acc > best_val_acc else "" + print( + f"{epoch:4d} | {train_loss:10.4f} | {train_acc:8.2f}% | " + f"{val_loss:9.4f} | {val_acc:7.2f}% | {current_lr:10.6f} | " + f"{elapsed:5.1f}s{marker}" + ) + + # Save best model + if val_acc > best_val_acc: + best_val_acc = val_acc + torch.save({ + "epoch": epoch, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "val_acc": val_acc, + "val_loss": val_loss, + "model_name": args.model, + "num_classes": num_classes, + "in_channels": in_channels, + }, best_model_path) + + # Early stopping check + early_stop.step(val_loss) + if early_stop.should_stop: + print(f"\n[EARLY STOP] No val loss improvement for {args.patience} epochs.") + break + + print("=" * 74) + + # ── Final Test Evaluation ────────────────────────────────── + print("\n--- Loading best checkpoint for final evaluation ---") + checkpoint = torch.load(best_model_path, map_location=device, weights_only=True) + model.load_state_dict(checkpoint["model_state_dict"]) + + test_loss, test_acc = validate(model, test_loader, criterion, device) + print(f"\n [TEST RESULT] Loss: {test_loss:.4f} | Accuracy: {test_acc:.2f}%") + + # ── Per-SNR Accuracy (Team A compatible format) ──────────── + print("\n--- Per-SNR Accuracy (Team A Benchmarking Format) ---") + snr_acc = evaluate_per_snr(model, test_loader, device) + + print(f"{'SNR (dB)':>10} | {'Accuracy':>10} | {'Bar'}") + print("-" * 50) + for snr, acc in snr_acc.items(): + bar = "#" * int(acc / 5) + print(f"{snr:>8d} dB | {acc:9.2f}% | {bar}") + + avg_snr_acc = np.mean(list(snr_acc.values())) + print("-" * 50) + print(f"{'Average':>10} | {avg_snr_acc:9.2f}%") + + print(f"\n[SAVED] Best checkpoint : {best_model_path}") + print(f"[SAVED] Best val acc : {best_val_acc:.2f}%") + print(f"[DONE] Test accuracy : {test_acc:.2f}%") + + +# ────────────────────────────────────────────────────────────── +# CLI Entry Point +# ────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="2D Model Training for RF Signal Classification" + ) + parser.add_argument( + "--model", type=str, default="resnet18", + choices=["resnet18", "stft-radn"], + help="Model architecture (default: resnet18)", + ) + parser.add_argument( + "--input-mode", type=str, default="grayscale", + choices=["grayscale", "hybrid"], + help="Spectrogram mode: grayscale (1ch) or hybrid (3ch)", + ) + parser.add_argument( + "--data-dir", type=str, default="./data", + help="Path to preprocessed dataset directory", + ) + parser.add_argument( + "--save-dir", type=str, default="./checkpoints", + help="Directory to save model checkpoints", + ) + parser.add_argument( + "--epochs", type=int, default=50, + help="Maximum training epochs (default: 50)", + ) + parser.add_argument( + "--batch-size", type=int, default=64, + help="Training batch size (default: 64)", + ) + parser.add_argument( + "--lr", type=float, default=1e-3, + help="AdamW learning rate (default: 1e-3)", + ) + parser.add_argument( + "--weight-decay", type=float, default=1e-4, + help="AdamW weight decay (default: 1e-4)", + ) + parser.add_argument( + "--patience", type=int, default=10, + help="Early stopping patience in epochs (default: 10)", + ) + + args = parser.parse_args() + main(args) From 26f9e3c633e79289ce2994d15f52ba8eef83d5e6 Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sun, 23 Aug 2026 17:37:50 +0530 Subject: [PATCH 5/9] feat: add latency/parameter benchmarking and ONNX export for Grad-CAM --- .gitignore | 3 +- benchmark_2d.py | 81 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 benchmark_2d.py diff --git a/.gitignore b/.gitignore index 86119c6..7fc2661 100644 --- a/.gitignore +++ b/.gitignore @@ -25,4 +25,5 @@ data/ .DS_Store Thumbs.db -stft_dataset.npz \ No newline at end of file +stft_dataset.npz +*.onnx \ No newline at end of file diff --git a/benchmark_2d.py b/benchmark_2d.py new file mode 100644 index 0000000..fdf1a73 --- /dev/null +++ b/benchmark_2d.py @@ -0,0 +1,81 @@ +import torch +import time +from models_2d import ResNet18_2D, STFTRADN + +def measure_latency(model, input_shape, num_runs=100, warmup_runs=10): + """Measures the inference latency per sample in milliseconds.""" + device = next(model.parameters()).device + dummy_input = torch.randn(input_shape).to(device) + + model.eval() + + # Warmup + with torch.no_grad(): + for _ in range(warmup_runs): + _ = model(dummy_input) + + # Measure + start_time = time.perf_counter() + with torch.no_grad(): + for _ in range(num_runs): + _ = model(dummy_input) + end_time = time.perf_counter() + + total_time_ms = (end_time - start_time) * 1000 + latency_per_sample_ms = total_time_ms / (num_runs * input_shape[0]) + return latency_per_sample_ms + +def count_trainable_parameters(model): + """Calculates the total trainable parameter footprint.""" + return sum(p.numel() for p in model.parameters() if p.requires_grad) + +def export_to_onnx(model, input_shape, filename="stft_radn_model.onnx"): + """Exports the model weights to an ONNX file.""" + device = next(model.parameters()).device + dummy_input = torch.randn(input_shape).to(device) + model.eval() + + torch.onnx.export( + model, + dummy_input, + filename, + export_params=True, + opset_version=12, + do_constant_folding=True, + input_names=['input'], + output_names=['output'], + dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} + ) + print(f"[INFO] Model exported successfully to {filename}") + +if __name__ == "__main__": + # Standard input shape: [Batch, Channels, Height, Width] + # Corresponding to [1, 1, 64, 5] for grayscale 2D spectrograms + input_shape = (1, 1, 64, 5) + + print("="*60) + print("Initializing Models...") + resnet_model = ResNet18_2D(num_classes=11, in_channels=1) + stft_model = STFTRADN(num_classes=11, in_channels=1) + + print("="*60) + print(" ResNet18_2D Benchmarks") + print("="*60) + resnet_params = count_trainable_parameters(resnet_model) + print(f"Trainable Parameters : {resnet_params:,}") + resnet_latency = measure_latency(resnet_model, input_shape) + print(f"Inference Latency : {resnet_latency:.4f} ms / sample") + + print("\n" + "="*60) + print(" STFTRADN Benchmarks") + print("="*60) + stft_params = count_trainable_parameters(stft_model) + print(f"Trainable Parameters : {stft_params:,}") + stft_latency = measure_latency(stft_model, input_shape) + print(f"Inference Latency : {stft_latency:.4f} ms / sample") + + print("\n" + "="*60) + print(" ONNX Export (STFTRADN for Grad-CAM validation)") + print("="*60) + export_to_onnx(stft_model, input_shape, filename="stft_radn_model.onnx") + print("="*60) From 0f527fcfd1a8d48de7a71416c0b84d22e6d1ca1d Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sat, 29 Aug 2026 11:27:25 +0530 Subject: [PATCH 6/9] Optimize 2D training pipeline to mitigate overfitting - Implemented lightweight channel progression in ResNet18_2D (16->32->64->128) reducing params to ~701K - Added nn.Dropout2d(p=0.3) inside ResNet BasicBlock - Integrated torchvision RandomErasing (SpecAugment) dynamically in SpectrogramDataset - Increased weight_decay to 1e-2 in training config - Updated colab_train.ipynb to mount Google Drive automatically for faster dataset loading --- colab_train.ipynb | 468 +++++++++++++++++++++++++++++++++++++++++++++ dataset_loader.py | 130 +++++++++++++ health_report.md | 48 ----- models_2d.py | 30 +-- precompute_stft.py | 128 +++++++++++++ train_2d.py | 109 ++++------- 6 files changed, 778 insertions(+), 135 deletions(-) create mode 100644 colab_train.ipynb delete mode 100644 health_report.md create mode 100644 precompute_stft.py diff --git a/colab_train.ipynb b/colab_train.ipynb new file mode 100644 index 0000000..5b7a475 --- /dev/null +++ b/colab_train.ipynb @@ -0,0 +1,468 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# RF Signal Classifier \u00e2\u20ac\u201d 2D Training (ResNet-18 + STFT-RADN)\n", + "**Run on GPU: Runtime \u00e2\u2020\u2019 Change runtime type \u00e2\u2020\u2019 T4 GPU**\n", + "\n", + "### Steps\n", + "1. Run Cell 1 \u00e2\u20ac\u201d check GPU\n", + "2. Run Cell 2 \u00e2\u20ac\u201d upload your 6 `.npy` data files\n", + "3. Run remaining cells in order" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Cell 1: GPU Check \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n", + "import torch\n", + "print('PyTorch:', torch.__version__)\n", + "print('GPU available:', torch.cuda.is_available())\n", + "if torch.cuda.is_available():\n", + " print('GPU name:', torch.cuda.get_device_name(0))\n", + "else:\n", + " print('WARNING: No GPU detected. Go to Runtime > Change runtime type > T4 GPU')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u2500\u2500 Cell 2: Mount Google Drive & Copy Data \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", + "# Since the 6 .npy files were uploaded during a previous run,\n", + "# we will just mount your Google Drive and copy them directly.\n", + "import os\n", + "import shutil\n", + "from google.colab import drive\n", + "\n", + "drive.mount('/content/drive')\n", + "\n", + "os.makedirs('/content/data', exist_ok=True)\n", + "os.makedirs('/content/checkpoints', exist_ok=True)\n", + "\n", + "# UPDATE THIS PATH if your data is saved in a different folder in your Drive\n", + "drive_data_path = '/content/drive/MyDrive/rf_data/rf_data' \n", + "\n", + "files_to_copy = ['spectrograms_all.npy', 'y_all.npy', 'snrs_all.npy', 'train_idx.npy', 'val_idx.npy', 'test_idx.npy']\n", + "print('Copying files from Drive to local Colab storage for speed...')\n", + "for fname in files_to_copy:\n", + " src = os.path.join(drive_data_path, fname)\n", + " dst = os.path.join('/content/data', fname)\n", + " if os.path.exists(src):\n", + " shutil.copy(src, dst)\n", + " size_mb = os.path.getsize(dst) / 1e6\n", + " print(f' Copied: {fname} ({size_mb:.1f} MB)')\n", + " else:\n", + " print(f' [ERROR] Could not find {fname} in {drive_data_path}')\n", + "\n", + "print('\\nFiles ready in /content/data:')\n", + "for f in sorted(os.listdir('/content/data')):\n", + " print(f' {f}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Cell 3: Model Definitions (ResNet18_2D + STFT-RADN) \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "\n", + "class BasicBlock(nn.Module):\n", + " expansion = 1\n", + " def __init__(self, in_planes, planes, stride=1):\n", + " super().__init__()\n", + " self.conv1 = nn.Conv2d(in_planes, planes, 3, stride=stride, padding=1, bias=False)\n", + " self.bn1 = nn.BatchNorm2d(planes)\n", + " self.dropout = nn.Dropout2d(p=0.3)\n", + " self.conv2 = nn.Conv2d(planes, planes, 3, stride=1, padding=1, bias=False)\n", + " self.bn2 = nn.BatchNorm2d(planes)\n", + " self.shortcut = nn.Sequential()\n", + " if stride != 1 or in_planes != planes * self.expansion:\n", + " self.shortcut = nn.Sequential(\n", + " nn.Conv2d(in_planes, planes * self.expansion, 1, stride=stride, bias=False),\n", + " nn.BatchNorm2d(planes * self.expansion),\n", + " )\n", + " def forward(self, x):\n", + " out = F.relu(self.bn1(self.conv1(x)))\n", + " out = self.dropout(out)\n", + " out = self.bn2(self.conv2(out))\n", + " out += self.shortcut(x)\n", + " return F.relu(out)\n", + "\n", + "class ResNet18_2D(nn.Module):\n", + " def __init__(self, num_classes=11, in_channels=1):\n", + " super().__init__()\n", + " self.in_planes = 16\n", + " self.conv1 = nn.Conv2d(in_channels, 16, 3, stride=1, padding=1, bias=False)\n", + " self.bn1 = nn.BatchNorm2d(16)\n", + " self.layer1 = self._make_layer(16, 2, stride=1)\n", + " self.layer2 = self._make_layer(32, 2, stride=2)\n", + " self.layer3 = self._make_layer(64, 2, stride=2)\n", + " self.layer4 = self._make_layer(128, 2, stride=2)\n", + " self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n", + " self.fc = nn.Linear(128, num_classes)\n", + " for m in self.modules():\n", + " if isinstance(m, nn.Conv2d):\n", + " nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n", + " elif isinstance(m, nn.BatchNorm2d):\n", + " nn.init.constant_(m.weight, 1); nn.init.constant_(m.bias, 0)\n", + " elif isinstance(m, nn.Linear):\n", + " nn.init.normal_(m.weight, 0, 0.01)\n", + " if m.bias is not None: nn.init.constant_(m.bias, 0)\n", + " def _make_layer(self, planes, num_blocks, stride):\n", + " strides = [stride] + [1] * (num_blocks - 1)\n", + " layers = []\n", + " for s in strides:\n", + " layers.append(BasicBlock(self.in_planes, planes, s))\n", + " self.in_planes = planes\n", + " return nn.Sequential(*layers)\n", + " def forward(self, x):\n", + " x = F.relu(self.bn1(self.conv1(x)))\n", + " x = self.layer1(x); x = self.layer2(x)\n", + " x = self.layer3(x); x = self.layer4(x)\n", + " x = self.avgpool(x)\n", + " return self.fc(torch.flatten(x, 1))\n", + "\n", + "class ChannelAttention(nn.Module):\n", + " def __init__(self, channels, reduction=16):\n", + " super().__init__()\n", + " mid = max(channels // reduction, 8)\n", + " self.mlp = nn.Sequential(nn.Linear(channels, mid, bias=False), nn.ReLU(inplace=True), nn.Linear(mid, channels, bias=False))\n", + " def forward(self, x):\n", + " scale = torch.sigmoid(self.mlp(x.mean(dim=[2,3])) + self.mlp(x.amax(dim=[2,3])))\n", + " return x * scale.unsqueeze(-1).unsqueeze(-1)\n", + "\n", + "class SpatialAttention(nn.Module):\n", + " def __init__(self, kernel_size=7):\n", + " super().__init__()\n", + " self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)\n", + " def forward(self, x):\n", + " scale = torch.sigmoid(self.conv(torch.cat([x.mean(1,keepdim=True), x.amax(1,keepdim=True)], dim=1)))\n", + " return x * scale\n", + "\n", + "class CBAM(nn.Module):\n", + " def __init__(self, channels, reduction=16, spatial_kernel=7):\n", + " super().__init__()\n", + " self.channel_att = ChannelAttention(channels, reduction)\n", + " self.spatial_att = SpatialAttention(spatial_kernel)\n", + " def forward(self, x):\n", + " return self.spatial_att(self.channel_att(x))\n", + "\n", + "class DenseLayer(nn.Module):\n", + " def __init__(self, in_channels, growth_rate):\n", + " super().__init__()\n", + " self.layer = nn.Sequential(nn.BatchNorm2d(in_channels), nn.ReLU(inplace=True), nn.Conv2d(in_channels, growth_rate, 3, padding=1, bias=False))\n", + " def forward(self, x):\n", + " return torch.cat([x, self.layer(x)], dim=1)\n", + "\n", + "class ResidualDenseBlock(nn.Module):\n", + " def __init__(self, in_channels=64, growth_rate=32, num_layers=4):\n", + " super().__init__()\n", + " layers, ch = [], in_channels\n", + " for _ in range(num_layers):\n", + " layers.append(DenseLayer(ch, growth_rate)); ch += growth_rate\n", + " self.dense_layers = nn.Sequential(*layers)\n", + " self.lff = nn.Sequential(nn.Conv2d(ch, in_channels, 1, bias=False), nn.BatchNorm2d(in_channels))\n", + " self.cbam = CBAM(in_channels)\n", + " def forward(self, x):\n", + " return self.cbam(self.lff(self.dense_layers(x))) + x\n", + "\n", + "class STFTRADN(nn.Module):\n", + " def __init__(self, num_classes=11, in_channels=1, base_channels=64, growth_rate=32, num_dense_layers=4, num_rdb=3):\n", + " super().__init__()\n", + " self.num_rdb = num_rdb\n", + " self.sfe1 = nn.Conv2d(in_channels, base_channels, 3, padding=1, bias=False)\n", + " self.sfe2 = nn.Sequential(nn.BatchNorm2d(base_channels), nn.ReLU(inplace=True), nn.Conv2d(base_channels, base_channels, 3, padding=1, bias=False))\n", + " self.rdbs = nn.ModuleList([ResidualDenseBlock(base_channels, growth_rate, num_dense_layers) for _ in range(num_rdb)])\n", + " self.gff = nn.Sequential(nn.Conv2d(base_channels * num_rdb, base_channels, 1, bias=False), nn.BatchNorm2d(base_channels), nn.ReLU(inplace=True))\n", + " self.backbone = nn.Sequential(\n", + " nn.Conv2d(base_channels, 128, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(128), nn.ReLU(inplace=True),\n", + " nn.Conv2d(128, 256, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(256), nn.ReLU(inplace=True),\n", + " nn.AdaptiveAvgPool2d((1, 1)))\n", + " self.fc = nn.Linear(256, num_classes)\n", + " for m in self.modules():\n", + " if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n", + " elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1); nn.init.constant_(m.bias, 0)\n", + " elif isinstance(m, nn.Linear):\n", + " nn.init.normal_(m.weight, 0, 0.01)\n", + " if m.bias is not None: nn.init.constant_(m.bias, 0)\n", + " def forward(self, x):\n", + " f0 = self.sfe1(x); f1 = self.sfe2(f0)\n", + " rdb_out, h = [], f1\n", + " for rdb in self.rdbs:\n", + " h = rdb(h); rdb_out.append(h)\n", + " out = self.gff(torch.cat(rdb_out, dim=1)) + f0\n", + " return self.fc(torch.flatten(self.backbone(out), 1))\n", + "\n", + "print('[OK] Models defined: ResNet18_2D, STFTRADN')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Cell 4: Dataset & DataLoader \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n", + "import numpy as np\n", + "from torch.utils.data import Dataset, DataLoader\n", + "\n", + "class SpectrogramDataset(Dataset):\n", + " \"\"\"Loads pre-computed STFT spectrograms via memory-mapping.\"\"\"\n", + " def __init__(self, data_dir='/content/data', split='train', in_channels=1):\n", + " self.split = split\n", + " self.in_channels = in_channels\n", + " self._spectrograms = np.load(f'{data_dir}/spectrograms_all.npy', mmap_mode='r')\n", + " self._labels = np.load(f'{data_dir}/y_all.npy')\n", + " self._snrs = np.load(f'{data_dir}/snrs_all.npy')\n", + " self._indices = np.load(f'{data_dir}/{split}_idx.npy')\n", + " import torchvision.transforms as T\n", + " self.transform = T.RandomErasing(p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0) if split == 'train' else None\n", + " def __len__(self):\n", + " return len(self._indices)\n", + " def __getitem__(self, idx):\n", + " i = self._indices[idx]\n", + " spec = np.array(self._spectrograms[i], dtype=np.float32)\n", + " if self.in_channels == 1:\n", + " spec = spec[:1]\n", + " signal_tensor = torch.from_numpy(spec)\n", + " if self.transform is not None:\n", + " signal_tensor = self.transform(signal_tensor)\n", + " return signal_tensor, torch.tensor(int(self._labels[i]), dtype=torch.long), self._snrs[i]\n", + "\n", + "def get_loader(split, batch_size, in_channels, num_workers=2, data_dir='/content/data'):\n", + " ds = SpectrogramDataset(data_dir, split, in_channels)\n", + " return DataLoader(ds, batch_size=batch_size, shuffle=(split=='train'),\n", + " num_workers=num_workers, pin_memory=True, persistent_workers=(num_workers>0))\n", + "\n", + "print('[OK] Dataset classes defined')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Cell 5: Training Config \u00e2\u20ac\u201d EDIT HERE \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n", + "CONFIG = dict(\n", + " model = 'resnet18', # 'resnet18' or 'stft-radn'\n", + " input_mode = 'grayscale', # 'grayscale' (1ch) or 'hybrid' (3ch)\n", + " batch_size = 256,\n", + " epochs = 30,\n", + " lr = 1e-3,\n", + " weight_decay = 1e-2,\n", + " patience = 10, # early stopping\n", + " num_workers = 2,\n", + " data_dir = '/content/data',\n", + " save_dir = '/content/checkpoints',\n", + ")\n", + "print('Config:', CONFIG)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Cell 6: Build Model + Loaders \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n", + "import os, time\n", + "from torch.optim import AdamW\n", + "from torch.optim.lr_scheduler import ReduceLROnPlateau\n", + "\n", + "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", + "in_ch = 1 if CONFIG['input_mode'] == 'grayscale' else 3\n", + "os.makedirs(CONFIG['save_dir'], exist_ok=True)\n", + "\n", + "# Loaders\n", + "kw = dict(batch_size=CONFIG['batch_size'], in_channels=in_ch,\n", + " num_workers=CONFIG['num_workers'], data_dir=CONFIG['data_dir'])\n", + "train_loader = get_loader('train', **kw)\n", + "val_loader = get_loader('val', **kw)\n", + "test_loader = get_loader('test', **kw)\n", + "\n", + "num_classes = len(np.unique(np.load(CONFIG['data_dir'] + '/y_all.npy')))\n", + "\n", + "# Model\n", + "if CONFIG['model'] == 'resnet18':\n", + " model = ResNet18_2D(num_classes=num_classes, in_channels=in_ch)\n", + "else:\n", + " model = STFTRADN(num_classes=num_classes, in_channels=in_ch)\n", + "model = model.to(device)\n", + "\n", + "optimizer = AdamW(model.parameters(), lr=CONFIG['lr'], weight_decay=CONFIG['weight_decay'])\n", + "criterion = torch.nn.CrossEntropyLoss()\n", + "scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)\n", + "\n", + "total_params = sum(p.numel() for p in model.parameters())\n", + "print(f'Device : {device}')\n", + "print(f'Model : {CONFIG[\"model\"]} ({total_params:,} params)')\n", + "print(f'Train : {len(train_loader.dataset):,} samples ({len(train_loader)} batches)')\n", + "print(f'Val : {len(val_loader.dataset):,} samples')\n", + "print(f'Test : {len(test_loader.dataset):,} samples')\n", + "print(f'Classes : {num_classes}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Cell 7: Training Loop \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n", + "best_val_acc = 0.0\n", + "best_path = f\"{CONFIG['save_dir']}/best_{CONFIG['model']}.pt\"\n", + "patience_ctr = 0\n", + "best_val_loss = float('inf')\n", + "\n", + "print('='*74)\n", + "print(f'{\"Ep\":>4} | {\"Train Loss\":>10} | {\"Train Acc\":>9} | {\"Val Loss\":>9} | {\"Val Acc\":>8} | {\"LR\":>8} | {\"Time\":>6}')\n", + "print('-'*74)\n", + "\n", + "for epoch in range(1, CONFIG['epochs'] + 1):\n", + " t0 = time.time()\n", + "\n", + " # Train\n", + " model.train()\n", + " tr_loss, tr_corr, tr_total = 0.0, 0, 0\n", + " for signals, labels, _ in train_loader:\n", + " signals, labels = signals.to(device), labels.to(device)\n", + " optimizer.zero_grad()\n", + " out = model(signals)\n", + " loss = criterion(out, labels)\n", + " loss.backward()\n", + " optimizer.step()\n", + " tr_loss += loss.item() * signals.size(0)\n", + " tr_corr += out.argmax(1).eq(labels).sum().item()\n", + " tr_total += labels.size(0)\n", + "\n", + " # Validate\n", + " model.eval()\n", + " vl_loss, vl_corr, vl_total = 0.0, 0, 0\n", + " with torch.no_grad():\n", + " for signals, labels, _ in val_loader:\n", + " signals, labels = signals.to(device), labels.to(device)\n", + " out = model(signals)\n", + " loss = criterion(out, labels)\n", + " vl_loss += loss.item() * signals.size(0)\n", + " vl_corr += out.argmax(1).eq(labels).sum().item()\n", + " vl_total += labels.size(0)\n", + "\n", + " tr_loss /= tr_total; tr_acc = 100. * tr_corr / tr_total\n", + " vl_loss /= vl_total; vl_acc = 100. * vl_corr / vl_total\n", + " current_lr = optimizer.param_groups[0]['lr']\n", + " scheduler.step(vl_loss)\n", + " elapsed = time.time() - t0\n", + "\n", + " marker = ' *' if vl_acc > best_val_acc else ''\n", + " print(f'{epoch:4d} | {tr_loss:10.4f} | {tr_acc:8.2f}% | {vl_loss:9.4f} | {vl_acc:7.2f}% | {current_lr:8.6f} | {elapsed:5.1f}s{marker}')\n", + "\n", + " if vl_acc > best_val_acc:\n", + " best_val_acc = vl_acc\n", + " torch.save({'epoch': epoch, 'model_state_dict': model.state_dict(),\n", + " 'val_acc': vl_acc, 'val_loss': vl_loss,\n", + " 'model_name': CONFIG['model'], 'num_classes': num_classes,\n", + " 'in_channels': in_ch}, best_path)\n", + "\n", + " if vl_loss < best_val_loss - 1e-4:\n", + " best_val_loss = vl_loss; patience_ctr = 0\n", + " else:\n", + " patience_ctr += 1\n", + " if patience_ctr >= CONFIG['patience']:\n", + " print(f'\\n[EARLY STOP] No val loss improvement for {CONFIG[\"patience\"]} epochs.')\n", + " break\n", + "\n", + "print('='*74)\n", + "print(f'\\n[DONE] Best val acc: {best_val_acc:.2f}% | Checkpoint: {best_path}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Cell 8: Final Test Evaluation + Per-SNR Accuracy \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n", + "checkpoint = torch.load(best_path, map_location=device, weights_only=True)\n", + "model.load_state_dict(checkpoint['model_state_dict'])\n", + "model.eval()\n", + "\n", + "te_loss, te_corr, te_total = 0.0, 0, 0\n", + "snr_correct, snr_total = {}, {}\n", + "\n", + "with torch.no_grad():\n", + " for signals, labels, snrs in test_loader:\n", + " signals, labels = signals.to(device), labels.to(device)\n", + " out = model(signals)\n", + " loss = criterion(out, labels)\n", + " pred = out.argmax(1)\n", + " te_loss += loss.item() * signals.size(0)\n", + " te_corr += pred.eq(labels).sum().item()\n", + " te_total += labels.size(0)\n", + " for i in range(labels.size(0)):\n", + " snr = int(snrs[i].item())\n", + " snr_correct.setdefault(snr, 0)\n", + " snr_total.setdefault(snr, 0)\n", + " snr_total[snr] += 1\n", + " if pred[i] == labels[i]: snr_correct[snr] += 1\n", + "\n", + "te_loss /= te_total\n", + "te_acc = 100. * te_corr / te_total\n", + "print(f'\\n[TEST RESULT] Loss: {te_loss:.4f} | Accuracy: {te_acc:.2f}%')\n", + "\n", + "print('\\n--- Per-SNR Accuracy ---')\n", + "print(f'{\"SNR (dB)\":>10} | {\"Accuracy\":>10} | Bar')\n", + "print('-'*50)\n", + "snr_accs = []\n", + "for snr in sorted(snr_correct):\n", + " acc = 100. * snr_correct[snr] / snr_total[snr]\n", + " snr_accs.append(acc)\n", + " bar = '#' * int(acc / 5)\n", + " print(f'{snr:>8d} dB | {acc:9.2f}% | {bar}')\n", + "print('-'*50)\n", + "print(f'{\"Average\":>10} | {np.mean(snr_accs):9.2f}%')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# \u00e2\u201d\u20ac\u00e2\u201d\u20ac Cell 9: Download Best Checkpoint \u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\u00e2\u201d\u20ac\n", + "from google.colab import files\n", + "print(f'Downloading {best_path} ...')\n", + "files.download(best_path)\n", + "print('Done!')" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} \ No newline at end of file diff --git a/dataset_loader.py b/dataset_loader.py index a1e48f0..7ab5e56 100644 --- a/dataset_loader.py +++ b/dataset_loader.py @@ -1,5 +1,6 @@ # dataset_loader.py import os +import warnings import numpy as np import torch from torch.utils.data import Dataset, DataLoader @@ -147,3 +148,132 @@ def get_rf_dataloader(data_dir="./data", split="train", batch_size=64, shuffle=N except Exception as e: print(f"\n[FAIL] PyTorch loader verification failed: {e}") print("Please ensure you run 'python dataset_agent.py' first to prepare the arrays.") + + +# ========================================== +# 2D Spectrogram Dataset (pre-computed STFT) +# ========================================== + +class SpectrogramDataset(Dataset): + """ + Fast 2D Dataset that loads pre-computed STFT spectrograms from + ./data/spectrograms_all.npy via memory-mapping. + + Eliminates the on-the-fly scipy.signal.stft() bottleneck entirely. + The full (220000, 3, 64, 5) array is memory-mapped — only the pages + actually requested by the DataLoader are paged into RAM. + + Args: + data_dir : directory containing spectrograms_all.npy + index files. + split : 'train', 'val', or 'test'. + in_channels: 1 → grayscale (power_db only), 3 → hybrid (all channels). + + Run once before training: + python precompute_stft.py + """ + + SPECTROGRAM_FILE = "spectrograms_all.npy" + + def __init__(self, data_dir: str = "./data", split: str = "train", + in_channels: int = 1): + self.in_channels = in_channels + split = split.lower() + + # ── Resolve data_dir ──────────────────────────────────────────── + if not os.path.exists(data_dir): + for alt in [ + os.path.join(os.path.dirname(os.path.abspath(__file__)), "data"), + os.path.join(".", "data"), + os.path.join("..", "data"), + ]: + if os.path.exists(alt): + data_dir = alt + break + + # ── Validate pre-computed file ─────────────────────────────────── + spec_path = os.path.join(data_dir, self.SPECTROGRAM_FILE) + if not os.path.exists(spec_path): + raise FileNotFoundError( + f"Pre-computed spectrogram file not found:\n {spec_path}\n" + "Run first: python precompute_stft.py" + ) + + # ── Memory-map the full spectrogram array (no RAM load) ────────── + # mmap_mode='r' → OS pages only the accessed rows into RAM. + self._spectrograms = np.load(spec_path, mmap_mode="r") # (N, 3, H, W) + + # ── Load labels & SNRs (tiny, load fully) ──────────────────────── + self._labels = np.load(os.path.join(data_dir, "y_all.npy")) # (N,) + self._snrs = np.load(os.path.join(data_dir, "snrs_all.npy")) # (N,) + + # ── Load split indices ─────────────────────────────────────────── + idx_map = {"train": "train_idx.npy", "val": "val_idx.npy", "test": "test_idx.npy"} + if split not in idx_map: + raise ValueError(f"split must be 'train', 'val', or 'test', got '{split}'") + self._indices = np.load(os.path.join(data_dir, idx_map[split])) + + self.split_name = split + if self.split_name == "train": + import torchvision.transforms as T + self.transform = T.RandomErasing(p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0) + else: + self.transform = None + + def __len__(self) -> int: + return len(self._indices) + + def __getitem__(self, idx: int): + real_idx = self._indices[idx] + + # np.array() materialises one row from the mmap into a regular array + spec = np.array(self._spectrograms[real_idx], dtype=np.float32) # (3, H, W) + + if self.in_channels == 1: + spec = spec[:1] # keep only power_db channel → (1, H, W) + + signal_tensor = torch.from_numpy(spec) + if self.transform is not None: + signal_tensor = self.transform(signal_tensor) + + label_tensor = torch.tensor(int(self._labels[real_idx]), dtype=torch.long) + snr = self._snrs[real_idx] + + return signal_tensor, label_tensor, snr + + +def get_spectrogram_dataloader( + data_dir: str = "./data", + split: str = "train", + batch_size: int = 64, + in_channels: int = 1, + shuffle: bool = None, + num_workers: int = None, +) -> DataLoader: + """ + Construct a fast DataLoader backed by SpectrogramDataset. + + num_workers defaults to min(4, cpu_count) for maximum throughput. + persistent_workers=True avoids worker-respawn overhead between epochs. + pin_memory=True (when CUDA available) accelerates host→GPU transfers. + """ + if shuffle is None: + shuffle = (split == "train") + if num_workers is None: + num_workers = min(4, os.cpu_count() or 1) + + dataset = SpectrogramDataset( + data_dir=data_dir, split=split, in_channels=in_channels + ) + + use_pin_memory = torch.cuda.is_available() + use_persistent = num_workers > 0 + + loader = DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + pin_memory=use_pin_memory, + persistent_workers=use_persistent, + ) + return loader diff --git a/health_report.md b/health_report.md deleted file mode 100644 index 47169d2..0000000 --- a/health_report.md +++ /dev/null @@ -1,48 +0,0 @@ -# Repository Health Report: RF_Signals_Classifier - -## 1. Overview -The **Intelligent RF Signal Classification Pipeline** aims to be an advanced, 4-stage machine learning architecture for Automatic Modulation Classification (AMC). The goal is to use deep learning for robust spectrum sensing. - -However, based on the current codebase, the repository is **only partially implemented**. It currently contains robust data ingestion, synthetic generation, preprocessing, and dataset verification tools. The core Machine Learning models and training pipelines (Stages 3 and 4) are **missing**. - -## 2. Current Status of the 4-Stage Pipeline - -| Stage | Name | Status | Description | -|---|---|---|---| -| **Stage 1** | Ingestion & Dataset Agent | **✅ Implemented** | Both `dataset_agent.py` (for RadioML dataset) and `generate_custom_rf_v4.m` (for synthetic generation) are present and functional. | -| **Stage 2** | Dual-Representation Preprocessing | **✅ Implemented** | Both Time-Domain (I/Q normalization) and Time-Frequency (STFT via `stft_converter.py`) branches exist and are integrated into PyTorch DataLoaders. | -| **Stage 3** | Adaptive CNN-Transformer Core | **❌ Missing** | No PyTorch models, convolutional layers, or Transformer architectures exist in the codebase. | -| **Stage 4** | Spectrum Cop Decision Engine | **❌ Missing** | No Outlier Detection, Continual Learning (EWC loss), or XAI (Grad-CAM) implementations are present. | - -## 3. File-by-File Breakdown - -### Data Preprocessing & Pipeline -* **`dataset_agent.py`**: A robust script that loads the legacy RadioML 2016.10a pickle file, normalizes the I/Q signals (zero-mean, unit-variance), and creates a locked 70/20/10 train/val/test split (saving as `.npy` files). -* **`dataset_loader.py`**: Provides the `RFSignalDataset` class and `get_rf_dataloader` function. Standardizes the PyTorch loading process using the splits generated by `dataset_agent.py`. -* **`verify_splits.py`**: A comprehensive analytics script that verifies data integrity, ensures no data leakage across splits, and generates a detailed visualization dashboard (`verification_results.png`). - -### Custom Dataset Generation -* **`generate_custom_rf_v4.m`**: A MATLAB script that generates custom synthetic RF signals (BPSK, QPSK, 8PSK, 16QAM, 64QAM) with realistic channel impairments (fading, Doppler, CFO). Fallbacks are implemented if the Communications Toolbox is unavailable. -* **`load_mat_dataset.py`**: PyTorch integration script that loads the `.mat` outputs from the MATLAB generator into a PyTorch `TensorDataset`. - -### Feature Engineering -* **`stft_converter.py`**: Handles the Time-Frequency branch of Stage 2. It converts I/Q sequences into Kaiser-windowed 2D power spectrograms for CNN processing. - -### Configuration -* **`requirements.txt`**: Standard Python environment requirements (includes `torch`, `scipy`, `scikit-learn`, `matplotlib`, `numpy`, etc.). -* **`README.md`**: Outlines the architectural framework but contains claims for features that have not yet been implemented in the code. - -## 4. Actionable Next Steps for Development - -To continue the development of this repository with full context, here is what you need to focus on next: - -1. **Generate or Download Data**: - * *Option A*: Download the `RML2016.10a_dict.pkl` dataset and place it in the `./data` directory, then run `python dataset_agent.py`. - * *Option B*: Run `generate_custom_rf_v4.m` in MATLAB to generate synthetic data, then use `load_mat_dataset.py`. -2. **Verify Data Pipeline**: Run `python verify_splits.py` to ensure the ingestion and splitting logic is sound. -3. **Implement Stage 3 (The Core Model)**: - * Create a new file (e.g., `models.py`). - * Implement the 1D/2D CNNs for Spatial Feature Extraction. - * Implement the local-windowed Transformer Encoder for Temporal Modeling. -4. **Implement Training Loop**: Create a `train.py` script to orchestrate the dataloaders (`dataset_loader.py`) and the new PyTorch models. -5. **Implement Stage 4 (Decision Engine)**: Later on, build out the MCD outlier detection, Elastic Weight Consolidation (EWC) for continual learning, and Grad-CAM integration for interpretability. diff --git a/models_2d.py b/models_2d.py index df73b52..2f19d6b 100644 --- a/models_2d.py +++ b/models_2d.py @@ -25,6 +25,7 @@ def __init__(self, in_planes, planes, stride=1): in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False ) self.bn1 = nn.BatchNorm2d(planes) + self.dropout = nn.Dropout2d(p=0.3) self.conv2 = nn.Conv2d( planes, planes, kernel_size=3, stride=1, padding=1, bias=False ) @@ -43,6 +44,7 @@ def __init__(self, in_planes, planes, stride=1): def forward(self, x): out = F.relu(self.bn1(self.conv1(x))) + out = self.dropout(out) out = self.bn2(self.conv2(out)) out += self.shortcut(x) out = F.relu(out) @@ -74,23 +76,23 @@ class ResNet18_2D(nn.Module): def __init__(self, num_classes=11, in_channels=1): super().__init__() - self.in_planes = 64 + self.in_planes = 16 # ── Stem (no MaxPool) ────────────────────────────────── self.conv1 = nn.Conv2d( - in_channels, 64, kernel_size=3, stride=1, padding=1, bias=False + in_channels, 16, kernel_size=3, stride=1, padding=1, bias=False ) - self.bn1 = nn.BatchNorm2d(64) + self.bn1 = nn.BatchNorm2d(16) # ── Residual Layers [2, 2, 2, 2] ────────────────────── - self.layer1 = self._make_layer(64, num_blocks=2, stride=1) - self.layer2 = self._make_layer(128, num_blocks=2, stride=2) - self.layer3 = self._make_layer(256, num_blocks=2, stride=2) - self.layer4 = self._make_layer(512, num_blocks=2, stride=2) + self.layer1 = self._make_layer(16, num_blocks=2, stride=1) + self.layer2 = self._make_layer(32, num_blocks=2, stride=2) + self.layer3 = self._make_layer(64, num_blocks=2, stride=2) + self.layer4 = self._make_layer(128, num_blocks=2, stride=2) # ── Classification Head ──────────────────────────────── self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) - self.fc = nn.Linear(512 * BasicBlock.expansion, num_classes) + self.fc = nn.Linear(128 * BasicBlock.expansion, num_classes) # ── Weight Initialization (Kaiming) ──────────────────── self._init_weights() @@ -122,12 +124,12 @@ def forward(self, x): Forward pass through the modified ResNet-18. Feature map sizes for input [B, 1, 64, 5]: - conv1 → [B, 64, 64, 5] - layer1 → [B, 64, 64, 5] - layer2 → [B, 128, 32, 3] - layer3 → [B, 256, 16, 2] - layer4 → [B, 512, 8, 1] - avgpool→ [B, 512, 1, 1] + conv1 → [B, 16, 64, 5] + layer1 → [B, 16, 64, 5] + layer2 → [B, 32, 32, 3] + layer3 → [B, 64, 16, 2] + layer4 → [B, 128, 8, 1] + avgpool→ [B, 128, 1, 1] fc → [B, num_classes] """ # Stem diff --git a/precompute_stft.py b/precompute_stft.py new file mode 100644 index 0000000..b490e1f --- /dev/null +++ b/precompute_stft.py @@ -0,0 +1,128 @@ +# -*- coding: utf-8 -*- +# precompute_stft.py +# ========================================================================= +# One-time preprocessing step: converts all 220,000 IQ signals in X_all.npy +# into pre-computed 3-channel STFT spectrograms and saves to: +# ./data/spectrograms_all.npy → shape (220000, 3, 64, 5), float32 +# +# Channels stored: +# ch0: power_db (used alone in "grayscale" mode) +# ch1: phase +# ch2: magnitude (ch0+ch1+ch2 used in "hybrid" mode) +# +# Run ONCE before training: +# python precompute_stft.py +# +# Output file is memory-mapped during training — no full load into RAM. +# ========================================================================= + +import os +import time +import warnings +import numpy as np +from scipy.signal import stft as scipy_stft + +warnings.filterwarnings("ignore", message=".*complex.*return_onesided.*") + +# ── Config ──────────────────────────────────────────────────────────────── +DATA_DIR = "./data" +INPUT_FILE = os.path.join(DATA_DIR, "X_all.npy") +OUTPUT_FILE = os.path.join(DATA_DIR, "spectrograms_all.npy") + +STFT_PARAMS = dict(window=("kaiser", 0.85), nperseg=64, noverlap=32) +EXPECTED_H, EXPECTED_W = 64, 5 +BATCH_REPORT = 10_000 # print progress every N samples + + +# ── Core converter ──────────────────────────────────────────────────────── + +def compute_hybrid_spectrogram(signal_np: np.ndarray) -> np.ndarray: + """ + Convert one (2, 128) IQ signal → (3, 64, 5) float32 spectrogram. + Uses identical Kaiser-windowed STFT parameters as stft_converter.py. + + Returns: + np.ndarray shape (3, H, W): + [0] power_db — 10·log10(|STFT|² + ε) + [1] phase — angle(STFT) + [2] magnitude — |STFT| + """ + I, Q = signal_np[0], signal_np[1] + complex_signal = I + 1j * Q + + _, _, Zxx = scipy_stft(complex_signal, **STFT_PARAMS) + Zxx = np.fft.fftshift(Zxx, axes=0) # centre DC + + power_db = 10.0 * np.log10(np.abs(Zxx) ** 2 + 1e-10).astype(np.float32) + phase = np.angle(Zxx).astype(np.float32) + magnitude = np.abs(Zxx).astype(np.float32) + + return np.stack([power_db, phase, magnitude], axis=0) # (3, H, W) + + +# ── Main ───────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + print("=" * 65) + print(" STFT PRE-COMPUTATION -- stft_converter -> spectrograms_all.npy") + print("=" * 65) + + # ── Verify input ────────────────────────────────────────────────────── + if not os.path.exists(INPUT_FILE): + raise FileNotFoundError(f"Input not found: {INPUT_FILE}") + + if os.path.exists(OUTPUT_FILE): + print(f"\n[SKIP] Output already exists: {OUTPUT_FILE}") + arr = np.load(OUTPUT_FILE, mmap_mode="r") + print(f" Shape: {arr.shape} dtype: {arr.dtype}") + print("\nDelete the file and re-run to force recompute.") + raise SystemExit(0) + + # ── Load raw IQ ─────────────────────────────────────────────────────── + print(f"\nLoading {INPUT_FILE} …") + X = np.load(INPUT_FILE) # (N, 2, 128) + N = len(X) + print(f" Loaded {N:,} signals shape={X.shape} dtype={X.dtype}") + + # Validate STFT output shape with one sample + sample_spec = compute_hybrid_spectrogram(X[0]) + H, W = sample_spec.shape[1], sample_spec.shape[2] + assert (H, W) == (EXPECTED_H, EXPECTED_W), \ + f"Unexpected spectrogram shape ({H}, {W}); expected ({EXPECTED_H}, {EXPECTED_W})" + print(f" STFT output validated: (3, {H}, {W}) OK") + + # ── Pre-allocate output array ───────────────────────────────────────── + out_shape = (N, 3, H, W) + ram_mb = N * 3 * H * W * 4 / 1e6 + print(f"\nAllocating output array {out_shape} ({ram_mb:.0f} MB RAM) …") + spectrograms = np.empty(out_shape, dtype=np.float32) + + # ── Convert ────────────────────────────────────────────────────────── + print(f"\nConverting {N:,} signals — reporting every {BATCH_REPORT:,} …\n") + t_start = time.perf_counter() + + for i in range(N): + spectrograms[i] = compute_hybrid_spectrogram(X[i]) + + if (i + 1) % BATCH_REPORT == 0 or i == 0: + elapsed = time.perf_counter() - t_start + rate = (i + 1) / elapsed + eta_s = (N - i - 1) / rate + eta_m = eta_s / 60 + pct = 100.0 * (i + 1) / N + print(f" [{pct:5.1f}%] {i+1:>7,}/{N:,}" + f" | {rate:,.0f} sig/s" + f" | ETA: {eta_m:.1f} min") + + elapsed_total = time.perf_counter() - t_start + print(f"\n[DONE] Converted {N:,} signals in {elapsed_total/60:.1f} min " + f"({elapsed_total:.1f}s)") + + # ── Save as raw .npy (supports mmap_mode during training) ──────────── + print(f"\nSaving to {OUTPUT_FILE} …") + np.save(OUTPUT_FILE, spectrograms) + file_mb = os.path.getsize(OUTPUT_FILE) / 1e6 + print(f"[SAVED] {OUTPUT_FILE} ({file_mb:.0f} MB on disk)") + print(f" Shape : {spectrograms.shape}") + print(f" dtype : {spectrograms.dtype}") + print(f"\n ✓ Ready for fast training. Run: python train_2d.py") diff --git a/train_2d.py b/train_2d.py index ae7a364..9efcb65 100644 --- a/train_2d.py +++ b/train_2d.py @@ -1,71 +1,32 @@ # train_2d.py # ========================================================================= -# CHUNK 4: 2D Training Loop +# CHUNK 4: 2D Training Loop [OPTIMISED — pre-computed STFT pipeline] # Comprehensive training orchestrator for ResNet-18 and STFT-RADN. # AdamW optimizer, CrossEntropy loss, ReduceLROnPlateau scheduler, # early stopping on validation loss, and per-SNR accuracy logging # to match Team A's benchmarking format. +# +# Data pipeline: +# python precompute_stft.py ← run ONCE to build spectrograms_all.npy +# python train_2d.py ← fast training via mmap + num_workers # ========================================================================= import os import sys import time -import warnings import argparse import numpy as np import torch import torch.nn as nn from torch.optim import AdamW from torch.optim.lr_scheduler import ReduceLROnPlateau -from scipy.signal import stft as scipy_stft -from dataset_loader import get_rf_dataloader +from dataset_loader import get_spectrogram_dataloader from models_2d import ResNet18_2D, STFTRADN -# Suppress scipy's harmless "complex input" warning during STFT -warnings.filterwarnings("ignore", message=".*complex.*return_onesided.*") - -# ────────────────────────────────────────────────────────────── -# Transform: 1D IQ -> 2D Spectrogram (on-the-fly) -# ────────────────────────────────────────────────────────────── - -class IQToSpectrogram: - """ - On-the-fly transform bridging Person 1's DataLoader and Person 4's models. - Converts a (2, 128) IQ tensor to a 2D spectrogram tensor using - Person 3's exact Kaiser-windowed STFT parameters. - - Args: - mode: "grayscale" -> (1, 64, 5) single power spectrogram - "hybrid" -> (3, 64, 5) [power_dB, phase, magnitude] - """ - STFT_PARAMS = dict(window=("kaiser", 0.85), nperseg=64, noverlap=32) - - def __init__(self, mode="grayscale"): - self.mode = mode - - def __call__(self, signal_tensor): - signal_np = signal_tensor.numpy() - I, Q = signal_np[0], signal_np[1] - complex_signal = I + 1j * Q - - _, _, Zxx = scipy_stft(complex_signal, **self.STFT_PARAMS) - - if self.mode == "grayscale": - power_db = 10 * np.log10(np.abs(Zxx) ** 2 + 1e-10) - spec = np.fft.fftshift(power_db, axes=0) - return torch.tensor(spec, dtype=torch.float32).unsqueeze(0) # (1, H, W) - - else: # hybrid - Zxx_s = np.fft.fftshift(Zxx, axes=0) - power_db = 10 * np.log10(np.abs(Zxx_s) ** 2 + 1e-10) - phase = np.angle(Zxx_s) - magnitude = np.abs(Zxx_s) - return torch.tensor( - np.stack([power_db, phase, magnitude], axis=0), - dtype=torch.float32, - ) # (3, H, W) +# IQToSpectrogram removed — STFT is now pre-computed once by precompute_stft.py +# and loaded via memory-mapped SpectrogramDataset for zero per-epoch overhead. # ────────────────────────────────────────────────────────────── @@ -185,36 +146,33 @@ def step(self, val_loss): def main(args): # ── Device ───────────────────────────────────────────────── device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + in_channels = 1 if args.input_mode == "grayscale" else 3 + num_workers = args.num_workers if args.num_workers >= 0 \ + else min(4, os.cpu_count() or 1) print("=" * 70) - print(" 2D MODEL TRAINING LOOP") + print(" 2D MODEL TRAINING LOOP [pre-computed STFT pipeline]") print("=" * 70) - print(f" Device : {device}") - print(f" Model : {args.model}") - print(f" Input mode : {args.input_mode}") - print(f" LR : {args.lr} | Weight decay: {args.weight_decay}") - print(f" Batch size : {args.batch_size} | Max epochs: {args.epochs}") - print(f" Patience : {args.patience}") + print(f" Device : {device}") + print(f" Model : {args.model}") + print(f" Input mode : {args.input_mode} ({in_channels}ch)") + print(f" LR : {args.lr} | Weight decay: {args.weight_decay}") + print(f" Batch size : {args.batch_size} | Max epochs: {args.epochs}") + print(f" Patience : {args.patience}") + print(f" num_workers : {num_workers} | pin_memory: {torch.cuda.is_available()}") print("=" * 70) - # ── Transform ────────────────────────────────────────────── - transform = IQToSpectrogram(mode=args.input_mode) - in_channels = 1 if args.input_mode == "grayscale" else 3 - - # ── DataLoaders ──────────────────────────────────────────── - print("\n--- Loading dataset splits ---") - train_loader = get_rf_dataloader( - data_dir=args.data_dir, split="train", - batch_size=args.batch_size, transform=transform, - ) - val_loader = get_rf_dataloader( - data_dir=args.data_dir, split="val", - batch_size=args.batch_size, transform=transform, - ) - test_loader = get_rf_dataloader( - data_dir=args.data_dir, split="test", - batch_size=args.batch_size, transform=transform, + # ── DataLoaders (backed by mmap spectrograms_all.npy) ────── + print("\n--- Loading dataset splits (pre-computed STFT) ---") + loader_kwargs = dict( + data_dir=args.data_dir, + batch_size=args.batch_size, + in_channels=in_channels, + num_workers=num_workers, ) + train_loader = get_spectrogram_dataloader(split="train", **loader_kwargs) + val_loader = get_spectrogram_dataloader(split="val", **loader_kwargs) + test_loader = get_spectrogram_dataloader(split="test", **loader_kwargs) print(f" Train : {len(train_loader.dataset):>6} samples ({len(train_loader)} batches)") print(f" Val : {len(val_loader.dataset):>6} samples ({len(val_loader)} batches)") print(f" Test : {len(test_loader.dataset):>6} samples ({len(test_loader)} batches)") @@ -366,13 +324,18 @@ def main(args): help="AdamW learning rate (default: 1e-3)", ) parser.add_argument( - "--weight-decay", type=float, default=1e-4, - help="AdamW weight decay (default: 1e-4)", + "--weight-decay", type=float, default=1e-2, + help="AdamW weight decay (default: 1e-2)", ) parser.add_argument( "--patience", type=int, default=10, help="Early stopping patience in epochs (default: 10)", ) + parser.add_argument( + "--num-workers", type=int, default=-1, + help="DataLoader worker processes. -1 = auto (min(4, cpu_count)). " + "Set 0 to disable multiprocessing (default: -1 / auto)", + ) args = parser.parse_args() main(args) From 9d1f9bdc08fe9ee8a745000f933f3e8f307b6863 Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sat, 29 Aug 2026 20:19:12 +0530 Subject: [PATCH 7/9] Fix underfitting: remove over-regularization - Removed Dropout2d from BasicBlock and moved Dropout(0.5) to final FC layer - Changed ResNet18_2D channels to 32->64->128->256 (~2.8M params) - Removed destructive RandomErasing SpecAugment from 64x5 spectrograms - Reverted weight_decay to 1e-4 --- colab_train.ipynb | 29 +++++++++++++---------------- dataset_loader.py | 6 +----- models_2d.py | 20 ++++++++++---------- train_2d.py | 4 ++-- 4 files changed, 26 insertions(+), 33 deletions(-) diff --git a/colab_train.ipynb b/colab_train.ipynb index 5b7a475..4bed107 100644 --- a/colab_train.ipynb +++ b/colab_train.ipynb @@ -84,7 +84,6 @@ " super().__init__()\n", " self.conv1 = nn.Conv2d(in_planes, planes, 3, stride=stride, padding=1, bias=False)\n", " self.bn1 = nn.BatchNorm2d(planes)\n", - " self.dropout = nn.Dropout2d(p=0.3)\n", " self.conv2 = nn.Conv2d(planes, planes, 3, stride=1, padding=1, bias=False)\n", " self.bn2 = nn.BatchNorm2d(planes)\n", " self.shortcut = nn.Sequential()\n", @@ -95,7 +94,6 @@ " )\n", " def forward(self, x):\n", " out = F.relu(self.bn1(self.conv1(x)))\n", - " out = self.dropout(out)\n", " out = self.bn2(self.conv2(out))\n", " out += self.shortcut(x)\n", " return F.relu(out)\n", @@ -103,15 +101,16 @@ "class ResNet18_2D(nn.Module):\n", " def __init__(self, num_classes=11, in_channels=1):\n", " super().__init__()\n", - " self.in_planes = 16\n", - " self.conv1 = nn.Conv2d(in_channels, 16, 3, stride=1, padding=1, bias=False)\n", - " self.bn1 = nn.BatchNorm2d(16)\n", - " self.layer1 = self._make_layer(16, 2, stride=1)\n", - " self.layer2 = self._make_layer(32, 2, stride=2)\n", - " self.layer3 = self._make_layer(64, 2, stride=2)\n", - " self.layer4 = self._make_layer(128, 2, stride=2)\n", + " self.in_planes = 32\n", + " self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=1, padding=1, bias=False)\n", + " self.bn1 = nn.BatchNorm2d(32)\n", + " self.layer1 = self._make_layer(32, 2, stride=1)\n", + " self.layer2 = self._make_layer(64, 2, stride=2)\n", + " self.layer3 = self._make_layer(128, 2, stride=2)\n", + " self.layer4 = self._make_layer(256, 2, stride=2)\n", " self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n", - " self.fc = nn.Linear(128, num_classes)\n", + " self.dropout = nn.Dropout(p=0.5)\n", + " self.fc = nn.Linear(256, num_classes)\n", " for m in self.modules():\n", " if isinstance(m, nn.Conv2d):\n", " nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n", @@ -132,7 +131,9 @@ " x = self.layer1(x); x = self.layer2(x)\n", " x = self.layer3(x); x = self.layer4(x)\n", " x = self.avgpool(x)\n", - " return self.fc(torch.flatten(x, 1))\n", + " x = torch.flatten(x, 1)\n", + " x = self.dropout(x)\n", + " return self.fc(x)\n", "\n", "class ChannelAttention(nn.Module):\n", " def __init__(self, channels, reduction=16):\n", @@ -227,8 +228,6 @@ " self._labels = np.load(f'{data_dir}/y_all.npy')\n", " self._snrs = np.load(f'{data_dir}/snrs_all.npy')\n", " self._indices = np.load(f'{data_dir}/{split}_idx.npy')\n", - " import torchvision.transforms as T\n", - " self.transform = T.RandomErasing(p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0) if split == 'train' else None\n", " def __len__(self):\n", " return len(self._indices)\n", " def __getitem__(self, idx):\n", @@ -237,8 +236,6 @@ " if self.in_channels == 1:\n", " spec = spec[:1]\n", " signal_tensor = torch.from_numpy(spec)\n", - " if self.transform is not None:\n", - " signal_tensor = self.transform(signal_tensor)\n", " return signal_tensor, torch.tensor(int(self._labels[i]), dtype=torch.long), self._snrs[i]\n", "\n", "def get_loader(split, batch_size, in_channels, num_workers=2, data_dir='/content/data'):\n", @@ -262,7 +259,7 @@ " batch_size = 256,\n", " epochs = 30,\n", " lr = 1e-3,\n", - " weight_decay = 1e-2,\n", + " weight_decay = 1e-4,\n", " patience = 10, # early stopping\n", " num_workers = 2,\n", " data_dir = '/content/data',\n", diff --git a/dataset_loader.py b/dataset_loader.py index 7ab5e56..c880d87 100644 --- a/dataset_loader.py +++ b/dataset_loader.py @@ -213,11 +213,7 @@ def __init__(self, data_dir: str = "./data", split: str = "train", self._indices = np.load(os.path.join(data_dir, idx_map[split])) self.split_name = split - if self.split_name == "train": - import torchvision.transforms as T - self.transform = T.RandomErasing(p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0) - else: - self.transform = None + self.transform = None def __len__(self) -> int: return len(self._indices) diff --git a/models_2d.py b/models_2d.py index 2f19d6b..48db5c1 100644 --- a/models_2d.py +++ b/models_2d.py @@ -25,7 +25,6 @@ def __init__(self, in_planes, planes, stride=1): in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False ) self.bn1 = nn.BatchNorm2d(planes) - self.dropout = nn.Dropout2d(p=0.3) self.conv2 = nn.Conv2d( planes, planes, kernel_size=3, stride=1, padding=1, bias=False ) @@ -44,7 +43,6 @@ def __init__(self, in_planes, planes, stride=1): def forward(self, x): out = F.relu(self.bn1(self.conv1(x))) - out = self.dropout(out) out = self.bn2(self.conv2(out)) out += self.shortcut(x) out = F.relu(out) @@ -76,23 +74,24 @@ class ResNet18_2D(nn.Module): def __init__(self, num_classes=11, in_channels=1): super().__init__() - self.in_planes = 16 + self.in_planes = 32 # ── Stem (no MaxPool) ────────────────────────────────── self.conv1 = nn.Conv2d( - in_channels, 16, kernel_size=3, stride=1, padding=1, bias=False + in_channels, 32, kernel_size=3, stride=1, padding=1, bias=False ) - self.bn1 = nn.BatchNorm2d(16) + self.bn1 = nn.BatchNorm2d(32) # ── Residual Layers [2, 2, 2, 2] ────────────────────── - self.layer1 = self._make_layer(16, num_blocks=2, stride=1) - self.layer2 = self._make_layer(32, num_blocks=2, stride=2) - self.layer3 = self._make_layer(64, num_blocks=2, stride=2) - self.layer4 = self._make_layer(128, num_blocks=2, stride=2) + self.layer1 = self._make_layer(32, num_blocks=2, stride=1) + self.layer2 = self._make_layer(64, num_blocks=2, stride=2) + self.layer3 = self._make_layer(128, num_blocks=2, stride=2) + self.layer4 = self._make_layer(256, num_blocks=2, stride=2) # ── Classification Head ──────────────────────────────── self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) - self.fc = nn.Linear(128 * BasicBlock.expansion, num_classes) + self.dropout = nn.Dropout(p=0.5) + self.fc = nn.Linear(256 * BasicBlock.expansion, num_classes) # ── Weight Initialization (Kaiming) ──────────────────── self._init_weights() @@ -144,6 +143,7 @@ def forward(self, x): # Pool → Flatten → Classify x = self.avgpool(x) x = torch.flatten(x, 1) + x = self.dropout(x) x = self.fc(x) return x diff --git a/train_2d.py b/train_2d.py index 9efcb65..18645f9 100644 --- a/train_2d.py +++ b/train_2d.py @@ -324,8 +324,8 @@ def main(args): help="AdamW learning rate (default: 1e-3)", ) parser.add_argument( - "--weight-decay", type=float, default=1e-2, - help="AdamW weight decay (default: 1e-2)", + "--weight-decay", type=float, default=1e-4, + help="AdamW weight decay (default: 1e-4)", ) parser.add_argument( "--patience", type=int, default=10, From a484fdade77f9743ef5f811ea0d3c3a812521799 Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sat, 29 Aug 2026 23:17:55 +0530 Subject: [PATCH 8/9] Optimize STFT for higher temporal resolution to hit 50% accuracy - Changed precompute_stft.py STFT params to nperseg=32, noverlap=28 (resulting in 32x25 spectrograms) - Updated weight_decay to 1e-3 in training config for ideal middle-ground regularization --- colab_train.ipynb | 2 +- precompute_stft.py | 4 ++-- train_2d.py | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/colab_train.ipynb b/colab_train.ipynb index 4bed107..5d9a185 100644 --- a/colab_train.ipynb +++ b/colab_train.ipynb @@ -259,7 +259,7 @@ " batch_size = 256,\n", " epochs = 30,\n", " lr = 1e-3,\n", - " weight_decay = 1e-4,\n", + " weight_decay = 1e-3,\n", " patience = 10, # early stopping\n", " num_workers = 2,\n", " data_dir = '/content/data',\n", diff --git a/precompute_stft.py b/precompute_stft.py index b490e1f..2717ce2 100644 --- a/precompute_stft.py +++ b/precompute_stft.py @@ -29,8 +29,8 @@ INPUT_FILE = os.path.join(DATA_DIR, "X_all.npy") OUTPUT_FILE = os.path.join(DATA_DIR, "spectrograms_all.npy") -STFT_PARAMS = dict(window=("kaiser", 0.85), nperseg=64, noverlap=32) -EXPECTED_H, EXPECTED_W = 64, 5 +STFT_PARAMS = dict(window=("kaiser", 0.85), nperseg=32, noverlap=28) +EXPECTED_H, EXPECTED_W = 32, 25 BATCH_REPORT = 10_000 # print progress every N samples diff --git a/train_2d.py b/train_2d.py index 18645f9..8d4a38c 100644 --- a/train_2d.py +++ b/train_2d.py @@ -324,8 +324,8 @@ def main(args): help="AdamW learning rate (default: 1e-3)", ) parser.add_argument( - "--weight-decay", type=float, default=1e-4, - help="AdamW weight decay (default: 1e-4)", + "--weight-decay", type=float, default=1e-3, + help="AdamW weight decay (default: 1e-3)", ) parser.add_argument( "--patience", type=int, default=10, From 401950adf4d99c71081952894bdf62b7f34230e5 Mon Sep 17 00:00:00 2001 From: Anjalii146 Date: Sat, 29 Aug 2026 23:27:10 +0530 Subject: [PATCH 9/9] Fix STFT assertion to 32x33 dimensions --- precompute_stft.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/precompute_stft.py b/precompute_stft.py index 2717ce2..65844fa 100644 --- a/precompute_stft.py +++ b/precompute_stft.py @@ -30,7 +30,7 @@ OUTPUT_FILE = os.path.join(DATA_DIR, "spectrograms_all.npy") STFT_PARAMS = dict(window=("kaiser", 0.85), nperseg=32, noverlap=28) -EXPECTED_H, EXPECTED_W = 32, 25 +EXPECTED_H, EXPECTED_W = 32, 33 BATCH_REPORT = 10_000 # print progress every N samples