{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91531,"databundleVersionId":11044547,"sourceType":"competition"},{"sourceId":10967741,"sourceType":"datasetVersion","datasetId":6824006}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install mlflow tonic","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T19:58:01.436834Z","iopub.execute_input":"2025-03-10T19:58:01.437108Z","iopub.status.idle":"2025-03-10T19:58:11.175835Z","shell.execute_reply.started":"2025-03-10T19:58:01.437087Z","shell.execute_reply":"2025-03-10T19:58:11.175050Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"configs = {\n \"device\": \"cuda:0\",\n \"mlflow_path\": \"/kaggle/working/mlruns\",\n \"experiment_name\": \"trial_experiment\",\n \"data_dir\": \"/kaggle/input/event-based-eye-tracking-cvpr-2025/event_data/event_data\",\n \"run_name\": \"trial_run\",\n \"architecture\": \"BBModel\",\n # \"lr\": 0.001,\n \"num_epochs\": 200,\n \"batch_size\": 32,\n \"spatial_factor\": 0.125,\n \"temporal_subsample_factor\": 1,\n \"val_interval\": 2,\n \"save_k_best\": 2,\n \"pixel_tolerances\": [5,10,15],\n \"sensor_width\": 640,\n \"sensor_height\": 480,\n \"train_stride\": 15,\n \"val_stride\": 30,\n \"train_length\": 30,\n \"val_length\": 30,\n \"n_time_bins\": 3,\n \"voxel_grid_ch_normaization\": False,\n \"loss\": \"euclidean\",\n \"n_time_bins\": 2, # Number of channels in the event representation.\n \"kernel_size_t\": 3, # Temporal convolution kernel size.\n \"kernel_size_spatial\": 3, # Spatial convolution kernel size.\n \"num_frames\": 50, # Number of event frames per batch.\n \"learning_rate\": 0.002,\n \"weight_decay\": 0.005,\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T18:34:54.326833Z","iopub.execute_input":"2025-03-10T18:34:54.327182Z","iopub.status.idle":"2025-03-10T18:34:54.332656Z","shell.execute_reply.started":"2025-03-10T18:34:54.327154Z","shell.execute_reply":"2025-03-10T18:34:54.331516Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import argparse, json, os, mlflow\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\n\n# Load config file and create args\n# config_file = 'train_baseline.json'\n# with open(os.path.join('./configs', config_file), 'r') as f:\n# config = json.load(f)\nargs = argparse.Namespace(**configs)\n\n# Set up MLflow tracking\nmlflow.set_tracking_uri(args.mlflow_path)\nmlflow.set_experiment(experiment_name=args.experiment_name)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-10T18:34:54.795736Z","iopub.execute_input":"2025-03-10T18:34:54.796075Z","iopub.status.idle":"2025-03-10T18:34:54.806602Z","shell.execute_reply.started":"2025-03-10T18:34:54.796047Z","shell.execute_reply":"2025-03-10T18:34:54.805512Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#CNN_GRU\nclass CNN_GRU(nn.Module):\n \"\"\"\n A baseline eye tracking which uses CNN + GRU to predict the pupil center coordinate\n \"\"\"\n def __init__(self, args):\n super().__init__() \n self.args = args\n self.conv1 = nn.Conv2d(args.n_time_bins, 32, kernel_size=3, stride=1, padding=1)\n self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)\n self.conv3 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1)\n self.pool = nn.MaxPool2d(kernel_size=3, stride=2)\n self.gru = nn.GRU(input_size=36192, hidden_size=128, num_layers=1, batch_first=True)\n self.fc = nn.Linear(128, 2)\n\n\n def forward(self, x):\n # input is of shape (batch_size, seq_len, channels, height, width)\n batch_size, seq_len, channels, height, width = x.shape\n x = x.view(batch_size*seq_len, channels, height, width)\n # permute height and width\n x = x.permute(0, 1, 3, 2)\n\n x= self.conv1(x)\n x= torch.relu(x)\n x= self.conv2(x)\n x= torch.relu(x)\n x= self.conv3(x)\n x= torch.relu(x)\n x= self.pool(x)\n\n x = x.view(batch_size, seq_len, -1)\n x, _ = self.gru(x)\n # output shape of x is (batch_size, seq_len, hidden_size)\n\n x = self.fc(x)\n # output is of shape (batch_size, seq_len, 2)\n return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-09T21:13:20.091486Z","iopub.execute_input":"2025-03-09T21:13:20.091780Z","iopub.status.idle":"2025-03-09T21:13:20.098524Z","shell.execute_reply.started":"2025-03-09T21:13:20.091757Z","shell.execute_reply":"2025-03-09T21:13:20.097494Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#ConvLSTM\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass ConvLSTMCell(nn.Module):\n \"\"\"\n Convolutional LSTM cell.\n It computes all gate operations via one convolution.\n \"\"\"\n def __init__(self, input_channels, hidden_channels, kernel_size, bias=True):\n super().__init__()\n self.hidden_channels = hidden_channels\n padding = kernel_size // 2 # use same padding to preserve spatial dimensions\n self.conv = nn.Conv2d(\n in_channels=input_channels + hidden_channels,\n out_channels=4 * hidden_channels,\n kernel_size=kernel_size,\n padding=padding,\n bias=bias\n )\n \n def forward(self, input, h_cur, c_cur):\n # Concatenate input and previous hidden state along the channel dimension\n combined = torch.cat([input, h_cur], dim=1)\n conv_output = self.conv(combined)\n # Split the output into 4 parts for input, forget, output, and cell gates\n cc_i, cc_f, cc_o, cc_g = torch.split(conv_output, self.hidden_channels, dim=1)\n i = torch.sigmoid(cc_i)\n f = torch.sigmoid(cc_f)\n o = torch.sigmoid(cc_o)\n g = torch.tanh(cc_g)\n c_next = f * c_cur + i * g\n h_next = o * torch.tanh(c_next)\n return h_next, c_next\n\nclass ConvLSTMStack(nn.Module):\n \"\"\"\n Stacks multiple ConvLSTM layers.\n The output of one layer is used as the input to the next.\n \"\"\"\n def __init__(self, input_channels, hidden_channels, kernel_size, num_layers):\n super().__init__()\n self.num_layers = num_layers\n self.cells = nn.ModuleList()\n for i in range(num_layers):\n in_ch = input_channels if i == 0 else hidden_channels\n self.cells.append(ConvLSTMCell(in_ch, hidden_channels, kernel_size))\n \n def forward(self, input):\n # Input shape: (batch, seq_len, channels, height, width)\n batch_size, seq_len, _, height, width = input.size()\n layer_output = input\n for cell in self.cells:\n h, c = self.init_hidden(batch_size, cell.hidden_channels, height, width, input.device)\n outputs = []\n # Process sequence step by step\n for t in range(seq_len):\n h, c = cell(layer_output[:, t], h, c)\n outputs.append(h)\n # Stack outputs along the time dimension for the next layer\n layer_output = torch.stack(outputs, dim=1)\n return layer_output # shape: (batch, seq_len, hidden_channels, height, width)\n \n def init_hidden(self, batch_size, hidden_channels, height, width, device):\n h = torch.zeros(batch_size, hidden_channels, height, width, device=device)\n c = torch.zeros(batch_size, hidden_channels, height, width, device=device)\n return h, c\n\nclass CNN_ConvLSTM(nn.Module):\n \"\"\"\n A convolutional LSTM based eye tracking model.\n \n The model first applies 2D convolutions (and pooling) to each event-frame\n to extract spatial features. Then, a 4-layer ConvLSTM stack processes the\n sequence of feature maps while preserving spatial structure. Finally, global\n average pooling and a fully connected layer regress the pupil (x, y) coordinates.\n \"\"\"\n def __init__(self, args):\n super().__init__()\n self.args = args\n # Initial convolution layers to extract spatial features\n # (Input channels is set to args.n_time_bins, e.g. the number of channels in the voxel representation)\n self.conv1 = nn.Conv2d(args.n_time_bins, 32, kernel_size=3, stride=1, padding=1)\n self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)\n self.conv3 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1)\n self.pool = nn.MaxPool2d(kernel_size=3, stride=2)\n \n # Stack 4 layers of ConvLSTM; here we use 32 hidden channels in each layer.\n self.num_conv_lstm_layers = 4\n self.hidden_channels = 32\n self.convlstm = ConvLSTMStack(\n input_channels=32, \n hidden_channels=self.hidden_channels, \n kernel_size=3, \n num_layers=self.num_conv_lstm_layers\n )\n \n # After the ConvLSTM, we perform global average pooling over the spatial dimensions\n # and then use a fully connected layer to regress the (x, y) coordinates.\n self.fc = nn.Linear(self.hidden_channels, 2)\n \n def forward(self, x):\n # x: (batch, seq_len, channels, height, width)\n batch_size, seq_len, channels, height, width = x.shape\n # Process each frame independently via convolution layers.\n # Merge batch and sequence dimensions.\n x = x.view(batch_size * seq_len, channels, height, width)\n x = F.relu(self.conv1(x))\n x = F.relu(self.conv2(x))\n x = F.relu(self.conv3(x))\n x = self.pool(x)\n # Get the new spatial dimensions after conv+pool.\n _, feat_channels, feat_h, feat_w = x.shape\n # Reshape back to sequence: (batch, seq_len, feat_channels, feat_h, feat_w)\n x = x.view(batch_size, seq_len, feat_channels, feat_h, feat_w)\n \n # Pass the sequence through the stacked ConvLSTM layers.\n x = self.convlstm(x) # Output shape: (batch, seq_len, hidden_channels, feat_h, feat_w)\n \n # For each time step, perform spatial global average pooling.\n x = x.mean(dim=[3, 4]) # Now shape: (batch, seq_len, hidden_channels)\n \n # Apply the final fully connected layer to regress the (x, y) coordinates.\n x = self.fc(x) # Shape: (batch, seq_len, 2)\n return x\n\n# # Example usage:\n# if __name__ == \"__main__\":\n# class Args:\n# # Dummy hyperparameters for demonstration purposes\n# n_time_bins = 3 # for example, 3 channels in the voxel representation\n# device = 'cuda' if torch.cuda.is_available() else 'cpu'\n# args = Args()\n# model = CNN_ConvLSTM(args).to(args.device)\n# # Create a dummy input: batch size = 2, sequence length = 10, 3 channels, height=80, width=60\n# dummy_input = torch.randn(2, 10, args.n_time_bins, 80, 60).to(args.device)\n# output = model(dummy_input)\n# print(\"Output shape:\", output.shape) # Expected: (2, 10, 2)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-09T21:13:20.501835Z","iopub.execute_input":"2025-03-09T21:13:20.502119Z","iopub.status.idle":"2025-03-09T21:13:20.514292Z","shell.execute_reply.started":"2025-03-09T21:13:20.502097Z","shell.execute_reply":"2025-03-09T21:13:20.513607Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#CTEM\nclass CETM(nn.Module):\n \"\"\"\n Consistent Eye Tracking Model (CETM)\n \n The model implements:\n 1. Representation enhancement: a two-layer convolutional block that refines the binary map representation.\n 2. Tracking predictor: three convolution blocks extract spatial features, which are then fed into a GRU to capture temporal context.\n Finally, a fully connected layer regresses the (x, y) coordinates.\n \n Input:\n - x: a tensor of shape (batch_size, seq_len, channels, height, width)\n where the channels correspond to the binary map representation (e.g., 1 or 2 channels)\n \n Output:\n - A tensor of shape (batch_size, seq_len, 2) with the predicted pupil coordinates.\n \"\"\"\n def __init__(self, args):\n super().__init__()\n self.args = args\n # --- Representation Enhancement Module ---\n # Enhance the input binary map (Bina-rep) with two convolutional layers.\n self.enhance_conv = nn.Sequential(\n nn.Conv2d(args.n_time_bins, 16, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(16),\n nn.ReLU(inplace=True),\n nn.Conv2d(16, args.n_time_bins, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(args.n_time_bins),\n nn.ReLU(inplace=True)\n )\n \n # --- Tracking Predictor ---\n # Three convolution blocks; each block: Conv2d -> BatchNorm2d -> ReLU -> Average Pooling.\n # These blocks are designed to extract pupil-space features from each frame.\n self.conv_block1 = nn.Sequential(\n nn.Conv2d(args.n_time_bins, 32, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(32),\n nn.ReLU(inplace=True),\n nn.AvgPool2d(kernel_size=2, stride=2)\n )\n self.conv_block2 = nn.Sequential(\n nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(64),\n nn.ReLU(inplace=True),\n nn.AvgPool2d(kernel_size=2, stride=2)\n )\n self.conv_block3 = nn.Sequential(\n nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(32),\n nn.ReLU(inplace=True),\n nn.AvgPool2d(kernel_size=2, stride=2)\n )\n \n # Assume the input event resolution is downsampled to 80x60.\n # For example, let’s assume height=60, width=80.\n # After conv_block1: (60/2, 80/2) -> (30, 40)\n # After conv_block2: (30/2, 40/2) -> (15, 20)\n # After conv_block3: (15/2, 20/2) -> (7, 10) [using floor division]\n # The number of output channels after block3 is 32.\n self.flattened_size = 32 * 7 * 10 # 32 channels * 7 * 10 spatial locations\n \n # GRU to capture temporal dependencies.\n self.gru = nn.GRU(input_size=self.flattened_size, hidden_size=128, num_layers=1, batch_first=True)\n \n # Fully connected layer to regress the final (x, y) coordinates.\n self.fc = nn.Linear(128, 2)\n \n def forward(self, x):\n # x shape: (batch_size, seq_len, channels, height, width)\n batch_size, seq_len, channels, height, width = x.shape\n \n # Representation enhancement: process each frame independently.\n x = x.view(batch_size * seq_len, channels, height, width)\n x = self.enhance_conv(x)\n \n # Tracking predictor: pass through three convolution blocks.\n x = self.conv_block1(x)\n x = self.conv_block2(x)\n x = self.conv_block3(x)\n \n # Flatten spatial dimensions: now shape becomes (batch_size, seq_len, flattened_size)\n x = x.view(batch_size, seq_len, -1)\n \n # Pass through GRU to capture temporal context.\n x, _ = self.gru(x) # Output shape: (batch_size, seq_len, hidden_size)\n \n # Final fully connected layer: regress pupil coordinates.\n x = self.fc(x) # Output shape: (batch_size, seq_len, 2)\n return x","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-03-09T21:13:21.139216Z","iopub.execute_input":"2025-03-09T21:13:21.139435Z","iopub.status.idle":"2025-03-09T21:13:21.147161Z","shell.execute_reply.started":"2025-03-09T21:13:21.139417Z","shell.execute_reply":"2025-03-09T21:13:21.146325Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#CETM Improved\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# ------------------------------\n# Define a Residual Convolution Block\n# ------------------------------\nclass ResidualConvBlock(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):\n super(ResidualConvBlock, self).__init__()\n self.conv = nn.Sequential(\n nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),\n nn.BatchNorm2d(out_channels),\n nn.ReLU(inplace=True),\n nn.Conv2d(out_channels, out_channels, kernel_size, stride, padding),\n nn.BatchNorm2d(out_channels)\n )\n self.relu = nn.ReLU(inplace=True)\n # If dimensions differ, use a 1x1 convolution to match\n if in_channels != out_channels:\n self.residual = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n else:\n self.residual = None\n\n def forward(self, x):\n identity = x\n out = self.conv(x)\n if self.residual is not None:\n identity = self.residual(identity)\n out += identity\n out = self.relu(out)\n return out\n\n# ------------------------------\n# Improved Consistent Eye Tracking Model (CETM_Improved)\n# ------------------------------\nclass CETM_Improved(nn.Module):\n \"\"\"\n Improved Consistent Eye Tracking Model (CETM_Improved)\n \n This model implements:\n 1. Representation enhancement using residual blocks.\n 2. A tracking predictor that uses three residual convolution blocks (each followed by average pooling)\n to extract spatial features from each frame.\n 3. A bidirectional GRU with dropout for temporal modeling.\n 4. A fully connected layer to regress the pupil (x, y) coordinates.\n \n Expected input shape: (batch, seq_len, channels, height, width)\n Output shape: (batch, seq_len, 2)\n \"\"\"\n def __init__(self, args):\n super(CETM_Improved, self).__init__()\n self.args = args\n # --- Representation Enhancement Module ---\n # Enhance the binary map representation with residual blocks.\n # (Assuming the input has args.n_time_bins channels.)\n self.enhance_conv = nn.Sequential(\n ResidualConvBlock(args.n_time_bins, 16),\n ResidualConvBlock(16, args.n_time_bins)\n )\n \n # --- Tracking Predictor ---\n # Use three residual blocks to extract spatial features.\n # For an input resolution of 80x60 (width x height):\n # After pool1: 40x30; after pool2: 20x15; after pool3: 10x7 (using kernel_size=2, stride=2).\n self.block1 = ResidualConvBlock(args.n_time_bins, 32)\n self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)\n self.block2 = ResidualConvBlock(32, 64)\n self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)\n self.block3 = ResidualConvBlock(64, 32)\n self.pool3 = nn.AvgPool2d(kernel_size=2, stride=2)\n \n # Compute the flattened size after three pooling layers:\n # Assuming input spatial dimensions are 80 (width) x 60 (height):\n # width: 80 -> 40 -> 20 -> 10; height: 60 -> 30 -> 15 -> 7 (using floor division)\n self.flattened_size = 32 * 10 * 7 # channels * width * height\n \n # --- Temporal Modeling ---\n # Use a bidirectional GRU with dropout to capture temporal context.\n # Bidirectional GRU doubles the output features (hidden_size * 2).\n self.gru = nn.GRU(\n input_size=self.flattened_size, \n hidden_size=128, \n num_layers=1, \n batch_first=True, \n bidirectional=True,\n dropout=0.3\n )\n # Fully connected layer maps GRU output (128*2=256) to (x, y)\n self.fc = nn.Linear(256, 2)\n \n def forward(self, x):\n # x shape: (batch, seq_len, channels, height, width)\n batch_size, seq_len, channels, height, width = x.shape\n # Process each frame independently by merging batch and sequence dimensions\n x = x.view(batch_size * seq_len, channels, height, width)\n \n # Representation enhancement\n x = self.enhance_conv(x)\n \n # Tracking predictor using residual convolution blocks and pooling\n x = self.block1(x)\n x = self.pool1(x)\n x = self.block2(x)\n x = self.pool2(x)\n x = self.block3(x)\n x = self.pool3(x)\n \n # Flatten the spatial dimensions\n x = x.view(batch_size, seq_len, -1) # (batch, seq_len, flattened_size)\n \n # Temporal modeling with bidirectional GRU\n x, _ = self.gru(x) # x shape: (batch, seq_len, 256)\n \n # Regress pupil coordinates for each time step\n x = self.fc(x) # (batch, seq_len, 2)\n return x\n\n# # ------------------------------\n# # Example Usage\n# # ------------------------------\n# if __name__ == \"__main__\":\n# # Dummy configuration (replace with your actual configuration)\n# class Args:\n# # Number of channels in the input event voxel representation (e.g., binary map)\n# n_time_bins = 1 \n# # Device configuration (use 'cuda' if available)\n# device = 'cuda' if torch.cuda.is_available() else 'cpu'\n# args = Args()\n \n# # Instantiate the improved model and move to device.\n# model = CETM_Improved(args).to(args.device)\n \n# # Create a dummy input: batch size = 4, sequence length = 10,\n# # channels = n_time_bins, height = 60, width = 80.\n# dummy_input = torch.randn(4, 10, args.n_time_bins, 60, 80).to(args.device)\n# output = model(dummy_input)\n# print(\"Output shape:\", output.shape) # Expected: (4, 10, 2)\n","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-03-10T14:59:34.457246Z","iopub.execute_input":"2025-03-10T14:59:34.457527Z","iopub.status.idle":"2025-03-10T14:59:34.468730Z","shell.execute_reply.started":"2025-03-10T14:59:34.457507Z","shell.execute_reply":"2025-03-10T14:59:34.467672Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CausalTemporalConv(nn.Module):\n \"\"\"\n Applies a causal convolution along the temporal dimension using 3D convolution.\n Input shape: (B, C, T, H, W). Uses padding only at the beginning of the temporal dimension.\n \"\"\"\n def __init__(self, in_channels, out_channels, kernel_size_t=3):\n super().__init__()\n self.kernel_size_t = kernel_size_t\n # For causal convolution, we pad (kernel_size_t - 1) at the beginning of the time dimension.\n self.conv = nn.Conv3d(\n in_channels,\n out_channels,\n kernel_size=(kernel_size_t, 1, 1),\n padding=(0, 0, 0) # Manual padding for causality.\n )\n \n def forward(self, x):\n # x: (B, C, T, H, W)\n pad = (0, 0, 0, 0, self.kernel_size_t - 1, 0) # Pad only at the beginning of time.\n x = F.pad(x, pad, mode='constant', value=0)\n return self.conv(x)\n\nclass SpatialConv(nn.Module):\n \"\"\"\n Applies a spatial convolution with kernel size (1, k, k) on input of shape (B, C, T, H, W).\n \"\"\"\n def __init__(self, in_channels, out_channels, kernel_size=3):\n super().__init__()\n padding = kernel_size // 2\n self.conv = nn.Conv3d(\n in_channels,\n out_channels,\n kernel_size=(1, kernel_size, kernel_size),\n padding=(0, padding, padding)\n )\n \n def forward(self, x):\n return self.conv(x)\n\nclass FactorizedSpatioTemporalBlock(nn.Module):\n \"\"\"\n A (1+2)D factorized spatio-temporal convolution block.\n First applies a causal temporal convolution, then a spatial convolution.\n \"\"\"\n def __init__(self, in_channels, out_channels, kernel_size_t=3, kernel_size_spatial=3):\n super().__init__()\n self.temporal_conv = CausalTemporalConv(in_channels, out_channels, kernel_size_t)\n self.spatial_conv = SpatialConv(out_channels, out_channels, kernel_size_spatial)\n self.bn = nn.BatchNorm3d(out_channels)\n self.relu = nn.ReLU(inplace=True)\n \n def forward(self, x):\n # x: (B, C, T, H, W)\n x = self.temporal_conv(x)\n x = self.spatial_conv(x)\n x = self.bn(x)\n x = self.relu(x)\n return x\n\n###############################################\n# Detector Head Inspired by CenterNet\n###############################################\n\nclass DetectorHead(nn.Module):\n \"\"\"\n The detector head consists of two spatial convolutional layers.\n It takes a feature map per time frame and outputs a grid (3x4) where each grid cell predicts:\n - a probability (pupil presence)\n - relative x and y offsets.\n \"\"\"\n def __init__(self, in_channels):\n super().__init__()\n self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)\n self.bn1 = nn.BatchNorm2d(in_channels)\n self.conv2 = nn.Conv2d(in_channels, 3, kernel_size=1) # 3 outputs per grid cell.\n \n def forward(self, x):\n # x: (B, C, H, W)\n x = F.relu(self.bn1(self.conv1(x)))\n x = self.conv2(x)\n return x\n\n###############################################\n# Complete BigBrains Model\n###############################################\n\nclass BBModel(nn.Module):\n \"\"\"\n Lightweight spatio-temporal network for online eye tracking.\n \n Implements:\n - A causal spatio-temporal backbone using (1+2)D factorized convolutions.\n - A detector head inspired by CenterNet that outputs a 3×4 grid for each frame.\n \n The temporal layers are causal and are designed to work with a FIFO buffer during streaming inference.\n \"\"\"\n def __init__(self, args):\n super().__init__()\n self.args = args\n in_channels = args.n_time_bins # e.g., number of channels in the event representation\n \n # Backbone: stack three factorized spatio-temporal blocks.\n self.block1 = FactorizedSpatioTemporalBlock(in_channels, 16, kernel_size_t=args.kernel_size_t, kernel_size_spatial=args.kernel_size_spatial)\n self.block2 = FactorizedSpatioTemporalBlock(16, 32, kernel_size_t=args.kernel_size_t, kernel_size_spatial=args.kernel_size_spatial)\n self.block3 = FactorizedSpatioTemporalBlock(32, 64, kernel_size_t=args.kernel_size_t, kernel_size_spatial=args.kernel_size_spatial)\n \n # Adaptive average pooling to produce a fixed spatial grid of 3x4 (time dimension remains unchanged).\n self.adapt_pool = nn.AdaptiveAvgPool3d((None, 3, 4))\n \n # Detector head to process each time frame.\n self.detector_head = DetectorHead(in_channels=64)\n \n def forward(self, x):\n # x: (B, T, channels, H, W)\n B, T, C, H, W = x.shape\n x = x.permute(0, 2, 1, 3, 4) # Convert to (B, channels, T, H, W)\n \n # Backbone: apply factorized spatio-temporal blocks.\n x = self.block1(x) # (B, 16, T, H, W)\n x = self.block2(x) # (B, 32, T, H, W)\n x = self.block3(x) # (B, 64, T, H, W)\n \n # Adaptive pooling to get fixed spatial grid of 3x4.\n x = self.adapt_pool(x) # (B, 64, T, 3, 4)\n \n # Process each time frame with the detector head.\n B, C, T, H_grid, W_grid = x.shape\n x = x.permute(0, 2, 1, 3, 4).contiguous() # (B, T, 64, 3, 4)\n x = x.view(B * T, C, H_grid, W_grid)\n x = self.detector_head(x) # (B*T, 3, 3, 4)\n x = x.view(B, T, 3, 3, 4)\n return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T18:34:42.803745Z","iopub.execute_input":"2025-03-10T18:34:42.804120Z","iopub.status.idle":"2025-03-10T18:34:42.819774Z","shell.execute_reply.started":"2025-03-10T18:34:42.804093Z","shell.execute_reply":"2025-03-10T18:34:42.818812Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Metrics\nimport torch\nimport torch.nn as nn\nimport numpy as np\n\n\ndef p_acc(target, prediction, width_scale, height_scale, pixel_tolerances=[1,3,5,10]):\n \"\"\"\n Calculate the accuracy of prediction\n :param target: (N, seq_len, 2) tensor, seq_len could be 1\n :param prediction: (N, seq_len, 2) tensor\n :return: a dictionary of p-total correct and batch size of this batch\n \"\"\"\n # flatten the N and seqlen dimension of target and prediction\n target = target.reshape(-1, 2)\n prediction = prediction.reshape(-1, 2)\n\n dis = target - prediction\n dis[:, 0] *= width_scale\n dis[:, 1] *= height_scale\n dist = torch.norm(dis, dim=-1)\n\n total_correct = {}\n for p_tolerance in pixel_tolerances:\n total_correct[f'p{p_tolerance}'] = torch.sum(dist < p_tolerance)\n\n bs_times_seqlen = target.shape[0]\n return total_correct, bs_times_seqlen\n\n\ndef p_acc_wo_closed_eye(target, prediction, width_scale, height_scale, pixel_tolerances=[1,3,5,10]):\n \"\"\"\n Calculate the accuracy of prediction, with p tolerance and only calculated on those with fully opened eyes\n :param target: (N, seqlen, 3) tensor\n :param prediction: (N, seqlen, 2) tensor, the last dimension is whether the eye is closed\n :return: a dictionary of p-total correct and batch size of this batch\n \"\"\"\n # flatten the N and seqlen dimension of target and prediction\n target = target.reshape(-1, 3)\n prediction = prediction.reshape(-1, 2)\n\n dis = target[:,:2] - prediction\n dis[:, 0] *= width_scale\n dis[:, 1] *= height_scale\n dist = torch.norm(dis, dim=-1)\n # check if there is nan in dist\n assert torch.sum(torch.isnan(dist)) == 0\n\n eye_closed = target[:,2] # 1 is closed eye\n # get the total number frames of those with fully opened eyes\n total_open_eye_frames = torch.sum(eye_closed == 0)\n\n # get the indices of those with closed eyes\n eye_closed_idx = torch.where(eye_closed == 1)[0]\n dist[eye_closed_idx] = np.inf\n total_correct = {}\n for p_tolerance in pixel_tolerances:\n total_correct[f'p{p_tolerance}'] = torch.sum(dist < p_tolerance)\n assert total_correct[f'p{p_tolerance}'] <= total_open_eye_frames\n\n return total_correct, total_open_eye_frames.item()\n\n\ndef px_euclidean_dist(target, prediction, width_scale, height_scale):\n \"\"\"\n Calculate the total pixel euclidean distance between target and prediction\n in a batch over the sequence length\n :param target: (N, seqlen, 3) tensor\n :param prediction: (N, seqlen, 2) tensor\n :return: a dictionary of p-total correct and batch size of this batch\n \"\"\"\n # flatten the N and seqlen dimension of target and prediction\n target = target.reshape(-1, 3)[:, :2]\n prediction = prediction.reshape(-1, 2)\n\n dis = target - prediction\n dis[:, 0] *= width_scale\n dis[:, 1] *= height_scale\n dist = torch.norm(dis, dim=-1)\n\n total_px_euclidean_dist = torch.sum(dist)\n sample_numbers = target.shape[0]\n return total_px_euclidean_dist, sample_numbers\n\n\nclass weighted_MSELoss(nn.Module):\n def __init__(self, weights, reduction='mean'):\n super().__init__()\n self.reduction = reduction\n self.weights = weights\n self.mseloss = nn.MSELoss(reduction='none')\n \n def forward(self, inputs, targets):\n batch_loss = self.mseloss(inputs, targets) * self.weights\n if self.reduction == 'mean':\n return torch.mean(batch_loss)\n elif self.reduction == 'sum':\n return torch.sum(batch_loss)\n else:\n return batch_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T19:59:13.191320Z","iopub.execute_input":"2025-03-10T19:59:13.191612Z","iopub.status.idle":"2025-03-10T19:59:13.201703Z","shell.execute_reply.started":"2025-03-10T19:59:13.191589Z","shell.execute_reply":"2025-03-10T19:59:13.200723Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# from model.BaselineEyeTrackingModel import CNN_GRU\n# from torch.utils.metrics import weighted_MSELoss\n\n# Initialize model and move it to the correct device (GPU/CPU)\nmodel = eval(args.architecture)(args).to(args.device)\n\n# Create the optimizer\noptimizer = optim.Adam(model.parameters(), lr=args.learning_rate)\n\n# # Select loss function based on args\n# if args.loss == \"mse\":\n# criterion = nn.MSELoss()\n# elif args.loss == \"weighted_mse\":\n# criterion = weighted_MSELoss(\n# weights=torch.tensor((args.sensor_width/args.sensor_height, 1)).to(args.device),\n# reduction='mean'\n# )\n# else:\n# raise ValueError(\"Invalid loss name\")\ndef euclidean_loss(pred, target, eps=1e-6):\n # Compute the Euclidean distance for each sample\n # pred and target are assumed to have shape (batch, seq_len, 2)\n loss = torch.sqrt(torch.sum((pred - target) ** 2, dim=-1) + eps)\n return torch.mean(loss)\n\nif args.loss == \"mse\":\n criterion = nn.MSELoss()\nelif args.loss == \"weighted_mse\":\n criterion = weighted_MSELoss(weights=torch.tensor((args.sensor_width/args.sensor_height, 1)).to(args.device), reduction='mean')\nelif args.loss == \"euclidean\":\n criterion = euclidean_loss # use our custom Euclidean distance loss\nelse:\n raise ValueError(\"Invalid loss name\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T18:35:34.879334Z","iopub.execute_input":"2025-03-10T18:35:34.879669Z","iopub.status.idle":"2025-03-10T18:35:36.863291Z","shell.execute_reply.started":"2025-03-10T18:35:34.879640Z","shell.execute_reply":"2025-03-10T18:35:36.862408Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n# import torch\nfrom tonic.slicers import (\n slice_events_by_time,\n)\nfrom typing import Any, List, Tuple\n\ndef custom_to_voxel_grid_numpy(events, sensor_size, n_time_bins=10):\n \"\"\"Build a voxel grid with bilinear interpolation in the time domain from a set of events.\n Implements the event volume from Zhu et al. 2019, Unsupervised event-based learning of optical\n flow, depth, and egomotion.\n\n Parameters:\n events: ndarray of shape [num_events, num_event_channels]\n sensor_size: size of the sensor that was used [W,H].\n n_time_bins: number of bins in the temporal axis of the voxel grid.\n\n Returns:\n numpy array of n event volumes (n,w,h,t)\n \"\"\"\n assert \"x\" and \"y\" and \"t\" and \"p\" in events.dtype.names\n assert sensor_size[2] == 2\n\n voxel_grid = np.zeros((n_time_bins, sensor_size[1], sensor_size[0]), float).ravel()\n\n # normalize the event timestamps so that they lie between 0 and n_time_bins\n time_diff = events[\"t\"][-1] - events[\"t\"][0]\n if time_diff < 1e-6: # avoid zero-vidision\n ts = np.zeros_like(events[\"t\"], dtype=float)\n else:\n ts = n_time_bins * (events[\"t\"].astype(float) - events[\"t\"][0]) / time_diff\n\n xs = events[\"x\"].astype(int)\n ys = events[\"y\"].astype(int)\n pols = events[\"p\"]\n pols[pols == 0] = -1 # polarity should be +1 / -1\n\n tis = ts.astype(int)\n dts = ts - tis\n vals_left = pols * (1.0 - dts)\n vals_right = pols * dts\n\n valid_indices = tis < n_time_bins\n np.add.at(\n voxel_grid,\n xs[valid_indices]\n + ys[valid_indices] * sensor_size[0]\n + tis[valid_indices] * sensor_size[0] * sensor_size[1],\n vals_left[valid_indices],\n )\n\n valid_indices = (tis + 1) < n_time_bins\n np.add.at(\n voxel_grid,\n xs[valid_indices]\n + ys[valid_indices] * sensor_size[0]\n + (tis[valid_indices] + 1) * sensor_size[0] * sensor_size[1],\n vals_right[valid_indices],\n )\n\n voxel_grid = np.reshape(\n voxel_grid, (n_time_bins, 1, sensor_size[1], sensor_size[0])\n )\n\n return voxel_grid\n\nclass SliceByTimeEventsTargets:\n \"\"\"\n Modified from tonic.slicers.SliceByTimeEventsTargets in the Tonic Library\n\n Slices an event array along fixed time window and overlap size. The number of bins depends\n on the length of the recording. Targets are copied.\n\n > \n >| window1 |\n > | window2 |\n\n Parameters:\n time_window (int): time for window length (same unit as event timestamps)\n overlap (int): overlap (same unit as event timestamps)\n include_incomplete (bool): include the last incomplete slice that has shorter time\n \"\"\"\n\n def __init__(self,time_window, overlap=0.0, seq_length=30, seq_stride=15, include_incomplete=False) -> None:\n self.time_window= time_window\n self.overlap= overlap\n self.seq_length=seq_length\n self.seq_stride=seq_stride\n self.include_incomplete=include_incomplete\n\n def slice(self, data: np.ndarray, targets: int) -> List[np.ndarray]:\n metadata = self.get_slice_metadata(data, targets)\n return self.slice_with_metadata(data, targets, metadata)\n\n def get_slice_metadata(\n self, data: np.ndarray, targets: int\n ) -> List[Tuple[int, int]]:\n t = data[\"t\"]\n stride = self.time_window - self.overlap\n assert stride > 0\n\n if self.include_incomplete:\n n_slices = int(np.ceil(((t[-1] - t[0]) - self.time_window) / stride) + 1)\n else:\n n_slices = int(np.floor(((t[-1] - t[0]) - self.time_window) / stride) + 1)\n n_slices = max(n_slices, 1) # for strides larger than recording time\n\n window_start_times = np.arange(n_slices) * stride + t[0]\n window_end_times = window_start_times + self.time_window\n indices_start = np.searchsorted(t, window_start_times)[:n_slices]\n indices_end = np.searchsorted(t, window_end_times)[:n_slices]\n\n if not self.include_incomplete:\n # get the strided indices for loading labels\n label_indices_start = np.arange(0, targets.shape[0]-self.seq_length, self.seq_stride)\n label_indices_end = label_indices_start + self.seq_length\n else:\n label_indices_start = np.arange(0, targets.shape[0], self.seq_stride)\n label_indices_end = label_indices_start + self.seq_length\n # the last label indices end should be the last label\n label_indices_end[-1] = targets.shape[0]\n\n assert targets.shape[0] >= label_indices_end[-1]\n\n return list(zip(zip(indices_start, indices_end), zip(label_indices_start, label_indices_end)))\n\n @staticmethod\n def slice_with_metadata(\n data: np.ndarray, targets: int, metadata: List[Tuple[Tuple[int, int], Tuple[int, int]]]\n ):\n return_data = []\n return_target = []\n for tuple1, tuple2 in metadata:\n return_data.append(data[tuple1[0]:tuple1[1]])\n return_target.append(targets[tuple2[0]:tuple2[1]])\n\n return return_data, return_target\n\n\nclass SliceLongEventsToShort:\n def __init__(self, time_window, overlap, include_incomplete):\n \"\"\"\n Initialize the transformation.\n\n Args:\n - time_window (int): The length of each sub-sequence.\n \"\"\"\n self.time_window = time_window\n self.overlap = overlap\n self.include_incomplete = include_incomplete\n\n def __call__(self, events):\n return slice_events_by_time(events, self.time_window, self.overlap, self.include_incomplete)\n\n\nclass EventSlicesToVoxelGrid:\n def __init__(self, sensor_size, n_time_bins, per_channel_normalize):\n \"\"\"\n Initialize the transformation.\n\n Args:\n - sensor_size (tuple): The size of the sensor.\n - n_time_bins (int): The number of time bins.\n \"\"\"\n self.sensor_size = sensor_size\n self.n_time_bins = n_time_bins\n self.per_channel_normalize = per_channel_normalize\n\n def __call__(self, event_slices):\n \"\"\"\n Apply the transformation to the given event slices.\n\n Args:\n - event_slices (Tensor): The input event slices.\n\n Returns:\n - Tensor: A batched tensor of voxel grids.\n \"\"\"\n voxel_grids = []\n for event_slice in event_slices:\n voxel_grid = custom_to_voxel_grid_numpy(event_slice, self.sensor_size, self.n_time_bins)\n voxel_grid = voxel_grid.squeeze(-3)\n if self.per_channel_normalize:\n # Calculate mean and standard deviation only at non-zero values\n non_zero_entries = (voxel_grid != 0)\n for c in range(voxel_grid.shape[0]):\n mean_c = voxel_grid[c][non_zero_entries[c]].mean()\n std_c = voxel_grid[c][non_zero_entries[c]].std()\n\n voxel_grid[c][non_zero_entries[c]] = (voxel_grid[c][non_zero_entries[c]] - mean_c) / (std_c + 1e-10)\n voxel_grids.append(voxel_grid)\n return np.array(voxel_grids).astype(np.float32)\n\n\nclass SplitSequence:\n def __init__(self, sub_seq_length, stride):\n \"\"\"\n Initialize the transformation.\n\n Args:\n - sub_seq_length (int): The length of each sub-sequence.\n - stride (int): The stride between sub-sequences.\n \"\"\"\n self.sub_seq_length = sub_seq_length\n self.stride = stride\n\n def __call__(self, sequence, labels):\n \"\"\"\n Apply the transformation to the given sequence and labels.\n\n Args:\n - sequence (Tensor): The input sequence of frames.\n - labels (Tensor): The corresponding labels.\n\n Returns:\n - Tensor: A batched tensor of sub-sequences.\n - Tensor: A batched tensor of corresponding labels.\n \"\"\"\n\n sub_sequences = []\n sub_labels = []\n\n for i in range(0, len(sequence) - self.sub_seq_length + 1, self.stride):\n sub_seq = sequence[i:i + self.sub_seq_length]\n sub_seq_labels = labels[i:i + self.sub_seq_length]\n sub_sequences.append(sub_seq)\n sub_labels.append(sub_seq_labels)\n\n return np.stack(sub_sequences), np.stack(sub_labels)\n \n\nclass SplitLabels:\n def __init__(self, sub_seq_length, stride):\n \"\"\"\n Initialize the transformation.\n\n Args:\n - sub_seq_length (int): The length of each sub-sequence.\n - stride (int): The stride between sub-sequences.\n \"\"\"\n self.sub_seq_length = sub_seq_length\n self.stride = stride\n # print(f\"stride is {self.stride}\")\n\n def __call__(self, labels):\n \"\"\"\n Apply the transformation to the given sequence and labels.\n\n Args:\n - labels (Tensor): The corresponding labels.\n\n Returns:\n - Tensor: A batched tensor of corresponding labels.\n \"\"\"\n sub_labels = []\n \n for i in range(0, len(labels) - self.sub_seq_length + 1, self.stride):\n sub_seq_labels = labels[i:i + self.sub_seq_length]\n sub_labels.append(sub_seq_labels)\n\n return np.stack(sub_labels)\n\nclass ScaleLabel:\n def __init__(self, scaling_factor):\n \"\"\"\n Initialize the transformation.\n\n Args:\n - scaling_factor (float): How much the spatial scaling was done on input\n \"\"\"\n self.scaling_factor = scaling_factor\n\n\n def __call__(self, labels):\n \"\"\"\n Apply the transformation to the given sequence and labels.\n\n Args:\n - labels (Tensor): The corresponding labels.\n\n Returns:\n - Tensor: A batched tensor of corresponding labels.\n \"\"\"\n labels[:,:2] = labels[:,:2] * self.scaling_factor\n return labels\n \nclass LabelTemporalSubsample:\n def __init__(self, temporal_subsample_factor):\n self.temp_subsample_factor = temporal_subsample_factor\n\n def __call__(self, labels):\n \"\"\"\n temorally subsample the labels\n \"\"\"\n interval = int(1/self.temp_subsample_factor)\n return labels[::interval]\n \n\nclass NormalizeLabel:\n def __init__(self, pseudo_width, pseudo_height):\n \"\"\"\n Initialize the transformation.\n\n Args:\n - scaling_factor (float): How much the spatial scaling was done on input\n \"\"\"\n self.pseudo_width = pseudo_width\n self.pseudo_height = pseudo_height\n \n def __call__(self, labels):\n \"\"\"\n Apply normalization on label, with pseudo width and height\n\n Args:\n - labels (Tensor): The corresponding labels.\n\n Returns:\n - Tensor: A batched tensor of corresponding labels.\n \"\"\"\n labels[:, 0] = labels[:, 0] / self.pseudo_width\n labels[:, 1] = labels[:, 1] / self.pseudo_height\n return labels\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T19:59:08.724645Z","iopub.execute_input":"2025-03-10T19:59:08.724941Z","iopub.status.idle":"2025-03-10T19:59:08.747224Z","shell.execute_reply.started":"2025-03-10T19:59:08.724919Z","shell.execute_reply":"2025-03-10T19:59:08.746459Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import tonic\nimport tonic.transforms as transforms\n# from dataset import ScaleLabel, NormalizeLabel, LabelTemporalSubsample\n\nfactor = args.spatial_factor # spatial downsample factor\ntemp_subsample_factor = args.temporal_subsample_factor # e.g., 5 to reduce 100Hz to 20Hz\n\nlabel_transform = transforms.Compose([\n ScaleLabel(factor),\n LabelTemporalSubsample(temp_subsample_factor),\n NormalizeLabel(pseudo_width=640*factor, pseudo_height=480*factor)\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T19:59:22.513592Z","iopub.execute_input":"2025-03-10T19:59:22.513878Z","iopub.status.idle":"2025-03-10T19:59:22.528006Z","shell.execute_reply.started":"2025-03-10T19:59:22.513857Z","shell.execute_reply":"2025-03-10T19:59:22.526675Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom typing import Any, Callable, Optional, Tuple\nimport h5py\nimport numpy as np\n\nfrom tonic.dataset import Dataset\n\nclass ThreeETplus_Eyetracking(Dataset):\n \"\"\"3ET DVS eye tracking `3ET `_\n ::\n\n @article{chen20233et,\n title={3ET: Efficient Event-based Eye Tracking using a Change-Based ConvLSTM Network},\n author={Chen, Qinyu and Wang, Zuowen and Liu, Shih-Chii and Gao, Chang},\n journal={arXiv preprint arXiv:2308.11771},\n year={2023}\n }\n\n authors: Qinyu Chen^{1,2}, Zuowen Wang^{1}\n affiliations: 1. Institute of Neuroinformatics, University of Zurich and ETH Zurich, Switzerland\n 2. Univeristy of Leiden, Netherlands\n\n Parameters:\n save_to (string): Location to save files to on disk.\n transform (callable, optional): A callable of transforms to apply to the data.\n split (string, optional): The dataset split to use, ``train`` or ``val``.\n target_transform (callable, optional): A callable of transforms to apply to the targets/labels.\n transforms (callable, optional): A callable of transforms that is applied to both data and\n labels at the same time.\n\n Returns:\n A dataset object that can be indexed or iterated over.\n One sample returns a tuple of (events, targets).\n \"\"\"\n\n sensor_size = (640, 480, 2)\n dtype = np.dtype([(\"t\", int), (\"x\", int), (\"y\", int), (\"p\", int)])\n ordering = dtype.names\n\n def __init__(\n self,\n save_to: str,\n split: str = \"train\",\n transform: Optional[Callable] = None,\n target_transform: Optional[Callable] = None,\n transforms: Optional[Callable] = None,\n ):\n super().__init__(\n save_to,\n transform=transform,\n target_transform=target_transform,\n transforms=transforms,\n )\n\n data_dir = save_to\n data_list_dir = '/kaggle/input/dataset'\n # Load filenames from the provided lists\n if split == \"train\":\n filenames = self.load_filenames(os.path.join(data_list_dir, \"train_files.txt\"))\n elif split == \"val\":\n filenames = self.load_filenames(os.path.join(data_list_dir, \"val_files.txt\"))\n elif split == \"test\":\n filenames = self.load_filenames(os.path.join(data_list_dir, \"test_files.txt\"))\n else:\n raise ValueError(\"Invalid split name\")\n\n # Get the data file paths and target file paths\n if split == \"train\" or split == \"val\":\n self.data = [os.path.join(data_dir, \"train\", f, f + \".h5\") for f in filenames]\n self.targets = [os.path.join(data_dir, \"train\", f, \"label.txt\") for f in filenames]\n elif split == \"test\":\n self.data = [os.path.join(data_dir, \"test\", f, f + \".h5\") for f in filenames]\n # for test set, we load the placeholder labels with all zeros\n self.targets = [os.path.join(data_dir, \"test\", f, \"label_zeros.txt\") for f in filenames]\n\n def __getitem__(self, index: int) -> Tuple[Any, Any]:\n \"\"\"\n Returns:\n (events, target) where target is index of the target class.\n \"\"\"\n # get events from .h5 file\n with h5py.File(self.data[index], \"r\") as f:\n # original events.dtype is dtype([('t', ' args.save_k_best:\n # sort all model checkpoints by validation loss in ascending order\n model_checkpoints = sorted([f for f in os.listdir(artifact_uri) if f.startswith(\"model_best_ep\")], \\\n key=lambda x: float(x.split(\"_\")[-1][:-4]))\n # delete the model checkpoint with the largest validation loss\n os.remove(os.path.join(artifact_uri, model_checkpoints[-1]))\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T18:36:28.014727Z","iopub.execute_input":"2025-03-10T18:36:28.015034Z","iopub.status.idle":"2025-03-10T18:36:28.035419Z","shell.execute_reply.started":"2025-03-10T18:36:28.015002Z","shell.execute_reply":"2025-03-10T18:36:28.034686Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# from utils.training_utils import train_epoch, validate_epoch, top_k_checkpoints\nfrom tqdm import tqdm\n\ndef train(model, train_loader, val_loader, criterion, optimizer, args):\n best_val_loss = float(\"inf\")\n for epoch in range(args.num_epochs):\n train_pbar = tqdm(train_loader, desc=f\"Training Epoch {epoch+1}/{args.num_epochs}\")\n model, train_loss, metrics = train_epoch(model, train_pbar, criterion, optimizer, args)\n mlflow.log_metric(\"train_loss\", train_loss, step=epoch)\n mlflow.log_metrics(metrics['tr_p_acc_all'], step=epoch)\n mlflow.log_metrics(metrics['tr_p_euc_error_all'], step=epoch)\n\n if args.val_interval > 0 and (epoch + 1) % args.val_interval == 0:\n val_pbar = tqdm(val_loader, desc=f\"Validation Epoch {epoch+1}/{args.num_epochs}\")\n val_loss, val_metrics = validate_epoch(model, val_pbar, criterion, args)\n if val_loss < best_val_loss:\n best_val_loss = val_loss\n torch.save(model.state_dict(), os.path.join(mlflow.get_artifact_uri(),\n f\"model_best_ep{epoch}_val_loss_{val_loss:.4f}.pth\"))\n top_k_checkpoints(args, mlflow.get_artifact_uri())\n print(f\"[Validation] Epoch {epoch+1}: Val Loss: {val_loss:.4f}\")\n mlflow.log_metric(\"val_loss\", val_loss, step=epoch)\n mlflow.log_metrics(val_metrics['val_p_acc_all'], step=epoch)\n mlflow.log_metrics(val_metrics['val_p_euc_error_all'], step=epoch)\n print(f\"Epoch {epoch+1}: Train Loss: {train_loss:.4f}\")\n return model\n\n# Run the training within an MLflow run:\nwith mlflow.start_run(run_name=args.run_name):\n mlflow.log_params(vars(args))\n with open(os.path.join(mlflow.get_artifact_uri(), \"args.json\"), 'w') as f:\n json.dump(vars(args), f)\n model = train(model, train_loader, val_loader, criterion, optimizer, args)\n torch.save(model.state_dict(), os.path.join(mlflow.get_artifact_uri(), f\"model_last_epoch{args.num_epochs}.pth\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T18:36:28.036261Z","iopub.execute_input":"2025-03-10T18:36:28.036479Z","iopub.status.idle":"2025-03-10T18:37:07.051428Z","shell.execute_reply.started":"2025-03-10T18:36:28.036460Z","shell.execute_reply":"2025-03-10T18:37:07.050104Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_config = {\n \"device\": \"cuda:0\",\n \"data_dir\": \"/kaggle/input/event-based-eye-tracking-cvpr-2025/event_data/event_data\",\n \"checkpoint\": \"/kaggle/working/mlruns/636469345475994249/c53becb1ad65418fb6a52bd597a34c15/artifacts/model_best_ep35_val_loss_0.0824.pth\",\n \"architecture\": \"CETM_Improved\",\n \"batch_size\": 1,\n \"spatial_factor\": 0.125,\n \"temporal_subsample_factor\": 1,\n \"pixel_tolerances\": [5,10,15],\n \"sensor_width\": 640,\n \"sensor_height\": 480,\n \"test_stride\": 30,\n \"test_length\": 30,\n \"n_time_bins\": 3,\n \"voxel_grid_ch_normaization\": False\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T17:34:41.714569Z","iopub.execute_input":"2025-03-10T17:34:41.714893Z","iopub.status.idle":"2025-03-10T17:34:41.719716Z","shell.execute_reply.started":"2025-03-10T17:34:41.714867Z","shell.execute_reply":"2025-03-10T17:34:41.718760Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nAuthor: Zuowen Wang\nAffiliation: Institute of Neuroinformatics, University of Zurich and ETH Zurich\nEmail: wangzu@ethz.ch\n\"\"\"\n\nimport argparse, json, os, mlflow, csv\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\n# from model.BaselineEyeTrackingModel import CNN_GRU\n# from dataset import ThreeETplus_Eyetracking, ScaleLabel, NormalizeLabel, \\\n# LabelTemporalSubsample, NormalizeLabel, SliceLongEventsToShort, \\\n# EventSlicesToVoxelGrid, SliceByTimeEventsTargets\nimport tonic.transforms as transforms\nfrom tonic import SlicedDataset, DiskCachedDataset\n\ndef main():\n # Load hyperparameters from JSON configuration file\n args = argparse.Namespace(**test_config)\n # if args.config_file:\n # with open(os.path.join('./configs', args.config_file), 'r') as f:\n # config = json.load(f)\n # # Overwrite hyperparameters with command-line arguments if provided\n # for key, value in vars(args).items():\n # if value is not None:\n # config[key] = value\n # args = argparse.Namespace(**config)\n # else:\n # raise ValueError(\"Please provide a JSON configuration file.\")\n\n # Dump the args to a JSON file in the MLflow artifact directory\n with open('/kaggle/working/mlruns/636469345475994249/c53becb1ad65418fb6a52bd597a34c15/artifacts/args.json', 'w') as f:\n json.dump(vars(args), f)\n\n # Define your model and move it to the desired device\n model = eval(args.architecture)(args).to(args.device)\n\n # For testing, we need to use the same spatial and temporal factors as in training.\n factor = args.spatial_factor\n temp_subsample_factor = args.temporal_subsample_factor\n\n # Define label transformation: scale, subsample temporally, and normalize the labels\n label_transform = transforms.Compose([\n ScaleLabel(factor),\n LabelTemporalSubsample(temp_subsample_factor),\n NormalizeLabel(pseudo_width=640*factor, pseudo_height=480*factor)\n ])\n\n # Load the raw test event data and apply spatial downsampling and label transformation\n test_data_orig = ThreeETplus_Eyetracking(\n save_to=args.data_dir, split=\"test\",\n transform=transforms.Downsample(spatial_factor=factor),\n target_transform=label_transform\n )\n\n # Compute slicing parameters for the test data (time window in microseconds)\n slicing_time_window = args.test_length * int(10000 / temp_subsample_factor) # microseconds\n\n # Create a slicer for the test sequences (no overlap; include incomplete sequences)\n test_slicer = SliceByTimeEventsTargets(\n slicing_time_window, overlap=0,\n seq_length=args.test_length, seq_stride=args.test_stride, include_incomplete=True\n )\n\n # Post-slicer transformation: convert slices to an event voxel grid representation\n post_slicer_transform = transforms.Compose([\n SliceLongEventsToShort(time_window=int(10000 / temp_subsample_factor), overlap=0, include_incomplete=True),\n EventSlicesToVoxelGrid(\n sensor_size=(int(640*factor), int(480*factor), 2),\n n_time_bins=args.n_time_bins, per_channel_normalize=args.voxel_grid_ch_normaization\n )\n ])\n\n # Create the sliced dataset for test data\n test_data = SlicedDataset(test_data_orig, test_slicer, transform=post_slicer_transform)\n\n # Optionally, you can cache the dataset to disk for faster subsequent runs\n # test_data = DiskCachedDataset(test_data, cache_path=f'./cached_dataset/test_l{args.test_length}s{args.test_stride}_ch{args.n_time_bins}')\n\n # For testing, ensure batch_size is 1 (this works in combination with include_incomplete=True)\n assert args.batch_size == 1, \"Batch size must be 1 for testing to avoid collate function errors.\"\n test_loader = DataLoader(\n test_data, batch_size=args.batch_size, shuffle=False,\n num_workers=int(os.cpu_count()-2)\n )\n\n # Load model weights from the provided checkpoint\n if args.checkpoint:\n model.load_state_dict(torch.load('/kaggle/working/mlruns/636469345475994249/c53becb1ad65418fb6a52bd597a34c15/artifacts/model_best_ep35_val_loss_0.0824.pth'))\n else:\n raise ValueError(\"Please provide a checkpoint file.\")\n\n # Run inference on the test set and write predictions to a CSV file.\n # with open('/kaggle/working/submission.csv', 'w', newline='') as csvfile:\n # csv_writer = csv.writer(csvfile, delimiter=',')\n # # Write header: row_id, x, y\n # csv_writer.writerow(['row_id', 'x', 'y'])\n # row_id = 0\n # for batch_idx, (data, target_placeholder) in enumerate(test_loader):\n # data = data.to(args.device)\n # output = model(data)\n \n # # Cast the output back to the downsampled sensor space (e.g., 640x480 scaled by factor)\n # output = output * torch.tensor((640*factor, 480*factor)).to(args.device)\n \n # # Use the output shape for iteration\n # n_samples, n_frames, _ = output.shape\n # for sample in range(n_samples):\n # for frame_id in range(n_frames):\n # row_to_write = output[sample][frame_id].tolist()\n # # Prepend the row_id to the output row\n # row_to_write.insert(0, row_id)\n # csv_writer.writerow(row_to_write)\n # row_id += 1\n # print(f\"Processed row {row_id}\")\n\n # Run inference on the test set and write predictions to a CSV file.\n with open('/kaggle/working/submission.csv', 'w', newline='') as csvfile:\n csv_writer = csv.writer(csvfile, delimiter=',')\n # Write header: row_id, x, y\n csv_writer.writerow(['row_id', 'x', 'y'])\n row_id = 0\n for batch_idx, (data, target_placeholder) in enumerate(test_loader):\n data = data.to(args.device)\n output = model(data)\n\n # Cast the output back to the downsampled sensor space (e.g., 640x480 scaled by factor)\n output = output * torch.tensor((640*factor, 480*factor)).to(args.device)\n\n n_samples = target_placeholder.shape[0]\n n_frames = target_placeholder.shape[1]\n for sample in range(n_samples):\n for frame_id in range(n_frames):\n row_to_write = output[sample][frame_id].tolist()\n # Prepend the row_id to the output row\n row_to_write.insert(0, row_id)\n csv_writer.writerow(row_to_write)\n row_id += 1\n print(f\"Processed row {row_id}\")\n\nif __name__ == \"__main__\":\n # parser = argparse.ArgumentParser()\n \n # Path to the configuration JSON file\n # parser.add_argument(\"--config_file\", type=str, default='test_config', \n # help=\"Path to JSON configuration file\")\n # # Path to the model checkpoint\n # parser.add_argument(\"--checkpoint\", type=str, help=\"Path to checkpoint\")\n # # Output CSV file path for submission predictions\n # parser.add_argument(\"--output_path\", type=str, default='./submission.csv')\n\n # args = parser.parse_args()\n\n main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T17:46:44.628027Z","iopub.execute_input":"2025-03-10T17:46:44.628386Z","iopub.status.idle":"2025-03-10T17:48:40.288587Z","shell.execute_reply.started":"2025-03-10T17:46:44.628345Z","shell.execute_reply":"2025-03-10T17:48:40.287462Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\n\nimport warnings\n\nwarnings.formatwarning = lambda message, category, filename, lineno, line=None: \\\n f'{category.__name__}: {message}\\n'\n\nclass CausalGroupNorm(nn.GroupNorm):\n \"\"\"A GroupNorm that does not use temporal statistics, to ensure causality\n \"\"\"\n def __init__(self, num_groups, num_channels, **kwargs):\n super().__init__(num_groups, num_channels, **kwargs)\n \n def forward(self, input):\n x = input.moveaxis(1, 2) # (B, T, C, H, W)\n x_shape = x.shape\n x = x.flatten(0, 1) # (B * T, C, H, W)\n x = super().forward(x).reshape(x_shape)\n return x.moveaxis(1, 2) # (B, C, T, H, W)\n\n\nact_layer = lambda: nn.ReLU()\nbn_block = lambda features: nn.Sequential(nn.BatchNorm3d(features), act_layer())\ngn_block = lambda features: nn.Sequential(CausalGroupNorm(4, features), act_layer())\npw_conv = lambda in_channels, out_channels: nn.Conv3d(in_channels, out_channels, 1, bias=False)\n\n\nclass SpatialBlock(nn.Module):\n def __init__(self, \n in_channels, \n out_channels, \n depthwise=False, \n kernel_size=1,\n full_conv3d=False, \n norms='mixed'):\n super().__init__()\n kernel = (kernel_size,3,3)\n self.kernel_size = kernel_size\n self.full_conv3d = full_conv3d\n self.norms = norms\n self.streaming_mode = False\n self.fifo = None # for streaming inference\n\n if self.norms=='all_gn':\n norm_block = gn_block\n else :\n norm_block = bn_block\n\n if depthwise:\n self.block = nn.Sequential(\n nn.Conv3d(in_channels, in_channels, kernel, (1, 2, 2), (0, 1, 1), groups=in_channels, bias=False), \n norm_block(in_channels), \n pw_conv(in_channels, out_channels), \n norm_block(out_channels), \n )\n \n else:\n self.block = nn.Sequential(\n nn.Conv3d(in_channels, out_channels, kernel, (1, 2, 2), (0, 1, 1), bias=False), \n norm_block(out_channels), \n )\n \n def streaming(self, enabled=True):\n if enabled:\n assert not self.training, \"Can only use streaming mode during evaluation.\"\n self.streaming_mode = enabled\n \n def reset_memory(self):\n self.fifo = None\n \n def forward(self, input):\n if self.full_conv3d: \n if self.streaming_mode:\n return self._streaming_forward(input)\n input = F.pad(input, (0, 0, 0, 0, self.kernel_size - 1, 0))\n return self.block(input)\n else: \n return self.block(input)\n \n def _streaming_forward(self, input):\n if self.fifo is None:\n self.fifo = torch.zeros(*input.shape[:2], self.kernel_size, *input.shape[3:]).type_as(input)\n self.fifo = torch.cat([self.fifo[:, :, 1:], input], dim=2)\n return self.block(self.fifo)\n\n\nclass TemporalBlock(nn.Module):\n def __init__(self, \n in_channels, \n out_channels, \n kernel_size=3, \n depthwise=False,\n full_conv3d=False, \n norms='mixed'):\n super().__init__()\n assert out_channels % 4 == 0 # needed for group norm to work\n self.kernel_size = kernel_size\n self.depthwise = depthwise\n self.norms = norms\n kernel = (kernel_size,3,3) if full_conv3d else (kernel_size,1,1)\n \n self.streaming_mode = False\n self.fifo = None # for streaming inference\n \n if self.norms=='mixed':\n norm1_block = bn_block\n norm2_block = gn_block\n elif self.norms=='all_bn':\n norm1_block = bn_block\n norm2_block = bn_block\n elif self.norms=='all_gn':\n norm1_block = gn_block\n norm2_block = gn_block\n\n if depthwise:\n self.block = nn.Sequential(\n nn.Conv3d(in_channels, in_channels, kernel, groups=in_channels, bias=False), \n norm1_block(in_channels), \n pw_conv(in_channels, out_channels), \n norm2_block(out_channels), \n )\n \n else:\n self.block = nn.Sequential(\n nn.Conv3d(in_channels, out_channels, kernel, bias=False), \n norm2_block(out_channels), \n )\n\n def streaming(self, enabled=True):\n if enabled:\n assert not self.training, \"Can only use streaming mode during evaluation.\"\n self.streaming_mode = enabled\n \n def reset_memory(self):\n self.fifo = None\n \n def forward(self, input):\n if self.streaming_mode:\n return self._streaming_forward(input)\n \n input = F.pad(input, (0, 0, 0, 0, self.kernel_size - 1, 0))\n return self.block(input)\n \n def _streaming_forward(self, input):\n if self.fifo is None:\n self.fifo = torch.zeros(*input.shape[:2], self.kernel_size, *input.shape[3:]).type_as(input)\n self.fifo = torch.cat([self.fifo[:, :, 1:], input], dim=2)\n return self.block(self.fifo)\n \n\nclass TennSt(nn.Module):\n def __init__(\n self, \n channels, \n t_kernel_size, \n n_depthwise_layers, \n detector_head, \n detector_depthwise, \n full_conv3d=False,\n norms='mixed',\n ):\n super().__init__()\n self.detector = detector_head\n \n depthwises = [False] * (10 - n_depthwise_layers) + [True] * n_depthwise_layers\n temporals = [True, False] * 5\n \n self.backbone = nn.Sequential()\n for i in range(len(depthwises)):\n in_channels, out_channels = channels[i], channels[i+1]\n depthwise = depthwises[i]\n temporal = temporals[i]\n \n if temporal:\n self.backbone.append(TemporalBlock(in_channels, out_channels, \n kernel_size=t_kernel_size, depthwise=depthwise,\n full_conv3d=full_conv3d, norms=norms))\n else:\n self.backbone.append(SpatialBlock(in_channels, out_channels, depthwise=depthwise,\n full_conv3d=full_conv3d,\n kernel_size=t_kernel_size if full_conv3d else 1,\n norms=norms))\n \n if detector_head:\n self.head = nn.Sequential(\n TemporalBlock(channels[-1], channels[-1], t_kernel_size, depthwise=detector_depthwise), \n nn.Conv3d(channels[-1], channels[-1], (1, 3, 3), (1, 1, 1), (0, 1, 1)), \n act_layer(), \n nn.Conv3d(channels[-1], 3, 1), \n )\n else:\n self.head = nn.Sequential(\n nn.Conv1d(channels[-1], channels[-1], 1), \n act_layer(), \n nn.Conv1d(channels[-1], 2, 1), \n )\n \n def streaming(self, enabled=True):\n if enabled:\n warnings.warn(\"You have enabled the streaming mode of the network. It is expected, but not checked, that the input will be of shape (batch, 1, H, W).\")\n for name, module in self.named_modules():\n if name and hasattr(module, 'streaming'):\n module.streaming(enabled)\n \n def reset_memory(self):\n for name, module in self.named_modules():\n if name and hasattr(module, 'reset_memory'):\n module.reset_memory()\n \n def forward(self, input: torch.Tensor) -> torch.Tensor:\n if self.detector:\n return self.head((self.backbone(input)))\n else:\n return self.head(self.backbone(input).mean((-2, -1)))\n ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T19:58:37.019041Z","iopub.execute_input":"2025-03-10T19:58:37.019353Z","iopub.status.idle":"2025-03-10T19:58:38.578384Z","shell.execute_reply.started":"2025-03-10T19:58:37.019329Z","shell.execute_reply":"2025-03-10T19:58:38.577722Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import argparse, json, os, mlflow\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nimport tonic.transforms as transforms\nfrom tonic import SlicedDataset, DiskCachedDataset\n\n# Import your dataset helper functions/classes.\n# from dataset import ThreeETplus_Eyetracking, ScaleLabel, NormalizeLabel, LabelTemporalSubsample, SliceLongEventsToShort, EventSlicesToVoxelGrid, SliceByTimeEventsTargets\n\n# Import the BigBrains model code (TennSt) from your module or paste the code above in a cell.\n# from bigbrains_model import TennSt\n# For the purpose of this notebook, assume TennSt is already defined in a previous cell.\n\n# -------------------------------\n# Custom Loss Function: Grid Loss\n# -------------------------------\n\ndef smooth_l1_loss(pred, target, beta=0.11):\n diff = torch.abs(pred - target)\n loss = torch.where(diff < beta, 0.5 * (diff ** 2) / beta, diff - 0.5 * beta)\n return loss\n\ndef focal_loss(pred, target, gamma=2):\n # pred: predicted probability, target: binary (0 or 1)\n eps = 1e-6\n loss = torch.where(target==1, -((1 - pred) ** gamma) * torch.log(pred + eps),\n -(pred ** gamma) * torch.log(1 - pred + eps))\n return loss\n\ndef grid_loss(pred, target, gamma=2, beta=0.11):\n \"\"\"\n pred and target shapes: (B, T, grid_y, grid_x, 3)\n The last dimension: index 0 = predicted probability,\n indices 1,2 = predicted offsets.\n \"\"\"\n p_pred = pred[..., 0]\n offsets_pred = pred[..., 1:]\n \n p_target = target[..., 0]\n offsets_target = target[..., 1:]\n \n focal = focal_loss(p_pred, p_target, gamma)\n reg = smooth_l1_loss(offsets_pred, offsets_target, beta)\n # Only apply regression loss for grid-cells with pupil presence (p_target == 1)\n reg = reg * p_target\n total_loss = focal + reg\n return total_loss.mean()\n\n# -------------------------------\n# Hyperparameter Configuration\n# -------------------------------\n\ndefault_config = {\n \"n_time_bins\": 2, # Number of channels in event representation.\n \"kernel_size_t\": 3, # Temporal conv kernel size.\n \"kernel_size_spatial\": 3, # Spatial conv kernel size.\n \"batch_size\": 32,\n \"num_frames\": 50, # Number of event frames per sample.\n \"learning_rate\": 0.002,\n \"weight_decay\": 0.005,\n \"num_epochs\": 200,\n \"spatial_factor\": 1, # Factor for spatial downsampling.\n \"temporal_subsample_factor\": 1, # Factor for temporal subsampling.\n \"train_length\": 50, # Number of frames in training sequence.\n \"train_stride\": 5, # Stride for slicing training data.\n \"voxel_grid_ch_normaization\": True,\n \"data_dir\": \"./data\", # Path to dataset.\n \"mlflow_path\": \"./mlruns\", # MLflow tracking directory.\n \"experiment_name\": \"BigBrains_EET\",\n \"run_name\": \"BigBrains_Training\",\n}\n\n# Optionally override defaults via argparse (if running from command line)\nparser = argparse.ArgumentParser()\nfor key, value in default_config.items():\n parser.add_argument(f\"--{key}\", type=type(value), default=value)\nargs = parser.parse_args()\nargs.device = 'cuda' if torch.cuda.is_available() else 'cpu'\n\n# -------------------------------\n# Training Pipeline\n# -------------------------------\n\ndef main():\n # Save hyperparameters to MLflow.\n mlflow.set_tracking_uri(args.mlflow_path)\n mlflow.set_experiment(args.experiment_name)\n with mlflow.start_run(run_name=args.run_name):\n with open(os.path.join(mlflow.get_artifact_uri(), \"args.json\"), \"w\") as f:\n json.dump(vars(args), f)\n \n # Instantiate the model (TennSt) with detector_head enabled.\n model = TennSt(\n channels=[args.n_time_bins, 16, 32, 64, 64], # Example channel progression.\n t_kernel_size=args.kernel_size_t,\n n_depthwise_layers=2, # Example value; adjust as needed.\n detector_head=True,\n detector_depthwise=False,\n full_conv3d=False,\n norms='mixed',\n ).to(args.device)\n \n # Define optimizer and (optionally) scheduler.\n optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)\n \n # Set up dataset and dataloader.\n factor = args.spatial_factor\n temp_subsample_factor = args.temporal_subsample_factor\n \n label_transform = transforms.Compose([\n ScaleLabel(factor),\n LabelTemporalSubsample(temp_subsample_factor),\n NormalizeLabel(pseudo_width=640*factor, pseudo_height=480*factor)\n ])\n \n train_data_orig = ThreeETplus_Eyetracking(\n save_to=args.data_dir, split=\"train\",\n transform=transforms.Downsample(spatial_factor=factor),\n target_transform=label_transform\n )\n \n slicing_time_window = args.train_length * int(10000 / temp_subsample_factor) # in microseconds\n train_slicer = SliceByTimeEventsTargets(\n slicing_time_window,\n overlap=slicing_time_window - int(10000 / temp_subsample_factor * args.train_stride),\n seq_length=args.train_length,\n seq_stride=args.train_stride,\n include_incomplete=False\n )\n \n post_slicer_transform = transforms.Compose([\n SliceLongEventsToShort(time_window=int(10000 / temp_subsample_factor), overlap=0, include_incomplete=True),\n EventSlicesToVoxelGrid(sensor_size=(int(640*factor), int(480*factor), 2),\n n_time_bins=args.n_time_bins,\n per_channel_normalize=args.voxel_grid_ch_normaization)\n ])\n \n train_data = SlicedDataset(train_data_orig, train_slicer, transform=post_slicer_transform)\n train_data = DiskCachedDataset(train_data, cache_path=f'./cached_dataset/train_l{args.train_length}_s{args.train_stride}_ch{args.n_time_bins}')\n \n train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)\n \n # Training Loop\n for epoch in range(args.num_epochs):\n model.train()\n epoch_loss = 0.0\n for batch_idx, (data, target) in enumerate(train_loader):\n data = data.to(args.device) # Expect shape: (B, T, C, H, W)\n target = target.to(args.device) # Expect shape: (B, T, grid_y, grid_x, 3)\n optimizer.zero_grad()\n output = model(data) # Output shape: (B, T, 3, 3, 4)\n loss = grid_loss(output, target, gamma=2, beta=0.11)\n loss.backward()\n optimizer.step()\n epoch_loss += loss.item()\n avg_loss = epoch_loss / len(train_loader)\n print(f\"Epoch {epoch+1}/{args.num_epochs}, Loss: {avg_loss:.4f}\")\n mlflow.log_metric(\"train_loss\", avg_loss, step=epoch)\n \n # Save final model.\n torch.save(model.state_dict(), os.path.join(mlflow.get_artifact_uri(), \"model_final.pth\"))\n \nif __name__ == \"__main__\":\n main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T19:59:50.231785Z","iopub.execute_input":"2025-03-10T19:59:50.232105Z","iopub.status.idle":"2025-03-10T19:59:50.250843Z","shell.execute_reply.started":"2025-03-10T19:59:50.232068Z","shell.execute_reply":"2025-03-10T19:59:50.249735Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}