Samy

.gitignore

@@ -26,6 +26,12 @@ share/python-wheels/
 *.egg
 MANIFEST
 
+# Data Directory/Files
+data/
+
+# Model Weights Files
+modelWeights
+
 # PyInstaller
 #  Usually these files are written by a python script from a template
 #  before PyInstaller builds the exe, so as to inject date/other infos into it.
@@ -127,6 +133,7 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+uv.*
 
 # Spyder project settings
 .spyderproject

requirements.txt

@@ -0,0 +1,5 @@
+torch
+edit_distance
+numpy
+pytest
+hydra-core

scripts/train_model.py

@@ -1,28 +1,47 @@
 
+#import sys
+#import pathlib
+
+# Ensure the project's `src/` directory is on sys.path so package imports work
+# when running this script directly (without installing the package).
+#sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / 'src'))
+
 modelName = 'speechBaseline4'
 
 args = {}
-args['outputDir'] = '/oak/stanford/groups/henderj/stfan/logs/speech_logs/' + modelName
-args['datasetPath'] = '/oak/stanford/groups/henderj/fwillett/speech/ptDecoder_ctc'
+args['outputDir'] = '../models/' + modelName
+args['datasetPath'] = '../data/ptDecoder_ctc'
 args['seqLen'] = 150
 args['maxTimeSeriesLen'] = 1200
-args['batchSize'] = 64
-args['lrStart'] = 0.02
+args['batchSize'] = 128
+args['optimizer'] = 'ADAM'
+args['warmupSteps'] = 500
+args['decayType'] = 'linear'
+args['lrStart'] = 0.05
 args['lrEnd'] = 0.02
-args['nUnits'] = 1024
+args['optimizerEps'] = 0.1
+args['SGDMomentum'] = 0
+args['l2_decay'] = 1e-5
+args['nUnits'] = 256
 args['nBatch'] = 10000 #3000
+args['patience'] = 10000 # Reduce to implement early stopping
 args['nLayers'] = 5
 args['seed'] = 0
 args['nClasses'] = 40
-args['nInputFeatures'] = 256
-args['dropout'] = 0.4
+args['nInputFeatures'] = 256 # Number of electrodes * 2
+args['nThresholdCrossings'] = 128
+args['nSpikeBandPowers'] = 128
+args['dropout'] = 0.2
 args['whiteNoiseSD'] = 0.8
 args['constantOffsetSD'] = 0.2
 args['gaussianSmoothWidth'] = 2.0
 args['strideLen'] = 4
 args['kernelLen'] = 32
-args['bidirectional'] = True
-args['l2_decay'] = 1e-5
+args['bidirectional'] = False
+
+# --- Time masking hyperparameters ---
+args['timeMaskNum'] = 20        # N: number of masks per trial
+args['timeMaskMaxFrac'] = 0.075 # M: max mask length as fraction of trial length
 
 from neural_decoder.neural_decoder_trainer import trainModel
 

setup.cfg

@@ -40,13 +40,15 @@ python_requires = >=3.9
 # For more information, check out https://semver.org/.
 install_requires =
     importlib-metadata; python_version<"3.8"
-    torch==1.13.1 
+    # torch==1.13.1 
     hydra-core==1.3.2
     hydra-submitit-launcher==1.1.5
     hydra-optuna-sweeper==1.2.0
-    numpy==1.25.0
-    scipy==1.11.1
-    numba==0.58.1
+    numpy==1.26.4
+    # scipy==1.11.1
+    scipy==1.16.0
+    # numba==0.58.1
+    numba==0.62.0
     scikit-learn==1.3.2
     g2p_en==2.1.0
     edit_distance==1.0.6
@@ -56,4 +58,4 @@ install_requires =
 where = src
 exclude =
     tests
-    examples
+    examples

src/neural_decoder/__init__.py

@@ -0,0 +1,12 @@
+"""neural_decoder package initializer.
+
+This file makes `src/neural_decoder` a regular Python package so relative
+and absolute imports work reliably when running scripts from the repository.
+"""
+
+__all__ = [
+    "augmentations",
+    "dataset",
+    "model",
+    "neural_decoder_trainer",
+]

src/neural_decoder/augmentations.py

@@ -24,6 +24,30 @@ def forward(self, x):
         noise = torch.randn(1, C) * self.std
         return x + noise
 
+class SpeckleNoise(nn.Module):
+    """
+    Speckled masking / coordinated dropout over neural inputs.
+
+    This applies elementwise dropout on an input tensor of shape [B, T, C]
+    (batch, time, channels) with probability p. During training, each element
+    is zeroed with probability p and the remaining elements are scaled by
+    1 / (1 - p). At eval time, the input is returned unchanged.
+    """
+    def __init__(self, p: float = 0.3):
+        super().__init__()
+        self.p = p
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x can be [B, T, C] or any shape; we treat all elements independently.
+        if (not self.training) or self.p <= 0.0:
+            return x
+        # mask: 1 with prob (1-p), 0 with prob p
+        mask = (torch.rand_like(x) > self.p).to(x.dtype)
+        # Inverted dropout scaling so expectation is preserved
+        x = x * mask / (1.0 - self.p)
+        return x
+
+
 class GaussianSmoothing(nn.Module):
     """
     Apply gaussian smoothing on a

src/neural_decoder/conf/config.yaml

@@ -16,6 +16,7 @@ datasetPath: /oak/stanford/groups/henderj/stfan/data/ptDecoder_ctc
 
 seed: 0
 batchSize: 64
+rnn_type: "lstm"
 lrStart: 0.02
 lrEnd: 0.02
 l2_decay: 1e-5
@@ -25,6 +26,8 @@ whiteNoiseSD: 0.8
 constantOffsetSD: 0.2
 gaussianSmoothWidth: 2.0
 
+max_grad_norm: 5.0      
+speckle_prob: 0.3      
 nUnits: 1024
 nLayers: 5
 nInputFeatures: 256

src/neural_decoder/dataset.py

@@ -1,9 +1,10 @@
+import numpy as np
 import torch
 from torch.utils.data import Dataset
 
 
 class SpeechDataset(Dataset):
-    def __init__(self, data, transform=None):
+    def __init__(self, data, transform=None, num_threshold_crossings=128, num_spike_band_powers=128):
         self.data = data
         self.transform = transform
         self.n_days = len(data)
@@ -14,13 +15,38 @@ def __init__(self, data, transform=None):
         self.neural_time_bins = []
         self.phone_seq_lens = []
         self.days = []
-        for day in range(self.n_days):
-            for trial in range(len(data[day]["sentenceDat"])):
-                self.neural_feats.append(data[day]["sentenceDat"][trial])
-                self.phone_seqs.append(data[day]["phonemes"][trial])
-                self.neural_time_bins.append(data[day]["sentenceDat"][trial].shape[0])
-                self.phone_seq_lens.append(data[day]["phoneLens"][trial])
-                self.days.append(day)
+
+        if (num_threshold_crossings == 128) & (num_spike_band_powers == 128):
+            for day in range(self.n_days):
+                for trial in range(len(data[day]["sentenceDat"])):
+                    self.neural_feats.append(data[day]["sentenceDat"][trial])
+                    self.phone_seqs.append(data[day]["phonemes"][trial])
+                    self.neural_time_bins.append(data[day]["sentenceDat"][trial].shape[0])
+                    self.phone_seq_lens.append(data[day]["phoneLens"][trial])
+                    self.days.append(day)
+        else:
+            pcs_tc = np.load("../data/threshold_crossing_principal_components.npy")
+            pcs_sbp = np.load("../data/spike_band_power_principal_components.npy")
+
+            tc_mean = np.load("../data/threshold_crossings_mean.npy")
+            sbp_mean = np.load("../data/spike_band_power_mean.npy")
+
+            for day in range(self.n_days):
+                for trial in range(len(data[day]["sentenceDat"])):
+                    trial_data = data[day]["sentenceDat"][trial]
+                    sample_tcs = trial_data.T[:128]
+                    sample_sbp = trial_data.T[128:]
+
+                    tc_new = (pcs_tc[:, :num_threshold_crossings].T @ (sample_tcs - tc_mean)).T
+                    sbp_new = (pcs_sbp[:, :num_spike_band_powers].T @ (sample_sbp - sbp_mean)).T
+
+                    trial_data_new = np.concatenate([tc_new, sbp_new], axis=1)
+
+                    self.neural_feats.append(trial_data_new)
+                    self.phone_seqs.append(data[day]["phonemes"][trial])
+                    self.neural_time_bins.append(data[day]["sentenceDat"][trial].shape[0])
+                    self.phone_seq_lens.append(data[day]["phoneLens"][trial])
+                    self.days.append(day)
 
     def __len__(self):
         return self.n_trials