{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":102335,"databundleVersionId":12518947,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":12328761,"sourceType":"datasetVersion","datasetId":7771623},{"sourceId":242954653,"sourceType":"kernelVersion"}],"dockerImageVersionId":31041,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This notebook is based on [Salman Ahmed](https://www.kaggle.com/salmanahmedtamu)'s \"[PYTORCH TRAINING / INFER PIPELINE 0.72](https://www.kaggle.com/code/salmanahmedtamu/pytorch-training-infer-pipeline-0-72)\", with modifications to the model, additional data augmentation, and the addition of 5-fold cross-validation. ","metadata":{}},{"cell_type":"code","source":"import os, json, joblib, numpy as np, pandas as pd\nimport random, math\nfrom pathlib import Path\nimport warnings \nwarnings.filterwarnings(\"ignore\")\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\n\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import StandardScaler, LabelEncoder\nfrom sklearn.utils.class_weight import compute_class_weight\nfrom sklearn.model_selection import StratifiedKFold\nfrom timm.scheduler import CosineLRScheduler\nfrom scipy.signal import firwin\nfrom cmi_2025_metric_copy_for_import import CompetitionMetric\n\nimport polars as pl\n# Configuration\nTRAIN = False                     # ← set to True when you want to train\nRAW_DIR = Path(\"../input/cmi-detect-behavior-with-sensor-data\")\nPRETRAINED_DIR = Path(\"/kaggle/input/cmi3-models-p\") # used when TRAIN=False\nEXPORT_DIR = Path(\"./\")                                    # artefacts will be saved here\nBATCH_SIZE = 64\nPAD_PERCENTILE = 100\nmaxlen = PAD_PERCENTILE\nLR_INIT = 1e-3\nWD = 3e-3\n# MIXUP_ALPHA = 0.4\nPATIENCE = 40\nFOLDS = 5\nrandom_state = 42\nepochs_warmup = 20\nwarmup_lr_init = 1.822126131809773e-05\nlr_min = 3.810323058740104e-09\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"▶ imports ready · pytorch {torch.__version__} · device: {device}\")\n\n# ================================\n# Model Components\n# ================================\nmean = torch.tensor([\n    0,  0, 0, 0, 0,\n    0,  9.0319e-03,  1.0849e+00, -2.6186e-03,  3.7651e-03,\n    -5.3660e-03, -2.8177e-03,  1.3318e-03, -1.5876e-04,  6.3495e-01,\n     6.2877e-01,  6.0607e-01,  6.2142e-01,  6.3808e-01,  6.5420e-01,\n     7.4102e-03, -3.4159e-03, -7.5237e-03, -2.6034e-02,  2.9704e-02,\n    -3.1546e-02, -2.0610e-03, -4.6986e-03, -4.7216e-03, -2.6281e-02,\n     1.5799e-02,  1.0016e-02\n], dtype=torch.float32).view(1, -1, 1).to(device)         \n\nstd = torch.tensor([\n    1, 1, 1, 1, 1, 1, 0.2067, 0.8583, 0.3162,\n    0.2668, 0.2917, 0.2341, 0.3023, 0.3281, 1.0264, 0.8838, 0.8686, 1.0973,\n    1.0267, 0.9018, 0.4658, 0.2009, 0.2057, 1.2240, 0.9535, 0.6655, 0.2941,\n    0.3421, 0.8156, 0.6565, 1.1034, 1.5577\n], dtype=torch.float32).view(1, -1, 1).to(device) + 1e-8  \n\nclass ImuFeatureExtractor(nn.Module):\n    def __init__(self, fs=100., add_quaternion=False):\n        super().__init__()\n        self.fs = fs\n        self.add_quaternion = add_quaternion\n\n        k = 15\n        self.lpf = nn.Conv1d(6, 6, kernel_size=k, padding=k//2,\n                             groups=6, bias=False)\n        nn.init.kaiming_uniform_(self.lpf.weight, a=math.sqrt(5))\n\n        self.lpf_acc  = nn.Conv1d(3, 3, k, padding=k//2, groups=3, bias=False)\n        self.lpf_gyro = nn.Conv1d(3, 3, k, padding=k//2, groups=3, bias=False)\n\n    def forward(self, imu):\n        # imu: \n        B, C, T = imu.shape\n        acc  = imu[:, 0:3, :]                 # acc_x, acc_y, acc_z\n        gyro = imu[:, 3:6, :]                 # gyro_x, gyro_y, gyro_z\n        extra = imu[:, 6:, :]                 \n\n        # 1) magnitude\n        acc_mag  = torch.norm(acc,  dim=1, keepdim=True)          # (B,1,T)\n        gyro_mag = torch.norm(gyro, dim=1, keepdim=True)\n\n        # 2) jerk \n        jerk = F.pad(acc[:, :, 1:] - acc[:, :, :-1], (1,0))       # (B,3,T)\n        gyro_delta = F.pad(gyro[:, :, 1:] - gyro[:, :, :-1], (1,0))\n\n        # 3) energy\n        acc_pow  = acc ** 2\n        gyro_pow = gyro ** 2\n\n        # 4) LPF / HPF \n        acc_lpf  = self.lpf_acc(acc)\n        acc_hpf  = acc - acc_lpf\n        gyro_lpf = self.lpf_gyro(gyro)\n        gyro_hpf = gyro - gyro_lpf\n\n        features = [\n            acc, gyro,\n            acc_mag, gyro_mag,\n            jerk, gyro_delta,\n            acc_pow, gyro_pow,\n            acc_lpf, acc_hpf,\n            gyro_lpf, gyro_hpf,\n        ]\n        return torch.cat(features, dim=1)  # (B, C_out, T)\n\n\nclass SEBlock(nn.Module):\n    def __init__(self, channels, reduction=8):\n        super().__init__()\n        self.squeeze = nn.AdaptiveAvgPool1d(1)\n        self.excitation = nn.Sequential(\n            nn.Linear(channels, channels // reduction, bias=False),\n            nn.ReLU(inplace=True),\n            nn.Linear(channels // reduction, channels, bias=False),\n            nn.Sigmoid()\n        )\n    \n    def forward(self, x):\n        b, c, _ = x.size()\n        y = self.squeeze(x).view(b, c)\n        y = self.excitation(y).view(b, c, 1)\n        return x * y.expand_as(x)\n\nclass ResidualSECNNBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size, pool_size=2, dropout=0.3, weight_decay=1e-4):\n        super().__init__()\n        \n        # First conv block\n        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, padding=kernel_size//2, bias=False)\n        self.bn1 = nn.BatchNorm1d(out_channels)\n        \n        # Second conv block\n        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, padding=kernel_size//2, bias=False)\n        self.bn2 = nn.BatchNorm1d(out_channels)\n        \n        # SE block\n        self.se = SEBlock(out_channels)\n        \n        # Shortcut connection\n        self.shortcut = nn.Sequential()\n        if in_channels != out_channels:\n            self.shortcut = nn.Sequential(\n                nn.Conv1d(in_channels, out_channels, 1, bias=False),\n                nn.BatchNorm1d(out_channels)\n            )\n        \n        self.pool = nn.MaxPool1d(pool_size)\n        self.dropout = nn.Dropout(dropout)\n        \n    def forward(self, x):\n        shortcut = self.shortcut(x)\n        \n        # First conv\n        out = F.relu(self.bn1(self.conv1(x)))\n        # Second conv\n        out = self.bn2(self.conv2(out))\n        \n        # SE block\n        out = self.se(out)\n        \n        # Add shortcut\n        out += shortcut\n        out = F.relu(out)\n        \n        # Pool and dropout\n        out = self.pool(out)\n        out = self.dropout(out)\n        \n        return out\n\nclass AttentionLayer(nn.Module):\n    def __init__(self, hidden_dim):\n        super().__init__()\n        self.attention = nn.Linear(hidden_dim, 1)\n        \n    def forward(self, x):\n        # x shape: (batch, seq_len, hidden_dim)\n        scores = torch.tanh(self.attention(x))  # (batch, seq_len, 1)\n        weights = F.softmax(scores.squeeze(-1), dim=1)  # (batch, seq_len)\n        context = torch.sum(x * weights.unsqueeze(-1), dim=1)  # (batch, hidden_dim)\n        return context\n\nclass TwoBranchModel(nn.Module):\n    def __init__(self, pad_len, imu_dim_raw, tof_dim, n_classes, dropouts=[0.3, 0.3, 0.3, 0.3, 0.4, 0.5, 0.3], feature_engineering=True, **kwargs):\n        super().__init__()\n        self.feature_engineering = feature_engineering\n        if feature_engineering:\n            self.imu_fe = ImuFeatureExtractor(**kwargs)\n            imu_dim = 32            \n        else:\n            self.imu_fe = nn.Identity()\n            imu_dim = imu_dim_raw   \n            \n        self.imu_dim = imu_dim\n        self.tof_dim = tof_dim\n\n        self.fir_nchan = 7\n\n        weight_decay = 3e-3\n\n        numtaps = 33  \n        fir_coef = firwin(numtaps, cutoff=1.0, fs=10.0, pass_zero=False)\n        fir_kernel = torch.tensor(fir_coef, dtype=torch.float32).view(1, 1, -1)\n        fir_kernel = fir_kernel.repeat(7, 1, 1)  # (imu_dim, 1, numtaps)\n        self.register_buffer(\"fir_kernel\", fir_kernel)\n        \n        # IMU deep branch\n        self.imu_block1 = ResidualSECNNBlock(imu_dim, 64, 3, dropout=dropouts[0], weight_decay=weight_decay)\n        self.imu_block2 = ResidualSECNNBlock(64, 128, 5, dropout=dropouts[1], weight_decay=weight_decay)\n        \n        # TOF/Thermal lighter branch\n        self.tof_conv1 = nn.Conv1d(tof_dim, 64, 3, padding=1, bias=False)\n        self.tof_bn1 = nn.BatchNorm1d(64)\n        self.tof_pool1 = nn.MaxPool1d(2)\n        self.tof_drop1 = nn.Dropout(dropouts[2])\n        \n        self.tof_conv2 = nn.Conv1d(64, 128, 3, padding=1, bias=False)\n        self.tof_bn2 = nn.BatchNorm1d(128)\n        self.tof_pool2 = nn.MaxPool1d(2)\n        self.tof_drop2 = nn.Dropout(dropouts[3])\n        \n        # BiLSTM\n        self.bilstm = nn.LSTM(256, 128, bidirectional=True, batch_first=True)\n        self.lstm_dropout = nn.Dropout(dropouts[4])\n        \n        # Attention\n        self.attention = AttentionLayer(256)  # 128*2 for bidirectional\n        \n        # Dense layers\n        self.dense1 = nn.Linear(256, 256, bias=False)\n        self.bn_dense1 = nn.BatchNorm1d(256)\n        self.drop1 = nn.Dropout(dropouts[5])\n        \n        self.dense2 = nn.Linear(256, 128, bias=False)\n        self.bn_dense2 = nn.BatchNorm1d(128)\n        self.drop2 = nn.Dropout(dropouts[6])\n        \n        self.classifier = nn.Linear(128, n_classes)\n        \n    def forward(self, x):\n        # Split input\n        \n        imu = x[:, :, :self.fir_nchan].transpose(1, 2)  # (batch, imu_dim, seq_len)\n        tof = x[:, :, self.fir_nchan:].transpose(1, 2)  # (batch, tof_dim, seq_len)\n\n        imu = self.imu_fe(imu)   # (B, imu_dim, T)\n        filtered = F.conv1d(\n            imu[:, :self.fir_nchan, :],        # (B,7,T)\n            self.fir_kernel,\n            padding=self.fir_kernel.shape[-1] // 2,\n            groups=self.fir_nchan,\n        )\n        \n        imu = torch.cat([filtered, imu[:, self.fir_nchan:, :]], dim=1)  \n        imu = (imu - mean) / std \n        # IMU branch\n        x1 = self.imu_block1(imu)\n        x1 = self.imu_block2(x1)\n        \n        # TOF branch\n        x2 = F.relu(self.tof_bn1(self.tof_conv1(tof)))\n        x2 = self.tof_drop1(self.tof_pool1(x2))\n        x2 = F.relu(self.tof_bn2(self.tof_conv2(x2)))\n        x2 = self.tof_drop2(self.tof_pool2(x2))\n        \n        # Concatenate branches\n        merged = torch.cat([x1, x2], dim=1).transpose(1, 2)  # (batch, seq_len, 256)\n        \n        # BiLSTM\n        lstm_out, _ = self.bilstm(merged)\n        lstm_out = self.lstm_dropout(lstm_out)\n        \n        # Attention\n        attended = self.attention(lstm_out)\n        \n        # Dense layers\n        x = F.relu(self.bn_dense1(self.dense1(attended)))\n        x = self.drop1(x)\n        x = F.relu(self.bn_dense2(self.dense2(x)))\n        x = self.drop2(x)\n        \n        # Classification\n        logits = (self.classifier(x))\n        return logits\n\n\n# ================================\n# Data Handling\n# ================================\n\ndef pad_sequences_torch(sequences, maxlen, padding='post', truncating='post', value=0.0):\n    \"\"\"PyTorch equivalent of Keras pad_sequences\"\"\"\n    result = []\n    for seq in sequences:\n        if len(seq) >= maxlen:\n            if truncating == 'post':\n                seq = seq[:maxlen]\n            else:  # 'pre'\n                seq = seq[-maxlen:]\n        else:\n            pad_len = maxlen - len(seq)\n            if padding == 'post':\n                seq = np.concatenate([seq, np.full((pad_len, seq.shape[1]), value)])\n            else:  # 'pre'\n                seq = np.concatenate([np.full((pad_len, seq.shape[1]), value), seq])\n        result.append(seq)\n    return np.array(result, dtype=np.float32)\n\ndef preprocess_sequence(df_seq: pd.DataFrame, feature_cols: list, scaler: StandardScaler):\n    \"\"\"Normalizes and cleans the time series sequence\"\"\"\n    mat = df_seq[feature_cols].ffill().bfill().fillna(0).values\n    return scaler.transform(mat).astype('float32')\n\nclass CMI3Dataset(Dataset):\n    def __init__(self,\n                 X_list,\n                 y_list,\n                 maxlen,\n                 mode=\"train\",\n                 imu_dim=7,\n                 augment=None):\n        self.X_list = X_list\n        self.mode = mode\n        self.y_list = y_list\n        self.maxlen = maxlen\n        self.imu_dim = imu_dim     \n        self.augment = augment   \n\n    def pad_sequences_torch(self, seq, maxlen, padding='post', truncating='post', value=0.0):\n\n        if seq.shape[0] >= maxlen:\n            if truncating == 'post':\n                seq = seq[:maxlen]\n            else:  # 'pre'\n                seq = seq[-maxlen:]\n        else:\n            pad_len = maxlen - seq.shape[0]\n            if padding == 'post':\n                seq = np.concatenate([seq, np.full((pad_len, seq.shape[1]), value)])\n            else:  # 'pre'\n                seq = np.concatenate([np.full((pad_len, seq.shape[1]), value), seq])\n        return seq  \n        \n    def __getitem__(self, index):\n        X = self.X_list[index]\n        y = self.y_list[index]\n\n        # ---------- (A)  Augmentation ----------\n        if self.mode == \"train\" and self.augment is not None:\n            X = self.augment(X, self.imu_dim)     \n\n        X = self.pad_sequences_torch(X, self.maxlen, 'pre', 'pre')\n        return X, y\n    \n    def __len__(self):\n        return len(self.X_list)\n\n\nclass EarlyStopping:\n    \"\"\"Early stopping utility\"\"\"\n    def __init__(self, patience=7, min_delta=0, restore_best_weights=True):\n        self.patience = patience\n        self.min_delta = min_delta\n        self.restore_best_weights = restore_best_weights\n        self.best_loss = None\n        self.counter = 0\n        self.best_weights = None\n        \n    def __call__(self, val_loss, model):\n        if self.best_loss is None:\n            self.best_loss = val_loss\n            self.save_checkpoint(model)\n        elif val_loss < self.best_loss - self.min_delta:\n            self.best_loss = val_loss\n            self.counter = 0\n            self.save_checkpoint(model)\n        else:\n            self.counter += 1\n            \n        if self.counter >= self.patience:\n            if self.restore_best_weights:\n                model.load_state_dict(self.best_weights)\n            return True\n        return False\n    \n    def save_checkpoint(self, model):\n        self.best_weights = model.state_dict().copy()\n\nclass EMA:\n    def __init__(self, model, decay=0.999):\n        self.decay = decay\n        self.shadow = {}\n        self.backup = {}\n\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                self.shadow[name] = param.data.clone()\n\n    def update(self, model):\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                assert name in self.shadow\n                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]\n                self.shadow[name] = new_average.clone()\n\n    def apply_shadow(self, model):\n        self.backup = {}\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                self.backup[name] = param.data.clone()\n                param.data = self.shadow[name]\n\n    def restore(self, model):\n        for name, param in model.named_parameters():\n            if param.requires_grad and name in self.backup:\n                param.data = self.backup[name]\n        self.backup = {}\n\ndef set_seed(seed: int = 42):\n    random.seed(seed)\n\n    os.environ['PYTHONHASHSEED'] = str(seed)\n\n    np.random.seed(seed)\n\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed) \n    # torch.backends.cudnn.deterministic = True\n    # torch.backends.cudnn.benchmark = False\n    # torch.use_deterministic_algorithms(True)\n\nclass Augment:\n    def __init__(self,\n                 p_jitter=0.8, sigma=0.02, scale_range=[0.9,1.1],\n                 p_dropout=0.3,\n                 p_moda=0.5,          \n                 drift_std=0.005,     \n                 drift_max=0.25):      \n        self.p_jitter  = p_jitter\n        self.sigma     = sigma\n        self.scale_min, self.scale_max = scale_range\n        self.p_dropout = p_dropout\n        self.p_moda    = p_moda\n        self.drift_std = drift_std\n        self.drift_max = drift_max\n\n    # ---------- Jitter & Scaling ----------\n    def jitter_scale(self, x: np.ndarray) -> np.ndarray:\n        noise  = np.random.randn(*x.shape) * self.sigma\n        scale  = np.random.uniform(self.scale_min,\n                                   self.scale_max,\n                                   size=(1, x.shape[1]))\n        return (x + noise) * scale\n\n    # ---------- Sensor Drop-out ----------\n    def sensor_dropout(self,\n                       x: np.ndarray,\n                       imu_dim: int) -> np.ndarray:\n\n        if random.random() < self.p_dropout:\n            x[:, imu_dim:] = 0.0\n        return x\n\n    def motion_drift(self, x: np.ndarray, imu_dim: int) -> np.ndarray:\n\n        T = x.shape[0]\n\n        drift = np.cumsum(\n            np.random.normal(scale=self.drift_std, size=(T, 1)),\n            axis=0\n        )\n        drift = np.clip(drift, -self.drift_max, self.drift_max)   \n\n        x[:, :6] += drift\n\n        if imu_dim > 6:\n            x[:, 6:imu_dim] += drift     \n        return x\n    \n    # ---------- master call ----------\n    def __call__(self,\n                 x: np.ndarray,\n                 imu_dim: int) -> np.ndarray:\n        if random.random() < self.p_jitter:\n            x = self.jitter_scale(x)\n\n        if random.random() < self.p_moda:\n            x = self.motion_drift(x, imu_dim)\n\n        x = self.sensor_dropout(x, imu_dim)\n        return x\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-30T10:14:56.738128Z","iopub.execute_input":"2025-06-30T10:14:56.738691Z","iopub.status.idle":"2025-06-30T10:14:56.788074Z","shell.execute_reply.started":"2025-06-30T10:14:56.738668Z","shell.execute_reply":"2025-06-30T10:14:56.787219Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ================================\n# Training Pipeline\n# ================================\nif TRAIN:\n    print(\"▶ TRAIN MODE – loading dataset …\")\n    df = pd.read_csv(RAW_DIR / \"train.csv\")\n\n    # Label encoding\n    le = LabelEncoder()\n    df['gesture_int'] = le.fit_transform(df['gesture'])\n    np.save(EXPORT_DIR / \"gesture_classes.npy\", le.classes_)\n\n    # Feature list\n    meta_cols = {'gesture', 'gesture_int', 'sequence_type', 'behavior', 'orientation',\n                 'row_id', 'subject', 'phase', 'sequence_id', 'sequence_counter'}\n    feature_cols = [c for c in df.columns if c not in meta_cols]\n\n    imu_cols = [c for c in feature_cols if not (c.startswith('thm_') or c.startswith('tof_'))]\n    tof_cols = [c for c in feature_cols if c.startswith('thm_') or c.startswith('tof_')]\n    print(f\"  IMU {len(imu_cols)} | TOF/THM {len(tof_cols)} | total {len(feature_cols)} features\")\n\n    # Global scaler\n    scaler = StandardScaler().fit(df[feature_cols].ffill().bfill().fillna(0).values)\n    joblib.dump(scaler, EXPORT_DIR / \"scaler.pkl\")\n\n    # Build sequences\n    seq_gp = df.groupby('sequence_id')\n    X_list, y_list, id_list = [], [], []\n    for seq_id, seq in seq_gp:\n        mat = preprocess_sequence(seq, feature_cols, scaler)\n        X_list.append(mat)\n        y_list.append(seq['gesture_int'].iloc[0])\n        id_list.append(seq_id)\n        # lens.append(len(mat))\n    \n    pad_len = PAD_PERCENTILE#int(np.percentile(lens, PAD_PERCENTILE))\n    print(pad_len)\n    np.save(EXPORT_DIR / \"sequence_maxlen.npy\", pad_len)\n    np.save(EXPORT_DIR / \"feature_cols.npy\", np.array(feature_cols))\n    id_list = np.array(id_list)\n    X_list_all = pad_sequences_torch(X_list, maxlen=pad_len, padding='pre', truncating='pre')\n    y_list_all = np.eye(len(le.classes_))[y_list].astype(np.float32)  # One-hot encoding\n\n    augmenter = Augment(\n        p_jitter=0.9844818619033621, sigma=0.03291295776089293, scale_range=(0.7542342630597011,1.1625052821731077),\n        p_dropout=0.41782786013520684,\n        p_moda=0.3910622476959722, drift_std=0.0040285239353308015, drift_max=0.3929358950258158    \n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-30T10:14:56.789557Z","iopub.execute_input":"2025-06-30T10:14:56.789787Z","iopub.status.idle":"2025-06-30T10:14:56.810949Z","shell.execute_reply.started":"2025-06-30T10:14:56.789770Z","shell.execute_reply":"2025-06-30T10:14:56.810175Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"EPOCHS = 125\nif TRAIN:\n    # Split\n    skf = StratifiedKFold(n_splits=FOLDS, shuffle=True, random_state=random_state)\n    models = []\n    \n    for fold, (train_idx, val_idx) in enumerate(skf.split(id_list, np.argmax(y_list_all, axis=1))):\n\n        train_list= X_list_all[train_idx]\n        train_y_list= y_list_all[train_idx]\n        val_list = X_list_all[val_idx]\n        val_y_list= y_list_all[val_idx]\n\n        \n        # Data loaders\n        train_dataset = CMI3Dataset(train_list, train_y_list, maxlen, mode=\"train\", imu_dim=len(imu_cols),\n                                augment=augmenter)\n        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4,drop_last=True)\n    \n        val_dataset = CMI3Dataset(val_list, val_y_list, maxlen, mode=\"val\")\n        val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4,drop_last=True)\n\n    \n        # Model\n        model = TwoBranchModel(maxlen, len(imu_cols), len(tof_cols), \n                      len(le.classes_)).to(device)\n        ema = EMA(model, decay=0.999)\n        # Optimizer and scheduler\n        optimizer = Adam(model.parameters(), lr=LR_INIT, weight_decay=WD)\n        \n        steps_per_epoch = len(train_loader)\n        nbatch = len(train_loader)\n        warmup = epochs_warmup * nbatch\n        nsteps = EPOCHS * nbatch\n        scheduler = CosineLRScheduler(optimizer,\n                          warmup_t=warmup, warmup_lr_init=warmup_lr_init, warmup_prefix=True,\n                          t_initial=(nsteps - warmup), lr_min=lr_min) \n    \n        early_stopping = EarlyStopping(patience=PATIENCE, restore_best_weights=True)\n    \n        train_loss = 0.0\n        train_acc = 0.0\n        val_loss = 0.0\n        val_acc = 0.0\n        val_best_acc = 0.0\n        i_scheduler = 0\n        \n        # Training loop\n        print(\"▶ Starting training...\")\n        for epoch in range(EPOCHS):\n            model.train()\n            train_preds = []\n            train_targets = []\n            for X, y in (train_loader):  \n                X, y = X.float().to(device), y.to(device)\n                optimizer.zero_grad()\n                logits = model(X)\n    \n                loss = -torch.sum(F.log_softmax(logits, dim=1) * y, dim=1).mean()\n                loss.backward()\n                optimizer.step()\n                ema.update(model)\n                train_preds.extend(logits.argmax(dim=1).cpu().numpy())\n                train_targets.extend(y.argmax(dim=1).cpu().numpy())\n                scheduler.step(i_scheduler)\n                i_scheduler +=1\n    \n                train_loss += loss.item()\n                \n            model.eval()\n            with torch.inference_mode():\n                val_preds = []\n                val_targets = []\n                for X, y in (val_loader):  \n                    half = BATCH_SIZE // 2         \n\n                    x_front = X[:half]               \n                    x_back  = X[half:].clone()      \n                    \n                    x_back[:, :, 7:] = 0.0    \n                    X = torch.cat([x_front, x_back], dim=0)  # (B, C, T)\n                    X, y = X.float().to(device), y.to(device)\n                    \n                    logits = model(X)\n                    val_preds.extend(logits.argmax(dim=1).cpu().numpy())\n                    val_targets.extend(y.argmax(dim=1).cpu().numpy())\n                    \n                    loss = F.cross_entropy(logits, y)\n                    val_loss += loss.item()\n    \n            train_acc = CompetitionMetric().calculate_hierarchical_f1(\n                pd.DataFrame({'gesture': le.classes_[train_targets]}),\n                pd.DataFrame({'gesture': le.classes_[train_preds]}))\n            val_acc = CompetitionMetric().calculate_hierarchical_f1(\n                pd.DataFrame({'gesture': le.classes_[val_targets]}),\n                pd.DataFrame({'gesture': le.classes_[val_preds]}))\n            train_loss /= len(train_loader)\n            val_loss /= len(val_loader)\n        models.append(model)\n        # Save model\n        torch.save({\n            'model_state_dict': model.state_dict(),\n            'imu_dim': len(imu_cols),\n            'tof_dim': len(tof_cols),\n            'n_classes': len(le.classes_),\n            'pad_len': pad_len\n        }, EXPORT_DIR / f\"gesture_two_branch_fold{fold}.pth\")\n        print(f\"fold: {fold} val_all_acc: {val_acc:.4f}\")\n        print(\"✔ Training done – artefacts saved in\", EXPORT_DIR)\n\nelse:\n    print(\"▶ INFERENCE MODE – loading artefacts from\", PRETRAINED_DIR)\n    feature_cols = np.load(PRETRAINED_DIR / \"feature_cols.npy\", allow_pickle=True).tolist()\n    pad_len = int(np.load(PRETRAINED_DIR / \"sequence_maxlen.npy\"))\n    scaler = joblib.load(PRETRAINED_DIR / \"scaler.pkl\")\n    gesture_classes = np.load(PRETRAINED_DIR / \"gesture_classes.npy\", allow_pickle=True)\n\n    imu_cols = [c for c in feature_cols if not (c.startswith('thm_') or c.startswith('tof_'))]\n    tof_cols = [c for c in feature_cols if c.startswith('thm_') or c.startswith('tof_')]\n\n    \n    # Load model\n    MODELS = [f'gesture_two_branch_fold{i}.pth' for i in range(5)]\n    \n    models = []\n    for path in MODELS:\n        checkpoint = torch.load(PRETRAINED_DIR / path, map_location=device)\n        \n        model = TwoBranchModel(\n            checkpoint['pad_len'], \n            checkpoint['imu_dim'], \n            checkpoint['tof_dim'], \n            checkpoint['n_classes']\n            ).to(device)\n        \n        model.load_state_dict(checkpoint['model_state_dict'])\n        model.eval()\n        models.append(model)\n\n    print(\"  model, scaler, pads loaded – ready for evaluation\")\n\n# Make sure gesture_classes exists in both modes\nif TRAIN:\n    gesture_classes = le.classes_\n\ndef predict(sequence: pl.DataFrame, demographics: pl.DataFrame) -> str:\n    \"\"\"Prediction function for Kaggle competition\"\"\"\n    global gesture_classes\n    if gesture_classes is None:\n        gesture_classes = np.load(PRETRAINED_DIR / \"gesture_classes.npy\", allow_pickle=True)\n\n    df_seq = sequence.to_pandas()\n    mat = preprocess_sequence(df_seq, feature_cols, scaler)\n    pad = pad_sequences_torch([mat], maxlen=pad_len, padding='pre', truncating='pre')\n    \n    with torch.no_grad():\n        x = torch.FloatTensor(pad).to(device)\n        outputs = None\n        for model in models:\n            model.eval()\n            p = torch.softmax(model(x), dim=1)\n            if outputs is None: outputs = p\n            else: outputs += p\n        outputs /= len(models)\n        \n        idx = int(outputs.argmax(dim=1)[0].cpu().numpy())\n    \n    return str(gesture_classes[idx])\n\n# Kaggle competition interface\nimport kaggle_evaluation.cmi_inference_server\ninference_server = kaggle_evaluation.cmi_inference_server.CMIInferenceServer(predict)\n\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    inference_server.serve()\nelse:\n    inference_server.run_local_gateway(\n        data_paths=(\n            '/kaggle/input/cmi-detect-behavior-with-sensor-data/test.csv',\n            '/kaggle/input/cmi-detect-behavior-with-sensor-data/test_demographics.csv',\n        )\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-30T10:14:56.811782Z","iopub.execute_input":"2025-06-30T10:14:56.812023Z","iopub.status.idle":"2025-06-30T10:15:00.456048Z","shell.execute_reply.started":"2025-06-30T10:14:56.812004Z","shell.execute_reply":"2025-06-30T10:15:00.455190Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}