{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":102335,"databundleVersionId":12518947,"sourceType":"competition"},{"sourceId":543540,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":418271,"modelId":430123},{"sourceId":539596,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":417626,"modelId":430123}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"341558b0","cell_type":"code","source":"%%capture\n!pip install azureml-mlflow","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-23T07:25:53.510997Z","iopub.execute_input":"2025-08-23T07:25:53.511256Z","iopub.status.idle":"2025-08-23T07:26:03.486162Z","shell.execute_reply.started":"2025-08-23T07:25:53.511229Z","shell.execute_reply":"2025-08-23T07:26:03.484899Z"}},"outputs":[],"execution_count":null},{"id":"6b1b109b","cell_type":"code","source":"import pandas as pd\nimport polars as pl\nimport numpy as np\nimport math\nfrom tqdm import tqdm\nfrom typing import Optional, Tuple, List\nfrom scipy.fft import rfft\nfrom kaggle_secrets import UserSecretsClient\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.model_selection import StratifiedGroupKFold\nimport mlflow\nimport os\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import ExponentialLR, OneCycleLR\nfrom torch.utils.data import DataLoader\nfrom torchmetrics import F1Score\nimport kaggle_evaluation.cmi_inference_server","metadata":{"execution":{"iopub.status.busy":"2025-08-23T07:26:27.098407Z","iopub.execute_input":"2025-08-23T07:26:27.099152Z","iopub.status.idle":"2025-08-23T07:26:27.103908Z","shell.execute_reply.started":"2025-08-23T07:26:27.099127Z","shell.execute_reply":"2025-08-23T07:26:27.103292Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"61ff989b-7e3c-40da-8b3b-976ae8577425","cell_type":"code","source":"user_secrets = UserSecretsClient()\nmlflow_tracking_uri = user_secrets.get_secret(\"MLFLOW_TRACKING_URI\")\nmlflow.set_tracking_uri(mlflow_tracking_uri)","metadata":{"execution":{"iopub.status.busy":"2025-08-23T07:26:29.713374Z","iopub.execute_input":"2025-08-23T07:26:29.713669Z","iopub.status.idle":"2025-08-23T07:26:29.908666Z","shell.execute_reply.started":"2025-08-23T07:26:29.713645Z","shell.execute_reply":"2025-08-23T07:26:29.907895Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"52caa64e-c1dc-4d67-bc8e-c489462f73d4","cell_type":"code","source":"os.environ[\"AZURE_CLIENT_ID\"] = user_secrets.get_secret(\"AZURE_CLIENT_ID\")\nos.environ[\"AZURE_CLIENT_SECRET\"] = user_secrets.get_secret(\"AZURE_CLIENT_SECRET\")\nos.environ[\"AZURE_TENANT_ID\"] = user_secrets.get_secret(\"AZURE_TENANT_ID\")","metadata":{"execution":{"iopub.status.busy":"2025-08-23T07:26:31.51082Z","iopub.execute_input":"2025-08-23T07:26:31.511532Z","iopub.status.idle":"2025-08-23T07:26:31.991658Z","shell.execute_reply.started":"2025-08-23T07:26:31.511504Z","shell.execute_reply":"2025-08-23T07:26:31.991109Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"51281bcf-8ac5-4ec0-bfaa-5064195693a6","cell_type":"code","source":"mlflow.set_experiment(\"best_sadeepsense_model\")\nmlflow.start_run(run_name=\"decrease_latent_space_strack_gru_3\")","metadata":{"execution":{"iopub.status.busy":"2025-08-23T08:09:43.047412Z","iopub.execute_input":"2025-08-23T08:09:43.048023Z","iopub.status.idle":"2025-08-23T08:09:44.910838Z","shell.execute_reply.started":"2025-08-23T08:09:43.04799Z","shell.execute_reply":"2025-08-23T08:09:44.910137Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"218645d8","cell_type":"markdown","source":"# Preprocessing data","metadata":{}},{"id":"6326ec9e-473c-4e80-8dee-95a953983981","cell_type":"code","source":"class CMIDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, config: Optional[dict] = None):\n        '''\n        Read and preprocess the dataset.\n\n        Args:\n            df (pd.DataFrame): DataFrame containing the dataset.\n            config (Optional[dict]): Configuration dictionary with preprocessing options.\n        \n        '''\n        self.df = df\n\n        if config is not None:\n            self.SEQUENCE_LEN = config.get('SEQUENCE_LEN', 512) # Maximum length of each sequence\n            self.N_SAMPLES_IN_CHUNK = config.get('N_SAMPLES_IN_CHUNK', 5) # Number of samples in each chunk\n            self.NAN_VALUE = config.get('NAN_VALUES', 0.0) # Value to fill missing data\n            self.TARGET_COL = config.get('TARGET_COL', None) # Target column for classification\n            self.TARGET_ENCODER = config.get('TARGET_ENCODER', LabelEncoder()) # Encoder for target column\n            self.SAMPLING_RATE = 10 # Sampling rate in Hz\n            self.N_TOFs = config.get('N_TOFs', 5) # Number of TOF sensors\n            self.MAX_NCHUNKS = config.get('MAX_NCHUNKS', 120) # Maximum number of chunks per sequence\n            self.NROWS = config.get('NROWS', None) # Number of rows to read from the dataset\n\n        self._init_feature_names(self.df.columns)\n\n        if self.TARGET_COL is not None:\n            print(\"Encoding target column:\", self.TARGET_COL)\n            self.df[self.TARGET_COL] = self.TARGET_ENCODER.fit_transform(self.df[self.TARGET_COL])\n\n        self.acc_data, self.gyr_data, self.thrm_data, self.tof_data, self.sequence_ids, self.n_chunks, self.labels = self._process_all_sequences(self.df)\n\n\n    def _init_feature_names(self, columns: pd.Index) -> None:\n        '''Initialize feature names for the dataset.'''\n\n        self.acc_cols = ['acc_x', 'acc_y', 'acc_z']\n        self.gyr_cols = ['rot_x', 'rot_y', 'rot_z', 'rot_w']\n        self.thrm_cols = [c for c in columns if c.startswith('thm')]\n        self.tof_cols = [c for c in columns if c.startswith('tof')]\n\n    def _process_all_sequences(self, df:pd.DataFrame) -> Tuple[List, List, List, List, List, List, List]:\n        '''\n        Process all sequences in the dataset to extract frequency domain features and reshape TOF data.\n        It creates arrays for accelerometer, gyroscope, thermometer, and TOF data, along with sequence IDs and lengths.\n        acc_data shape: (N_SEQUENCES, N_CHUNKS, ACC_CHANNELS, 2*f)\n        gyr_data shape: (N_SEQUENCES, N_CHUNKS, GYR_CHANNELS, 2*f)\n        thrm_data shape: (N_SEQUENCES, N_CHUNKS, THRM_CHANNELS, 2*f)\n        tof_data shape: (N_SEQUENCES, N_CHUNKS, N_TOFs, 8, 8)\n        '''\n\n        # Create arrays to hold the processed data\n        acc_data_list = [] \n        gyr_data_list = []\n        thrm_data_list = []\n        tof_data_list = []\n        sequence_id_list = []\n        n_chunk_list = []\n        label_list = []\n\n        print(\"Processing sequences...\")\n        for (sequence_id, ), sequence_data in tqdm(df.groupby([\"sequence_id\"])): \n            # Fill missing data for each sequence\n            sequence_data = sequence_data.ffill(axis=0).bfill(axis=0).fillna(self.NAN_VALUE)\n\n            # Process each sequence\n            acc_from_seq, gyr_from_seq, thrm_from_seq, tof_from_seq, n_chunk_seq = self._process_sequence(sequence_data)\n            acc_data_list.append(acc_from_seq)\n            gyr_data_list.append(gyr_from_seq)\n            thrm_data_list.append(thrm_from_seq)\n            tof_data_list.append(tof_from_seq)\n            sequence_id_list.append(sequence_id)\n            n_chunk_list.append(n_chunk_seq)\n\n            if self.TARGET_COL is not None:\n                label_list.append(sequence_data[self.TARGET_COL].values[0])\n\n        return acc_data_list, gyr_data_list, thrm_data_list, tof_data_list, sequence_id_list, n_chunk_list, label_list\n\n    def _process_sequence(self, sequence_data: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:\n        \"\"\"\n        Process a single sequence of sensor data.\n        Convert accelerometer, gyroscope, thermometer data to frequency domain features\n        Average each TOF sensor's data on each chunk.\n\n        Args:\n        sequence_data: pd.DataFrame\n            DataFrame containing the sensor data for a single sequence.\n\n        Returns:\n        acc_seq_freq: np.ndarray, shape (N_CHUNKS, ACC_CHANNELS, 2*f)\n            Frequency domain features for accelerometer data.\n        gyr_seq_freq: np.ndarray, shape (N_CHUNKS, GYR_CHANNELS, 2*f)\n            Frequency domain features for gyroscope data.\n        thrm_seq_freq: np.ndarray, shape (N_CHUNKS, THRM_CHANNELS, 2*f)\n            Frequency domain features for thermometer data.\n        tof_seq: np.ndarray, shape (N_CHUNKS, N_TOFs, 8, 8)\n            Reshaped TOF data.\n        num_chunks: int\n            Number of chunks created from the sequence data.\n        \"\"\"\n        # Sort by sequence_counter and reset_index\n        sequence_data = sequence_data.sort_values(by=\"sequence_counter\", ascending=True).reset_index()\n\n        # Ensure the sequence data has maximum length of SEQUENCE_LEN\n        if sequence_data.shape[0] > self.SEQUENCE_LEN:\n            sequence_data = sequence_data.iloc[:self.SEQUENCE_LEN, :]\n\n        # If the number of samples in the last chunk is not equal to N_SAMPLES_IN_CHUNK, pad the data\n        if sequence_data.shape[0] % self.N_SAMPLES_IN_CHUNK != 0:\n            pad_data = np.zeros(shape=(self.N_SAMPLES_IN_CHUNK - (sequence_data.shape[0] % self.N_SAMPLES_IN_CHUNK), sequence_data.shape[1]))\n            sequence_data = pd.concat([sequence_data, pd.DataFrame(data=pad_data, columns=sequence_data.columns)], ignore_index=True)\n        \n        acc_seq_freq = []\n        gyr_seq_freq = []\n        thrm_seq_freq = []\n        tof_seq = []\n\n        # Divide the sequence data into chunks of N_SAMPLES_IN_CHUNK\n        # and convert each chunk to frequency domain features\n        for start_chunk_index in range(0, sequence_data.shape[0], self.N_SAMPLES_IN_CHUNK):\n            chunk_data = sequence_data.iloc[start_chunk_index:start_chunk_index+self.N_SAMPLES_IN_CHUNK]\n            \n            acc_chunk_freq = self._convert_freq_domain(chunk_data.loc[:, self.acc_cols])\n            gyr_chunk_freq = self._convert_freq_domain(chunk_data.loc[:, self.gyr_cols])\n            thrm_chunk_freq = self._convert_freq_domain(chunk_data.loc[:, self.thrm_cols])\n\n            # Reshape TOF data to have shape \n            # And average over the time dimension for each TOF sensor\n            tof_chunk = chunk_data.loc[:, self.tof_cols].values.reshape(self.N_SAMPLES_IN_CHUNK, self.N_TOFs, 8, 8)\n            tof_chunk = tof_chunk.transpose(1, 0, 2, 3).mean(axis=1)\n\n            # Append the frequency domain features to the respective lists\n            acc_seq_freq.append(acc_chunk_freq)\n            gyr_seq_freq.append(gyr_chunk_freq)\n            thrm_seq_freq.append(thrm_chunk_freq)\n            tof_seq.append(tof_chunk)\n\n        return np.array(acc_seq_freq), np.array(gyr_seq_freq), np.array(thrm_seq_freq), np.array(tof_seq), sequence_data.shape[0] // self.N_SAMPLES_IN_CHUNK\n\n    def _convert_freq_domain(self, chunk_data: pd.DataFrame) -> np.ndarray:\n        '''\n        Convert time domain sensor data to frequency domain using FFT.\n\n        Args:\n            chunk_data (pd.DataFrame): Time domain sensor data.\n\n        Returns:\n            np.ndarray: Frequency domain representation of the sensor data.\n        '''\n        fourier = rfft(chunk_data, axis=0)\n        amplitude = 2*np.abs(np.array(fourier, dtype=complex))/chunk_data.shape[0]\n        phase = np.angle(np.array(fourier, dtype=complex), deg=False)\n\n        return np.concatenate([amplitude, phase], axis=0).T\n        \n\n\n    def __getitem__(self, idx):\n        '''\n        Get a single item from the dataset and padding\n        '''\n        acc_data = self.acc_data[idx]\n        gyr_data = self.gyr_data[idx]\n        thrm_data = self.thrm_data[idx]\n        tof_data = self.tof_data[idx]\n        n_chunks = self.n_chunks[idx]\n        if self.TARGET_COL is not None:\n            label = self.labels[idx]\n        else:\n            label = None\n\n        if n_chunks > self.MAX_NCHUNKS:\n            acc_data = acc_data[:self.MAX_NCHUNKS]\n            gyr_data = gyr_data[:self.MAX_NCHUNKS]\n            thrm_data = thrm_data[:self.MAX_NCHUNKS]\n            tof_data = tof_data[:self.MAX_NCHUNKS]\n            n_chunks = self.MAX_NCHUNKS\n        else:\n            n_padded_chunks = self.MAX_NCHUNKS - n_chunks\n\n            pad_acc_data = np.full(shape=(n_padded_chunks, acc_data.shape[1], acc_data.shape[2]), fill_value=self.NAN_VALUE)\n            pad_gyr_data = np.full(shape=(n_padded_chunks, gyr_data.shape[1], gyr_data.shape[2]), fill_value=self.NAN_VALUE)\n            pad_thrm_data = np.full(shape=(n_padded_chunks, thrm_data.shape[1], thrm_data.shape[2]), fill_value=self.NAN_VALUE)\n            pad_tof_data = np.full(shape=(n_padded_chunks, tof_data.shape[1], tof_data.shape[2], tof_data.shape[3]), fill_value=self.NAN_VALUE)\n\n            acc_data = np.concatenate([acc_data, pad_acc_data])\n            gyr_data = np.concatenate([gyr_data, pad_gyr_data])\n            thrm_data = np.concatenate([thrm_data, pad_thrm_data])\n            tof_data = np.concatenate([tof_data, pad_tof_data])\n\n        return torch.tensor(n_chunks), \\\n                torch.tensor(label, dtype=torch.long), \\\n               torch.tensor(acc_data, dtype=torch.float32), \\\n               torch.tensor(gyr_data, dtype=torch.float32), \\\n               torch.tensor(thrm_data, dtype=torch.float32), \\\n               torch.tensor(tof_data, dtype=torch.float32)\n\n    \n    def __len__(self):\n        '''\n        Get the length of the dataset.\n        '''\n        return len(self.sequence_ids)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-08-23T07:26:42.411364Z","iopub.execute_input":"2025-08-23T07:26:42.411833Z","iopub.status.idle":"2025-08-23T07:26:42.433668Z","shell.execute_reply.started":"2025-08-23T07:26:42.411808Z","shell.execute_reply":"2025-08-23T07:26:42.433002Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"cbcd9e8e","cell_type":"code","source":"train_config = {\n    'SEQUENCE_LEN': 700,\n    'N_SAMPLES_IN_CHUNK': 5,\n    'NAN_VALUES': 0.0, \n    'TARGET_COL': 'gesture',\n    'TARGET_ENCODER': LabelEncoder(),\n    'SAMPLING_RATE': 10,\n    'N_TOFs': 5, \n    'MAX_NCHUNKS': 140, \n    'NROWS': None \n}\n\ndata_path = \"/kaggle/input/cmi-detect-behavior-with-sensor-data/train.csv\"\nprint(f\"Reading data from file: {data_path}\")\ndf = pd.read_csv(data_path, nrows=train_config['NROWS'])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-23T07:26:49.233545Z","iopub.execute_input":"2025-08-23T07:26:49.234203Z","iopub.status.idle":"2025-08-23T07:27:19.739308Z","shell.execute_reply.started":"2025-08-23T07:26:49.234179Z","shell.execute_reply":"2025-08-23T07:27:19.73838Z"}},"outputs":[],"execution_count":null},{"id":"046f3ea0","cell_type":"code","source":"sgkfolds = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=42)\nfor i, (train_index, test_index) in enumerate(sgkfolds.split(df, df.loc[:, train_config['TARGET_COL']], df.loc[:, 'sequence_id'])):\n    print(f\"Fold {i+1}:\")\n    train_df = df.iloc[train_index, :]\n    test_df = df.iloc[test_index, :]\n\n    train_dataset = CMIDataset(train_df, config=train_config)\n    train_config['TARGET_ENCODER'] = train_dataset.TARGET_ENCODER\n    test_dataset = CMIDataset(test_df, config=train_config)\n\n    print(f\"Train dataset size: {len(train_dataset)}\")\n    print(f\"Test dataset size: {len(test_dataset)}\")\n\n    break\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-23T07:27:19.740692Z","iopub.execute_input":"2025-08-23T07:27:19.740952Z","iopub.status.idle":"2025-08-23T07:31:08.441903Z","shell.execute_reply.started":"2025-08-23T07:27:19.740929Z","shell.execute_reply":"2025-08-23T07:31:08.441212Z"}},"outputs":[],"execution_count":null},{"id":"3e1cc077","cell_type":"markdown","source":"# Create model\n","metadata":{}},{"id":"b7be1a84-1769-4eee-b2b8-7d826c94b6c2","cell_type":"code","source":"class SensorBranch(nn.Module):\n    def __init__(self, drop_prob=0.2, sensor_channels=3, num_kernels=64):\n        super(SensorBranch, self).__init__()\n        self.DROP_PROB = drop_prob\n        self.SENSOR_CHANNELS = sensor_channels\n        self.NUM_KERNELS = num_kernels\n        self.KERNEL_LENGTH = 2\n\n        # Output shape of (BATCH_SIZE, N_CHUNKS, NUM_KERNELS*(2*F - 3))\n        self.output_shape = (None, None, self.NUM_KERNELS*(6-3))\n\n\n        self.seq = nn.Sequential(\n            nn.Conv2d(in_channels=1, out_channels=self.NUM_KERNELS, kernel_size=(self.SENSOR_CHANNELS, self.KERNEL_LENGTH), stride=(self.SENSOR_CHANNELS, 1)),\n            nn.BatchNorm2d(num_features=self.NUM_KERNELS),\n            nn.ELU(),\n            nn.Dropout(p=self.DROP_PROB),\n\n            nn.Conv2d(in_channels=self.NUM_KERNELS, out_channels=self.NUM_KERNELS, kernel_size=(1, self.KERNEL_LENGTH), stride=(1, 1)),\n            nn.BatchNorm2d(num_features=self.NUM_KERNELS),\n            nn.ELU(),\n            nn.Dropout(p=self.DROP_PROB),\n\n            nn.Conv2d(in_channels=self.NUM_KERNELS, out_channels=self.NUM_KERNELS, kernel_size=(1, self.KERNEL_LENGTH), stride=(1, 1)),\n            nn.BatchNorm2d(num_features=self.NUM_KERNELS),\n            nn.ELU(),\n            nn.Dropout(p=self.DROP_PROB)\n        )\n\n\n    def forward(self, x:torch.Tensor):\n        # x_shape: (batch_size, n_chunks, sensor_channels, 2*F)\n        x = torch.reshape(x, shape=(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])) \n        x = torch.unsqueeze(x, dim=1) # x_shape: (batch_size, 1,  n_chunks* sensor_channels, 2*F)\n        \n        # Conv branch\n        x = self.seq(x)\n        # output_shape_conv1: (BATCH_SIZE, N_CONV,  N_CHUNKS, 2*F - 1)\n        # output_shape_conv2: (BATCH_SIZE, N_CONV,  N_CHUNKS, 2*F - 2)\n        # output_shape_conv3: (BATCH_SIZE, N_CONV,  N_CHUNKS, 2*F - 3)\n\n        # Flattern the output \n        x = x.permute(0, 2, 1, 3) # x.permute: (BATCH_SIZE, N_CHUNKS, N_CONV, 2*F - 3)\n        x = x.reshape(shape=(x.shape[0], x.shape[1], x.shape[2]*x.shape[3])) \n        # x.reshape: (BATCH_SIZE, N_CHUNKS, N_CONV * (2*F - 3))\n        \n        return x\n    \nclass TOFBranch(nn.Module):\n    def __init__(self, drop_prob=0.2, num_kernels=64):\n        super(TOFBranch, self).__init__()\n        self.DROP_PROB = drop_prob\n        self.NUM_KERNELS = num_kernels\n\n        self.seq = nn.Sequential(\n            nn.Conv3d(in_channels=1, out_channels=self.NUM_KERNELS, kernel_size=(1, 3, 3)), \n            nn.BatchNorm3d(num_features=self.NUM_KERNELS),\n            nn.ELU(),\n            nn.Dropout(p=self.DROP_PROB),\n            #output_shape: [BATCH_SIZE, NUM_KERNELS, N_CHUNKS, 6, 6]\n\n            nn.Conv3d(in_channels=self.NUM_KERNELS, out_channels=self.NUM_KERNELS, kernel_size=(1, 3, 3)), \n            nn.BatchNorm3d(num_features=self.NUM_KERNELS),\n            nn.ELU(),\n            nn.Dropout(p=self.DROP_PROB),\n            #output_shape: [BATCH_SIZE, NUM_KERNELS, N_CHUNKS, 4, 4]\n\n            nn.Conv3d(in_channels=self.NUM_KERNELS, out_channels=self.NUM_KERNELS, kernel_size=(1, 4, 2)), \n            nn.BatchNorm3d(num_features=self.NUM_KERNELS),\n            nn.ELU(),\n            nn.Dropout(p=self.DROP_PROB)\n            #output_shape: [BATCH_SIZE, NUM_KERNELS, N_CHUNKS, 1, 3],\n        )\n        # Output shape after all convolutions and flattening: [BATCH_SIZE, N_CHUNKS, NUM_KERNELS * 3]\n        self.output_shape = (None, None, self.NUM_KERNELS * 3)\n\n    def forward(self, tofs_data:torch.Tensor):\n        # input shape: [BATCH_SIZE, N_CHUNKS, N_TOFS,  8, 8]\n        n_tofs = tofs_data.shape[2]\n\n        tof_data_list = []\n        for idx in range(n_tofs):\n            tof_data = tofs_data[:, :, idx].unsqueeze(1) #output_shape: [BATCH_SIZE, 1, N_CHUNKS,  8, 8]\n\n            tof_data = self.seq(tof_data) #output_shape: [BATCH_SIZE, NUM_KERNELS, N_CHUNKS, 1, 3]\n\n            # Flattern the output\n            tof_data = tof_data.permute(0, 2, 1, 3, 4)\n            tof_data = tof_data.reshape(shape=(tof_data.shape[0], tof_data.shape[1], \n                                               tof_data.shape[2]*tof_data.shape[3]*tof_data.shape[4]))\n\n            #output_shape: [BATCH_SIZE, N_CHUNKS, NUM_KERNELS * 3]\n            tof_data_list.append(tof_data)\n        \n        # Average over TOF data\n        stacked_tof_data = torch.stack(tof_data_list, dim=2)\n        average_tof_data = stacked_tof_data.mean(dim=2)\n        #output_shape: [BATCH_SIZE, N_CHUNKS, NUM_KERNELS * 3]\n\n        return average_tof_data\n\nclass SensorSelfAttention(nn.Module):\n    \"\"\"\n    Lớp Self-Attention để gộp các nhánh sensor dựa trên bài báo SADeepSense.\n    Nó tính toán sự tương quan giữa các sensor để xác định trọng số chú ý,\n    sau đó tạo ra một biểu diễn (representation) tổng hợp.\n\n    Tham khảo: Mục III-A và Hình 1 trong bài báo.\n    \"\"\"\n    def __init__(self, input_dim: int, latent_space_dim: int):\n        \"\"\"\n        Khởi tạo lớp SelfAttention.\n\n        Args:\n            input_dim (int): Chiều của feature vector từ mỗi sensor (DIM).\n            latent_space_dim (int): Chiều của không gian tiềm ẩn (latent space) để tính tương quan.\n                         Bài báo đề nghị latent_space_dim bằng số lớp phân loại.\n        \"\"\"\n        super().__init__()\n        self.INPUT_DIM = input_dim\n        self.LATENT_SPACE_DIM = latent_space_dim\n\n        # Phép biến đổi f(X) để trích xuất feature cục bộ cho từng sensor\n        # Tương đương với một lớp fully-connected áp dụng cho mỗi sensor\n        self.f_transform = nn.Linear(self.INPUT_DIM, self.LATENT_SPACE_DIM)\n\n        # Phép biến đổi g(X) để trích xuất feature toàn cục từ tất cả sensor\n        # Tương đương với việc lấy trung bình các sensor rồi qua một lớp fully-connected\n        self.g_transform = nn.Linear(self.INPUT_DIM, self.LATENT_SPACE_DIM)\n\n        # Hằng số để scale, giúp ổn định quá trình huấn luyện\n        self.LATENT_SPACE_DIM_SQRT = math.sqrt(float(self.LATENT_SPACE_DIM))\n\n        # Softmax để chuẩn hóa trọng số chú ý trên các sensor\n        # dim=2 tương ứng với chiều N_SENSORS trong tensor [B, T, N, D]\n        self.softmax = nn.Softmax(dim=2)\n\n    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Thực hiện quá trình forward pass.\n\n        Args:\n            x (torch.Tensor): Tensor đầu vào với shape [B, T, N, D]\n                              B: Batch size\n                              T: Số lượng đoạn thời gian (N_CHUNKS)\n                              N: Số lượng sensor (N_SENSORS)\n                              D: Chiều của feature (DIM)\n\n        Returns:\n            tuple[torch.Tensor, torch.Tensor]:\n            - output (torch.Tensor): Tensor đầu ra sau khi gộp, shape [B, T, D]\n            - attention_weights (torch.Tensor): Trọng số chú ý đã tính, shape [B, T, N]\n        \"\"\"\n        # x có shape: [B, T, N, D]\n        \n        # 1. Biến đổi f(X) và g(X)\n        # f_x: Biểu diễn cục bộ cho mỗi sensor\n        f_x = self.f_transform(x)  # Shape: [B, T, N, h_dim]\n\n        # g_x: Biểu diễn toàn cục\n        # Lấy trung bình các sensor để có context chung\n        global_context = torch.mean(x, dim=2, keepdim=True) # Shape: [B, T, 1, D]\n        g_x = self.g_transform(global_context) # Shape: [B, T, 1, h_dim]\n\n        # 2. Tính toán Heat-Map Tương quan (M) - Công thức (2)\n        # Nhân element-wise, g_x sẽ được broadcast qua chiều N\n        M = (f_x * g_x) / self.LATENT_SPACE_DIM_SQRT # Shape: [B, T, N, h_dim]\n\n        # 3. Tính toán Trọng số Chú ý (w) - Công thức (3)\n        # Lấy trung bình heat-map theo chiều feature h_dim để ra logits\n        logits = torch.mean(M, dim=-1) # Shape: [B, T, N]\n        \n        # Áp dụng softmax để có trọng số\n        attention_weights = self.softmax(logits) # Shape: [B, T, N]\n        \n        # 4. Áp dụng Trọng số và Gộp\n        # Reshape trọng số để nhân với input ban đầu\n        # attention_weights.unsqueeze(-1) -> [B, T, N, 1]\n        weighted_input = x * attention_weights.unsqueeze(-1) # Shape: [B, T, N, D]\n        \n        # Tính tổng có trọng số để gộp chiều sensor (dim=2)\n        output = torch.sum(weighted_input, dim=2) # Shape: [B, T, D]\n\n        return output, attention_weights\n    \nclass MergeBranch(nn.Module):\n    def __init__(self, input_dim:int, latent_space_dim:int = 120, drop_prob=0.2, num_kernels=64, kernel_len:int = 32):\n        super(MergeBranch, self).__init__()\n        self.DROP_PROB = drop_prob\n        self.NUM_KERNELS = num_kernels\n        self.INPUT_DIMS = input_dim\n        self.LATENT_SPACE_DIM = latent_space_dim\n        self.KERNEL_LENGTH = kernel_len\n\n        self.self_attention = SensorSelfAttention(input_dim=self.INPUT_DIMS, latent_space_dim=self.LATENT_SPACE_DIM)\n\n        self.merge_conv1 = nn.Conv2d(in_channels=1, out_channels=self.NUM_KERNELS, kernel_size=(1, self.KERNEL_LENGTH))\n        self.merge_batch_norm1 = nn.BatchNorm2d(num_features=self.NUM_KERNELS)\n\n        self.merge_conv2 = nn.Conv2d(in_channels=self.NUM_KERNELS, out_channels=self.NUM_KERNELS, kernel_size=(1, self.KERNEL_LENGTH))\n        self.merge_batch_norm2 = nn.BatchNorm2d(num_features=self.NUM_KERNELS)\n\n        self.merge_conv3 = nn.Conv2d(in_channels=self.NUM_KERNELS, out_channels=self.NUM_KERNELS, kernel_size=(1, self.KERNEL_LENGTH))\n        self.merge_batch_norm3 = nn.BatchNorm2d(num_features=self.NUM_KERNELS)\n\n        self.activation = nn.ELU()\n        self.dropout = nn.Dropout(p=self.DROP_PROB)\n\n        # Output shape after all convolutions and flattening: \n        # [BATCH_SIZE, N_CHUNKS, OUTPUT_DIMS]    \n        self.output_shape = (None, None, self.NUM_KERNELS * (self.INPUT_DIMS - 3*(self.KERNEL_LENGTH-1)))\n\n    def forward(self, acc_branch, gyr_branch, thrm_branch, tof_branch):\n        # acc_branch, gry_branch, thrm_branch, tof_branch:  [BATCH_SIZE, N_CHUNKS, INPUT_DIMS]\n\n        # stack inputs\n        merge_input = torch.stack([acc_branch, gyr_branch, thrm_branch, tof_branch], dim=2)  # shape [BATCH_SIZE, N_CHUNKS, N_SENSORS, INPUT_DIMS]\n\n        # Apply self-attention\n        merge_output, _ = self.self_attention(merge_input) # shape: [BATCH_SIZE, N_CHUNKS, INPUT_DIMS]\n        merge_output = merge_output.unsqueeze(dim=1)  # shape: [BATCH_SIZE, 1, N_CHUNKS, INPUT_DIMS]\n\n        merge_output = self.merge_conv1(merge_output)\n        merge_output = self.merge_batch_norm1(merge_output)\n        merge_output = self.activation(merge_output)\n        merge_output = self.dropout(merge_output) # shape: [BATCH_SIZE, NUM_KERNELS, N_CHUNKS, INPUT_DIMS - (KERNEL_LENGTH-1)]\n\n        merge_output = self.merge_conv2(merge_output)\n        merge_output = self.merge_batch_norm2(merge_output)\n        merge_output = self.activation(merge_output)\n        merge_output = self.dropout(merge_output) # shape: [BATCH_SIZE, NUM_KERNELS, N_CHUNKS, INPUT_DIMS - 2*(KERNEL_LENGTH-1)]\n\n        merge_output = self.merge_conv3(merge_output)\n        merge_output = self.merge_batch_norm3(merge_output)\n        merge_output = self.activation(merge_output)\n        merge_output = self.dropout(merge_output) # shape: [BATCH_SIZE, NUM_KERNELS, N_CHUNKS, INPUT_DIMS - 3*(KERNEL_LENGTH-1)]\n\n        # flattern output\n        merge_output = merge_output.permute(0, 2, 1, 3)  # shape: [BATCH_SIZE, N_CHUNKS, NUM_KERNELS, INPUT_DIMS - 3*(KERNEL_LENGTH-1)]\n        merge_output = merge_output.reshape(shape=(merge_output.shape[0], merge_output.shape[1], -1))  # shape: [BATCH_SIZE, N_CHUNKS, NUM_KERNELS * (INPUT_DIMS - 3*(KERNEL_LENGTH-1))]    \n\n        return merge_output\n    \nclass TemporalSelfAttention(nn.Module):\n    \"\"\"\n    Lớp Self-Attention theo chiều thời gian để gộp các output từ lớp GRU.\n    Lớp này áp dụng cùng một logic như trong SelfAttention cho sensor, nhưng \n    lần này là để gộp thông tin qua các đoạn thời gian (chunks).\n\n    Tham khảo: Mục III-B và Hình 2 trong bài báo.\n    \"\"\"\n    def __init__(self, input_dim: int, latent_space_dim: int):\n        \"\"\"\n        Khởi tạo lớp SelfTemporalAttention.\n\n        Args:\n            input_dim (int): Chiều của feature vector từ mỗi chunk thời gian (DIMS).\n            h_dim (int): Chiều của không gian tiềm ẩn (latent space) để tính tương quan.\n        \"\"\"\n        super().__init__()\n        self.INPUT_DIM = input_dim\n        self.LATENT_SPACE_DIM = latent_space_dim\n\n        # Phép biến đổi f(X) để trích xuất feature cục bộ cho từng chunk\n        self.f_transform = nn.Linear(input_dim, self.LATENT_SPACE_DIM)\n\n        # Phép biến đổi g(X) để trích xuất feature toàn cục từ tất cả các chunk\n        self.g_transform = nn.Linear(input_dim, self.LATENT_SPACE_DIM)\n        \n        # Hằng số để scale\n        self.LATENT_SPACE_DIM_SQRT = math.sqrt(float(self.LATENT_SPACE_DIM))\n        \n        # Softmax để chuẩn hóa trọng số chú ý trên các chunk thời gian\n        # dim=1 tương ứng với chiều N_CHUNKS trong tensor [B, T, D]\n        self.softmax = nn.Softmax(dim=1)\n\n    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Thực hiện quá trình forward pass.\n\n        Args:\n            x (torch.Tensor): Tensor đầu vào với shape [B, T, D]\n                              B: Batch size\n                              T: Số lượng đoạn thời gian (N_CHUNKS)\n                              D: Chiều của feature (DIMS)\n\n        Returns:\n            tuple[torch.Tensor, torch.Tensor]:\n            - output (torch.Tensor): Tensor đầu ra sau khi gộp, shape [B, D]\n            - attention_weights (torch.Tensor): Trọng số chú ý đã tính, shape [B, T]\n        \"\"\"\n        # x có shape: [B, T, D]\n        \n        # 1. Biến đổi f(X) và g(X)\n        f_x = self.f_transform(x)  # Shape: [B, T, h_dim]\n\n        # Tính bối cảnh toàn cục bằng cách lấy trung bình các chunk\n        global_context = torch.mean(x, dim=1, keepdim=True) # Shape: [B, 1, D]\n        g_x = self.g_transform(global_context) # Shape: [B, 1, h_dim]\n\n        # 2. Tính toán Heat-Map Tương quan (M)\n        # g_x sẽ được broadcast qua chiều T\n        M = (f_x * g_x) / self.LATENT_SPACE_DIM_SQRT # Shape: [B, T, h_dim]\n\n        # 3. Tính toán Trọng số Chú ý (w)\n        # Lấy trung bình heat-map theo chiều feature h_dim\n        logits = torch.mean(M, dim=-1) # Shape: [B, T]\n        \n        # Áp dụng softmax để có trọng số\n        attention_weights = self.softmax(logits) # Shape: [B, T]\n        \n        # 4. Áp dụng Trọng số và Gộp\n        # Reshape trọng số để nhân với input ban đầu\n        # attention_weights.unsqueeze(-1) -> [B, T, 1]\n        weighted_input = x * attention_weights.unsqueeze(-1) # Shape: [B, T, D]\n        \n        # Tính tổng có trọng số để gộp chiều thời gian (dim=1)\n        output = torch.sum(weighted_input, dim=1) # Shape: [B, D]\n\n        return output, attention_weights\n    \nclass SADeepSense(nn.Module):\n    def __init__(self, config):\n        super(SADeepSense, self).__init__()\n        self.ACC_CHANNELS = config.get(\"acc_channels\", 3)\n        self.GYR_CHANNELS = config.get(\"gyr_channels\", 4)\n        self.THRM_CHANNELS = config.get(\"thrm_channels\", 5)\n\n        self.SENSOR_BRANCH_NUM_KERNELS = config.get(\"sensor_branch_num_kernels\", 64)\n        self.SENSOR_BRANCH_DROP_PROB = config.get(\"sensor_branch_drop_prob\", 0.2)\n        self.CHUNKS = config.get(\"chunks\", 120)\n\n        self.MERGE_BRANCH_SA_LATENT_SPACE_DIM = config.get(\"merge_branch_sa_latent_space_dim\", 120)\n        self.MERGE_BRANCH_NUM_KERNELS = config.get(\"merge_branch_num_kernels\", 64)\n        self.MERGE_BRANCH_KERNEL_LENGTH = config.get(\"merge_branch_kernel_length\", 32)\n        self.MERGE_BRANCH_DROP_PROB = config.get(\"merge_branch_drop_prob\", 0.2)\n\n        self.GRU_HIDDEN_SIZE = config.get(\"gru_hidden_size\", 200)\n        self.GRU_NUM_LAYERS = config.get(\"gru_num_layers\", 2)\n\n        self.TEMPORAL_SA_LATENT_SPACE_DIM = config.get(\"temporal_sa_latent_space_dim\", 120)\n\n        self.N_LABELS = config.get(\"n_labels\", 18)\n\n        self.acc_branch = SensorBranch(drop_prob=self.SENSOR_BRANCH_DROP_PROB, sensor_channels=self.ACC_CHANNELS, num_kernels=self.SENSOR_BRANCH_NUM_KERNELS)\n        self.gyr_branch = SensorBranch(drop_prob=self.SENSOR_BRANCH_DROP_PROB, sensor_channels=self.GYR_CHANNELS, num_kernels=self.SENSOR_BRANCH_NUM_KERNELS)\n        self.thrm_branch = SensorBranch(drop_prob=self.SENSOR_BRANCH_DROP_PROB, sensor_channels=self.THRM_CHANNELS, num_kernels=self.SENSOR_BRANCH_NUM_KERNELS)\n        self.tof_branch = TOFBranch(drop_prob=self.SENSOR_BRANCH_DROP_PROB, num_kernels=self.SENSOR_BRANCH_NUM_KERNELS)\n\n        self.merge_branch = MergeBranch(input_dim=self.acc_branch.output_shape[2], \n                                        latent_space_dim=self.MERGE_BRANCH_SA_LATENT_SPACE_DIM, \n                                        drop_prob=self.MERGE_BRANCH_DROP_PROB, \n                                        num_kernels=self.MERGE_BRANCH_NUM_KERNELS,\n                                        kernel_len=self.MERGE_BRANCH_KERNEL_LENGTH)\n\n        self.gru_layer = nn.GRU(input_size=self.merge_branch.output_shape[2], hidden_size=self.GRU_HIDDEN_SIZE, num_layers=self.GRU_NUM_LAYERS, batch_first=True)\n        self.temporal_self_attention = TemporalSelfAttention(input_dim=self.GRU_HIDDEN_SIZE, latent_space_dim=self.TEMPORAL_SA_LATENT_SPACE_DIM)\n        self.full_connected_layer = nn.Linear(in_features=self.GRU_HIDDEN_SIZE, out_features=self.N_LABELS)\n\n    def forward(self, n_chunks, acc_data, gyr_data, thrm_data, tof_data):\n\n        # Forward each sensor data to each sensor branch\n        acc_branch_output = self.acc_branch(acc_data)\n        gyr_branch_output = self.gyr_branch(gyr_data)\n        thrm_branch_output = self.thrm_branch(thrm_data)\n        tof_branch_output = self.tof_branch(tof_data)\n\n        # Merge data from all sensor branches\n        merge_branch_output = self.merge_branch(acc_branch_output, gyr_branch_output, thrm_branch_output, tof_branch_output)\n\n        # Pack the merged branch output for the GRU layer\n        packed_merge_branch_output = torch.nn.utils.rnn.pack_padded_sequence(merge_branch_output, lengths=n_chunks, batch_first=True, enforce_sorted=False)\n        gru_packed_output, _ = self.gru_layer(packed_merge_branch_output)\n        gru_output, _ = torch.nn.utils.rnn.pad_packed_sequence(gru_packed_output, batch_first=True, total_length=self.CHUNKS)\n\n        sta_output, _ = self.temporal_self_attention(gru_output)\n        fcl_output = self.full_connected_layer(sta_output)\n\n        return fcl_output\n        \n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-08-23T07:31:58.743429Z","iopub.execute_input":"2025-08-23T07:31:58.743703Z","iopub.status.idle":"2025-08-23T07:31:58.775749Z","shell.execute_reply.started":"2025-08-23T07:31:58.743682Z","shell.execute_reply":"2025-08-23T07:31:58.775192Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"727dc880","cell_type":"code","source":"model_config = {\n    \"acc_channels\": 3,\n    \"gyr_channels\": 4,\n    \"thrm_channels\": 5,\n    \"sensor_branch_num_kernels\": 64,\n    \"sensor_branch_drop_prob\": 0.15,\n    \"chunks\": train_dataset.MAX_NCHUNKS,\n    \"merge_branch_sa_latent_space_dim\": 64,\n    \"merge_branch_num_kernels\": 64,\n    \"merge_branch_kernel_length\": 56,\n    \"merge_branch_drop_prob\": 0.15,\n    \"gru_hidden_size\": 512,\n    \"gru_num_layers\": 3,\n    \"temporal_sa_latent_space_dim\": 64,\n    \"n_labels\": len(train_dataset.TARGET_ENCODER.classes_)\n}\nmlflow.log_params(model_config)\nmodel = SADeepSense(model_config)\n\n\nprint(model)","metadata":{"execution":{"iopub.status.busy":"2025-08-23T08:10:35.655243Z","iopub.execute_input":"2025-08-23T08:10:35.655783Z","iopub.status.idle":"2025-08-23T08:10:36.008679Z","shell.execute_reply.started":"2025-08-23T08:10:35.655759Z","shell.execute_reply":"2025-08-23T08:10:36.008018Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"9edb4c5b","cell_type":"code","source":"# Print total number of parameters in the model\ntotal_params = sum(p.numel() for p in model.parameters())\nprint(f\"Total number of parameters: {total_params}\")\nmlflow.log_param(\"total_params\", total_params)","metadata":{"execution":{"iopub.status.busy":"2025-08-23T08:10:40.456315Z","iopub.execute_input":"2025-08-23T08:10:40.456889Z","iopub.status.idle":"2025-08-23T08:10:40.676142Z","shell.execute_reply.started":"2025-08-23T08:10:40.456864Z","shell.execute_reply":"2025-08-23T08:10:40.675543Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"2326d576","cell_type":"markdown","source":"# Setting training parameters","metadata":{}},{"id":"56500968","cell_type":"code","source":"class EarlyStopping:\n  def __init__(self, patience:int=3, min_delta:float=0.0):\n    self.patience = patience\n    self.min_delta = min_delta\n    self.counter = 0\n    self.min_loss = float('inf')\n\n  def early_stop(self, validation_loss):\n    if validation_loss - self.min_loss >= self.min_delta:\n      self.counter += 1\n      if self.counter >= self.patience:\n        return True\n    else:\n      self.counter = 0\n      self.min_loss = validation_loss\n    return False","metadata":{"execution":{"iopub.status.busy":"2025-08-23T07:32:23.218747Z","iopub.execute_input":"2025-08-23T07:32:23.219455Z","iopub.status.idle":"2025-08-23T07:32:23.223629Z","shell.execute_reply.started":"2025-08-23T07:32:23.219433Z","shell.execute_reply":"2025-08-23T07:32:23.223098Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"80768991","cell_type":"code","source":"class Trainer():\n    def __init__(self, model:SADeepSense, n_labels:int,\n                 train_data:CMIDataset, test_data:CMIDataset,\n                 lr:float=1e-4, max_lr:float=1e-3, num_epochs:int=100, \n                 batch_size = 16, \n                 turn_off_sensor_prob:float=0.0,\n                 early_stopping_patience:int=5,\n                 save_checkpoint_path:str=\"checkpoint.pth\"):\n        self.LEARNING_RATE = lr\n        self.MAX_LR = max_lr\n        self.NUM_EPOCHS = num_epochs\n        self.TRAIN_DATA = train_data\n        self.TEST_DATA = test_data\n        self.model = model\n        self.N_LABELS = n_labels\n        self.SAVE_CHECKPOINT_PATH = save_checkpoint_path\n        self.BATCH_SIZE = batch_size\n        self.TURN_OFF_SENSOR_PROB = turn_off_sensor_prob\n        self.EARLY_STOPPING_PATIENCE = early_stopping_patience\n\n    def setup(self):\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.train_loader = DataLoader(self.TRAIN_DATA, batch_size=self.BATCH_SIZE, shuffle=True)\n        self.test_loader = DataLoader(self.TEST_DATA, batch_size=self.BATCH_SIZE, shuffle=False)\n\n        self.criterion = nn.CrossEntropyLoss()\n        self.metrics = F1Score(task=\"multiclass\", num_classes=self.N_LABELS, average='weighted')\n        \n        self.optimizer = AdamW(self.model.parameters(), lr=self.LEARNING_RATE)\n        self.scheduler = OneCycleLR(self.optimizer, max_lr=self.MAX_LR, steps_per_epoch=len(self.train_loader), epochs=self.NUM_EPOCHS)\n        self.early_stopping = EarlyStopping(patience=self.EARLY_STOPPING_PATIENCE, min_delta=0.01)\n    \n    def train(self):\n        self.model.to(self.device)\n        self.metrics.to(self.device)\n\n        \n        for epoch in range(self.NUM_EPOCHS):\n            train_loss = 0.0\n            train_f1_score = 0.0\n            self.model.train()\n            for batch in tqdm(self.train_loader):\n               \n                n_chunks, labels, acc_data, gyr_data, thrm_data, tof_data = batch\n\n                if self.TURN_OFF_SENSOR_PROB > 0.0:\n                    if np.random.random() < self.TURN_OFF_SENSOR_PROB:\n                        # Randomly turn off one of the sensors\n                        sensor_to_turn_off = np.random.choice(['thrm', 'tof'])\n                        if sensor_to_turn_off == 'thrm':\n                            thrm_data = torch.zeros_like(thrm_data)\n                        elif sensor_to_turn_off == 'tof':\n                            tof_data = torch.zeros_like(tof_data)\n\n                #n_chunks = n_chunks.to(self.device)\n                labels = labels.to(self.device)\n                acc_data = acc_data.to(self.device)\n                gyr_data = gyr_data.to(self.device)\n                thrm_data = thrm_data.to(self.device)\n                tof_data = tof_data.to(self.device)\n\n                self.optimizer.zero_grad()\n\n                outputs = self.model(n_chunks, acc_data, gyr_data, thrm_data, tof_data)\n                loss = self.criterion(outputs, labels)\n                self.metrics.update(outputs, labels)\n                train_loss += loss.item()\n\n                loss.backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)\n                #torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=0.5)\n                self.optimizer.step()\n\n            # Calculate f1 score and loss for each epoch\n            train_f1_score = self.metrics.compute()\n            self.metrics.reset()\n            train_loss = train_loss/len(self.train_loader)\n            mlflow.log_metric(\"f1_score\", train_f1_score.item(), step=epoch)\n            mlflow.log_metric(\"loss\", train_loss, step=epoch)\n\n            # Update learning rate\n            self.scheduler.step()\n\n            # Save checkpoint\n            checkpoint = {\n                \"model_state_dict\": self.model.state_dict(),\n                \"epoch\": epoch,\n                \"label_encoder_classes\": train_dataset.TARGET_ENCODER.classes_.tolist(),\n            }\n            torch.save(checkpoint, self.SAVE_CHECKPOINT_PATH)\n            \n\n            self.model.eval()\n            eval_loss = 0.0\n            eval_f1_score = 0.0\n            for batch in tqdm(self.test_loader):\n                n_chunks, labels, acc_data, gyr_data, thrm_data, tof_data = batch\n\n                #n_chunks = n_chunks.to(self.device)\n                labels = labels.to(self.device)\n                acc_data = acc_data.to(self.device)\n                gyr_data = gyr_data.to(self.device)\n                thrm_data = thrm_data.to(self.device)\n                tof_data = tof_data.to(self.device)\n\n                with torch.no_grad():\n                    outputs = self.model(n_chunks, acc_data, gyr_data, thrm_data, tof_data)\n                    self.metrics.update(outputs, labels)\n\n                eval_loss += self.criterion(outputs, labels).item()\n\n            eval_loss = eval_loss / len(self.test_loader)\n            eval_f1_score = self.metrics.compute()\n            self.metrics.reset()\n            mlflow.log_metric(\"test_f1_score\", eval_f1_score.item(), step=epoch)\n            mlflow.log_metric(\"test_loss\", eval_loss, step=epoch)\n\n            print(f\"Epoch {epoch + 1}/{self.NUM_EPOCHS}, Learning Rate: {self.scheduler.get_last_lr()[0]:.6f}, Loss: {train_loss:.4f}, Batch F1 Score: {train_f1_score:.4f}, Test Loss: {eval_loss:.4f}, Test F1 Score: {eval_f1_score:.4f}\")\n\n            if self.early_stopping.early_stop(eval_loss):\n                print(\"Early stopping triggered.\")\n                break\n","metadata":{"execution":{"iopub.status.busy":"2025-08-23T08:11:18.474953Z","iopub.execute_input":"2025-08-23T08:11:18.475546Z","iopub.status.idle":"2025-08-23T08:11:18.497494Z","shell.execute_reply.started":"2025-08-23T08:11:18.475524Z","shell.execute_reply":"2025-08-23T08:11:18.496328Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"fb2c62bb","cell_type":"markdown","source":"# Training","metadata":{}},{"id":"a3cb51e8","cell_type":"code","source":"training_config = {\n    \"lr\": 1e-3,\n    \"max_lr\": 5e-3,\n    \"num_epochs\": 80,\n    \"batch_size\": 16,\n    \"turn_off_sensor_prob\": 0.2,\n    \"early_stopping_patience\": 15\n}\n\ntrainer = Trainer(model=model, n_labels=len(train_dataset.TARGET_ENCODER.classes_),\n                  train_data=train_dataset, test_data=test_dataset,\n                  **training_config)","metadata":{"execution":{"iopub.status.busy":"2025-08-23T08:11:22.604488Z","iopub.execute_input":"2025-08-23T08:11:22.604766Z","iopub.status.idle":"2025-08-23T08:11:22.609089Z","shell.execute_reply.started":"2025-08-23T08:11:22.604744Z","shell.execute_reply":"2025-08-23T08:11:22.608378Z"},"scrolled":true,"trusted":true},"outputs":[],"execution_count":null},{"id":"642a5ffc","cell_type":"code","source":"trainer.setup()","metadata":{"execution":{"iopub.status.busy":"2025-08-23T08:11:28.77641Z","iopub.execute_input":"2025-08-23T08:11:28.777218Z","iopub.status.idle":"2025-08-23T08:11:28.782789Z","shell.execute_reply.started":"2025-08-23T08:11:28.777189Z","shell.execute_reply":"2025-08-23T08:11:28.782076Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"4a6de703","cell_type":"code","source":"mlflow.log_params(training_config)\ntrainer.train()","metadata":{"execution":{"iopub.status.busy":"2025-08-23T08:11:31.216352Z","iopub.execute_input":"2025-08-23T08:11:31.216653Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"c3e34b5b","cell_type":"code","source":"mlflow.log_artifact(trainer.SAVE_CHECKPOINT_PATH, artifact_path=\"model_checkpoint\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T08:31:44.784219Z","iopub.execute_input":"2025-08-25T08:31:44.784509Z","iopub.status.idle":"2025-08-25T08:31:49.272022Z","shell.execute_reply.started":"2025-08-25T08:31:44.784485Z","shell.execute_reply":"2025-08-25T08:31:49.271237Z"}},"outputs":[],"execution_count":null},{"id":"9fb2e8dc","cell_type":"markdown","source":"# Reload model","metadata":{}},{"id":"d9eef6af","cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ncheckpoint = torch.load(f=\"/kaggle/input/sadeepsense/pytorch/baseline/2/checkpoint.pth\", map_location=device)\n\ntarget_encoder = LabelEncoder()\ntarget_encoder = target_encoder.fit(checkpoint[\"label_encoder_classes\"])\n\nmodel_config = {\n    \"acc_channels\": 3,\n    \"gyr_channels\": 4,\n    \"thrm_channels\": 5,\n    \"sensor_branch_num_kernels\": 64,\n    \"sensor_branch_drop_prob\": 0.2,\n    \"chunks\": 120,\n    \"sensor_sa_latent_space_dim\": 120,\n    \"merge_branch_sa_latent_space_dim\": 120,\n    \"merge_branch_num_kernels\": 64,\n    \"merge_branch_drop_prob\": 0.2,\n    \"gru_hidden_size\": 256,\n    \"gru_num_layers\": 2,\n    \"n_labels\": 18\n}\nmodel = SADeepSense(model_config).to(device)\nmodel.load_state_dict(checkpoint[\"model_state_dict\"])","metadata":{},"outputs":[],"execution_count":null},{"id":"47f3f9a9","cell_type":"markdown","source":"# Inference","metadata":{}},{"id":"17546061","cell_type":"code","source":"train_config = {\n    'SEQUENCE_LEN': 700,\n    'N_SAMPLES_IN_CHUNK': 5,\n    'NAN_VALUES': 0.0, \n    'TARGET_COL': 'gesture',\n    'TARGET_ENCODER': target_encoder,\n    'SAMPLING_RATE': 10,\n    'N_TOFs': 5, \n    'MAX_NCHUNKS': 140, \n    'NROWS': None \n}","metadata":{},"outputs":[],"execution_count":null},{"id":"8950a8ce","cell_type":"code","source":"def convert_freq_domain(chunk_data: pd.DataFrame) -> np.ndarray:\n    '''\n    Convert time domain sensor data to frequency domain using FFT.\n\n    Args:\n        chunk_data (pd.DataFrame): Time domain sensor data.\n\n    Returns:\n        np.ndarray: Frequency domain representation of the sensor data.\n    '''\n    fourier = rfft(chunk_data, axis=0)\n    amplitude = 2*np.abs(np.array(fourier, dtype=complex))/chunk_data.shape[0]\n    phase = np.angle(np.array(fourier, dtype=complex), deg=False)\n\n    return np.concatenate([amplitude, phase], axis=0).T","metadata":{},"outputs":[],"execution_count":null},{"id":"078e525e","cell_type":"code","source":"def process_sequence(sequence_data: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:\n    acc_cols = ['acc_x', 'acc_y', 'acc_z']\n    gyr_cols = ['rot_x', 'rot_y', 'rot_z', 'rot_w']\n    thrm_cols = [c for c in sequence_data.columns if c.startswith('thm')]\n    tof_cols = [c for c in sequence_data.columns if c.startswith('tof')]\n\n    # Sort by sequence_counter and reset_index\n    sequence_data = sequence_data.sort_values(by=\"sequence_counter\", ascending=True).reset_index()\n\n    # Ensure the sequence data has maximum length of SEQUENCE_LEN\n    if sequence_data.shape[0] > train_config[\"SEQUENCE_LEN\"]:\n        sequence_data = sequence_data.iloc[:train_config[\"SEQUENCE_LEN\"], :]\n\n    # If the number of samples in the last chunk is not equal to N_SAMPLES_IN_CHUNK, pad the data\n    if sequence_data.shape[0] % train_config[\"N_SAMPLES_IN_CHUNK\"] != 0:\n        pad_data = np.zeros(shape=(train_config[\"N_SAMPLES_IN_CHUNK\"] - (sequence_data.shape[0] % train_config[\"N_SAMPLES_IN_CHUNK\"]), sequence_data.shape[1]))\n        sequence_data = pd.concat([sequence_data, pd.DataFrame(data=pad_data, columns=sequence_data.columns)], ignore_index=True)\n    \n    acc_seq_freq = []\n    gyr_seq_freq = []\n    thrm_seq_freq = []\n    tof_seq = []\n\n    # Divide the sequence data into chunks of N_SAMPLES_IN_CHUNK\n    # and convert each chunk to frequency domain features\n    for start_chunk_index in range(0, sequence_data.shape[0], train_config[\"N_SAMPLES_IN_CHUNK\"]):\n        chunk_data = sequence_data.iloc[start_chunk_index:start_chunk_index+train_config[\"N_SAMPLES_IN_CHUNK\"]]\n\n        acc_chunk_freq = convert_freq_domain(chunk_data.loc[:, acc_cols])\n        gyr_chunk_freq = convert_freq_domain(chunk_data.loc[:, gyr_cols])\n        thrm_chunk_freq = convert_freq_domain(chunk_data.loc[:, thrm_cols])\n\n        # Reshape TOF data to have shape \n        # And average over the time dimension for each TOF sensor\n        tof_chunk = chunk_data.loc[:, tof_cols].values.reshape(train_config[\"N_SAMPLES_IN_CHUNK\"], train_config[\"N_TOFs\"], 8, 8)\n        tof_chunk = tof_chunk.transpose(1, 0, 2, 3).mean(axis=1)\n\n        # Append the frequency domain features to the respective lists\n        acc_seq_freq.append(acc_chunk_freq)\n        gyr_seq_freq.append(gyr_chunk_freq)\n        thrm_seq_freq.append(thrm_chunk_freq)\n        tof_seq.append(tof_chunk)\n\n    return np.array(acc_seq_freq), np.array(gyr_seq_freq), np.array(thrm_seq_freq), np.array(tof_seq), sequence_data.shape[0] // train_config[\"N_SAMPLES_IN_CHUNK\"]","metadata":{},"outputs":[],"execution_count":null},{"id":"b53d2ec1","cell_type":"code","source":"def predict(sequence: pl.DataFrame, demographics: pl.DataFrame) -> str:\n    sequence = sequence.to_pandas()\n    demographics = demographics.to_pandas()\n\n    # Preprocess the input data\n    acc_data, gyr_data, thrm_data, tof_data, num_chunks = process_sequence(sequence)\n\n    if num_chunks > train_config['MAX_NCHUNKS']:\n            acc_data = acc_data[:train_config['MAX_NCHUNKS']]\n            gyr_data = gyr_data[:train_config['MAX_NCHUNKS']]\n            thrm_data = thrm_data[:train_config['MAX_NCHUNKS']]\n            tof_data = tof_data[:train_config['MAX_NCHUNKS']]\n            num_chunks = train_config['MAX_NCHUNKS']\n    else:\n        n_padded_chunks = train_config['MAX_NCHUNKS'] - num_chunks\n\n        pad_acc_data = np.full(shape=(n_padded_chunks, acc_data.shape[1], acc_data.shape[2]), fill_value=train_config['NAN_VALUES'])\n        pad_gyr_data = np.full(shape=(n_padded_chunks, gyr_data.shape[1], gyr_data.shape[2]), fill_value=train_config['NAN_VALUES'])\n        pad_thrm_data = np.full(shape=(n_padded_chunks, thrm_data.shape[1], thrm_data.shape[2]), fill_value=train_config['NAN_VALUES'])\n        pad_tof_data = np.full(shape=(n_padded_chunks, tof_data.shape[1], tof_data.shape[2], tof_data.shape[3]), fill_value=train_config['NAN_VALUES'])\n\n        acc_data = np.concatenate([acc_data, pad_acc_data])\n        gyr_data = np.concatenate([gyr_data, pad_gyr_data])\n        thrm_data = np.concatenate([thrm_data, pad_thrm_data])\n        tof_data = np.concatenate([tof_data, pad_tof_data])\n\n    num_chunks = torch.tensor(num_chunks).unsqueeze(0)\n    acc_data = torch.tensor(acc_data, dtype=torch.float32).unsqueeze(0)  # Add batch dimension\n    gyr_data = torch.tensor(gyr_data, dtype=torch.float32).unsqueeze(0)  # Add batch dimension\n    thrm_data = torch.tensor(thrm_data, dtype=torch.float32).unsqueeze(0)  # Add batch dimension\n    tof_data = torch.tensor(tof_data, dtype=torch.float32).unsqueeze(0)  # Add batch dimension\n\n\n    model.eval()  # Set the model to evaluation mode\n    # Forward pass through the model\n    with torch.no_grad():\n        acc_data = acc_data.to(device)\n        gyr_data = gyr_data.to(device)\n        thrm_data = thrm_data.to(device)\n        tof_data = tof_data.to(device)\n        predictions = model(num_chunks, acc_data, gyr_data, thrm_data, tof_data)\n\n    predictions = predictions.argmax(dim=1).detach().cpu().numpy()\n    \n    # Get the predicted class indices\n    predicted_classes = target_encoder.inverse_transform(predictions)[0]\n    \n    # Postprocess the predictions\n    return predicted_classes","metadata":{},"outputs":[],"execution_count":null},{"id":"134ac7cc","cell_type":"code","source":"inference_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    )\n","metadata":{},"outputs":[],"execution_count":null},{"id":"4f9e93fd","cell_type":"code","source":"mlflow.end_run()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-25T08:32:08.341735Z","iopub.execute_input":"2025-08-25T08:32:08.342006Z","iopub.status.idle":"2025-08-25T08:32:09.912454Z","shell.execute_reply.started":"2025-08-25T08:32:08.341986Z","shell.execute_reply":"2025-08-25T08:32:09.911878Z"}},"outputs":[],"execution_count":null}]}