{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "b8e236f1-385a-4d93-bb39-bea3ee384d76", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "PID of this process = 1877774\n" ] } ], "source": [ "# Import packages and setup gpu configuration.\n", "# This code block shouldnt need to be adjusted!\n", "import os\n", "import sys\n", "import json\n", "import yaml\n", "import numpy as np\n", "import copy\n", "import math\n", "import time\n", "import random\n", "from tqdm.auto import tqdm\n", "import webdataset as wds\n", "import matplotlib.pyplot as plt\n", "\n", "import torch\n", "import torch.nn as nn\n", "from torchvision import transforms\n", "import utils\n", "from mae_utils.flat_models import *\n", "\n", "# tf32 data type is faster than standard float32\n", "torch.backends.cuda.matmul.allow_tf32 = True\n", "# following fixes a Conv3D CUDNN_NOT_SUPPORTED error\n", "torch.backends.cudnn.benchmark = True\n", "\n", "# ## MODEL TO LOAD ##\n", "# model_name = \"HCPflat_large_gsrFalse_\"\n", "# parquet_folder = \"epoch99\"\n", "\n", "# # outdir = os.path.abspath(f'checkpoints/{model_name}')\n", "# outdir = os.path.abspath(f'checkpoints/{model_name}')\n", "\n", "# print(\"outdir\", outdir)\n", "# # Load previous config.yaml if available\n", "# if os.path.exists(f\"{outdir}/config.yaml\"):\n", "# config = yaml.load(open(f\"{outdir}/config.yaml\", 'r'), Loader=yaml.FullLoader)\n", "# print(f\"Loaded config.yaml from ckpt folder {outdir}\")\n", "# # create global variables from the config\n", "# print(\"\\n__CONFIG__\")\n", "# for attribute_name in config.keys():\n", "# print(f\"{attribute_name} = {config[attribute_name]}\")\n", "# globals()[attribute_name] = config[f'{attribute_name}']\n", "# print(\"\\n\")\n", "\n", "# world_size = os.getenv('WORLD_SIZE')\n", "# if world_size is None: \n", "# world_size = 1\n", "# else:\n", "# world_size = int(world_size)\n", "# print(f\"WORLD_SIZE={world_size}\")\n", "\n", "# if utils.is_interactive():\n", "# # Following allows you to change functions in models.py or utils.py and \n", "# # have this notebook automatically update with your revisions\n", "# %load_ext autoreload\n", "# %autoreload 2\n", "\n", "# batch_size = probe_batch_size\n", "# num_epochs = probe_num_epochs\n", "\n", "# data_type = torch.float32 # change depending on your mixed_precision\n", "# global_batch_size = batch_size * world_size\n", "\n", "device = torch.device('cuda')\n", "\n", "hcp_flat_path = \"/weka/proj-medarc/shared/HCP-Flat\"\n", "seed = 42\n", "num_frames = 16\n", "gsr = False\n", "num_workers = 10\n", "\n", "print(\"PID of this process =\",os.getpid())\n", "utils.seed_everything(seed)" ] }, { "cell_type": "code", "execution_count": 2, "id": "fc00235a-a9e1-4dc4-8fd3-359c1d91bfc7", "metadata": {}, "outputs": [], "source": [ "from torch.utils.data import default_collate\n", "from mae_utils.flat import load_hcp_flat_mask\n", "from mae_utils.flat import create_hcp_flat\n", "from mae_utils.flat import batch_unmask\n", "import mae_utils.visualize as vis\n", "\n", "\n", "# batch_size = 1\n", "# print(f\"changed batch_size to {batch_size}\")\n", "\n", "# ## Test ##\n", "# datasets_to_include = \"HCP\"\n", "# assert \"HCP\" in datasets_to_include\n", "# test_dataset = create_hcp_flat(root=hcp_flat_path, \n", "# clip_mode=\"event\", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'test')\n", "# test_dl = wds.WebLoader(\n", "# test_dataset.batched(batch_size, partial=False, collation_fn=default_collate),\n", "# batch_size=None,\n", "# shuffle=False,\n", "# num_workers=num_workers,\n", "# pin_memory=True,\n", "# )\n", "\n", "# ## Train ##\n", "# assert \"HCP\" in datasets_to_include\n", "# train_dataset = create_hcp_flat(root=hcp_flat_path, \n", "# clip_mode=\"event\", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'train')\n", "# train_dl = wds.WebLoader(\n", "# train_dataset.batched(batch_size, partial=False, collation_fn=default_collate),\n", "# batch_size=None,\n", "# shuffle=False,\n", "# num_workers=num_workers,\n", "# pin_memory=True,\n", "# )" ] }, { "cell_type": "code", "execution_count": 26, "id": "99d41181-bb61-40c1-ab94-63e732119f03", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "changed batch_size to 1\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "7f4fdcd311f7459aa8d5bb5f8f5723a2", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/120000 [00:00 8\u001b[0m \u001b[43maaa\u001b[49m\n\u001b[1;32m 10\u001b[0m \u001b[38;5;66;03m# Step 2: Serialize the dictionary to a JSON string\u001b[39;00m\n\u001b[1;32m 11\u001b[0m meta_str \u001b[38;5;241m=\u001b[39m json\u001b[38;5;241m.\u001b[39mdumps(flatten_meta(meta_serializable), indent\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m4\u001b[39m)\n", "\u001b[0;31mNameError\u001b[0m: name 'aaa' is not defined" ] } ], "source": [ "for i, batch in tqdm(enumerate(test_dl), total = 12000):\n", " images = batch['image']\n", " meta = batch['meta']\n", " batch_size = images.shape[0]\n", " meta_serializable = meta.copy()\n", "\n", " print(images.shape)\n", " aaa\n", " \n", " # Step 2: Serialize the dictionary to a JSON string\n", " meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)\n", " meta_array = np.append(meta_array, meta_str)\n", " if flatmaps_dset is None:\n", " # Initialize datasets with unlimited (None) maxshape along the first axis\n", " flatmaps_shape = (0,) + images.shape[1:]\n", " flatmaps_maxshape = (None,) + images.shape\n", "\n", " flatmaps_dset = h5f.create_dataset(\n", " 'flatmaps',\n", " shape=flatmaps_shape,\n", " maxshape=flatmaps_maxshape,\n", " dtype=np.float16,\n", " chunks=True # Enable chunking for efficient resizing\n", " )\n", "\n", " # Resize datasets to accommodate new data\n", " flatmaps_dset.resize(total_samples + batch_size, axis=0)\n", "\n", " # Write data to the datasets\n", " flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)\n", "\n", " total_samples += batch_size" ] }, { "cell_type": "code", "execution_count": 22, "id": "27b52f76-e197-46df-8cd6-e1013bcc3d45", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(0, 1, 1, 16, 144, 320)" ] }, "execution_count": 22, "metadata": {}, "output_type": "execute_result" } ], "source": [ "flatmaps_shape" ] }, { "cell_type": "code", "execution_count": 3, "id": "5aab380c-fa81-4fe1-9c56-3ce75e69a964", "metadata": {}, "outputs": [], "source": [ "# all_keys = []\n", "# for sample in tqdm(train_dl, desc=\"Processing samples\"):\n", "# all_keys = all_keys + sample['meta']['key']\n", "\n", "# subset_size = int(0.1 * len(all_keys))\n", "# # Select a random subset\n", "# val_keys = random.sample(all_keys, subset_size)\n", "\n", "# import pickle\n", "# with open('val_keys.pickle', 'wb') as handle:\n", "# pickle.dump(val_keys, handle, protocol=pickle.HIGHEST_PROTOCOL)" ] }, { "cell_type": "code", "execution_count": 4, "id": "f12db98b-3bea-43d9-8310-cc2d9fa41466", "metadata": {}, "outputs": [], "source": [ "import pickle\n", "with open('val_keys.pickle', 'rb') as handle:\n", " val_keys = pickle.load(handle)" ] }, { "cell_type": "code", "execution_count": 5, "id": "40f04009-13f1-473d-9808-181a46d3e97e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Number of classes: 21\n" ] } ], "source": [ "from sklearn.preprocessing import LabelEncoder\n", "\n", "INCLUDE_CONDS = {\n", " \"fear\",\n", " \"neut\",\n", " \"math\",\n", " \"story\",\n", " \"lf\",\n", " \"lh\",\n", " \"rf\",\n", " \"rh\",\n", " \"t\",\n", " \"match\",\n", " \"relation\",\n", " \"mental\",\n", " \"rnd\",\n", " \"0bk_body\",\n", " \"2bk_body\",\n", " \"0bk_faces\",\n", " \"2bk_faces\",\n", " \"0bk_places\",\n", " \"2bk_places\",\n", " \"0bk_tools\",\n", " \"2bk_tools\",\n", "}\n", "\n", "# test_data = []\n", "\n", "# # Iterate over the DataLoader with a progress bar\n", "# for sample in tqdm(train_dl, desc=\"Processing samples\"):\n", "# x = sample['image']\n", "# y = sample['meta']['trial_type']\n", "# key = sample['meta']['key']\n", "# print(x.shape, y, key)\n", "# break\n", "# Initialize the label encoder\n", "label_encoder = LabelEncoder()\n", "label_encoder.fit(sorted(INCLUDE_CONDS)) # Ensure consistent ordering\n", "\n", "num_classes = len(label_encoder.classes_)\n", "print(f\"Number of classes: {num_classes}\")" ] }, { "cell_type": "code", "execution_count": 6, "id": "74bc856e-1eeb-4340-91d8-838f50f463d2", "metadata": {}, "outputs": [], "source": [ "# Assuming all_keys and val_keys are already defined as per your initial code\n", "val_keys = set(val_keys)" ] }, { "cell_type": "code", "execution_count": 7, "id": "1a2c6527-fa75-42ba-be43-3d753ed8100c", "metadata": {}, "outputs": [ { "ename": "NameError", "evalue": "name 'train_dl' is not defined", "output_type": "error", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", "Cell \u001b[0;32mIn[7], line 14\u001b[0m\n\u001b[1;32m 10\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out \u001b[38;5;66;03m# Raw logits\u001b[39;00m\n\u001b[1;32m 12\u001b[0m \u001b[38;5;66;03m# Determine the input dimension from a single sample\u001b[39;00m\n\u001b[1;32m 13\u001b[0m \u001b[38;5;66;03m# Assuming images are of shape [1, 16, 144, 320]\u001b[39;00m\n\u001b[0;32m---> 14\u001b[0m sample_batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mnext\u001b[39m(\u001b[38;5;28miter\u001b[39m(\u001b[43mtrain_dl\u001b[49m))\n\u001b[1;32m 15\u001b[0m sample_image \u001b[38;5;241m=\u001b[39m sample_batch[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mimage\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;241m0\u001b[39m] \u001b[38;5;66;03m# Shape: [1, 16, 144, 320]\u001b[39;00m\n\u001b[1;32m 16\u001b[0m input_dim \u001b[38;5;241m=\u001b[39m sample_image\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m0\u001b[39m)\n", "\u001b[0;31mNameError\u001b[0m: name 'train_dl' is not defined" ] } ], "source": [ "class LinearClassifier(nn.Module):\n", " def __init__(self, input_dim, num_classes):\n", " super(LinearClassifier, self).__init__()\n", " self.linear = nn.Linear(input_dim, num_classes)\n", " \n", " def forward(self, x):\n", " # Flatten the input except for the batch dimension\n", " x = x.view(x.size(0), -1)\n", " out = self.linear(x)\n", " return out # Raw logits\n", "\n", "# Determine the input dimension from a single sample\n", "# Assuming images are of shape [1, 16, 144, 320]\n", "sample_batch = next(iter(train_dl))\n", "sample_image = sample_batch['image'][0] # Shape: [1, 16, 144, 320]\n", "input_dim = sample_image.view(-1).size(0)\n", "print(f\"Input dimension: {input_dim}\")\n" ] }, { "cell_type": "code", "execution_count": 8, "id": "b8983a0b-c810-4424-8c31-32b9921ebd9c", "metadata": {}, "outputs": [], "source": [ "# Initialize the model\n", "model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n", "\n", "# Move the model to GPU if available\n", "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", "model.to(device)\n", "\n", "# Define loss function\n", "criterion = nn.CrossEntropyLoss()\n", "\n", "# Define optimizer with L2 regularization (weight_decay)\n", "learning_rate = 1e-4\n", "weight_decay = 1e-5 # Adjust based on your needs\n", "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)\n" ] }, { "cell_type": "code", "execution_count": 9, "id": "f43535f4-0f6e-4655-b4f0-926416685358", "metadata": {}, "outputs": [], "source": [ "# num_epochs = 20 # Adjust as needed\n", "\n", "# for epoch in range(num_epochs):\n", "# model.train()\n", "# running_train_loss = 0.0\n", "# correct_train = 0\n", "# total_train = 0\n", "# step = 0\n", "# # Training Phase\n", "# for batch in tqdm(train_dl, desc=f\"Epoch {epoch+1}/{num_epochs} - Training\"):\n", "# images = batch['image'].to(device) # Shape: [batch_size, 1, 16, 144, 320]\n", "# labels = batch['meta']['trial_type'] # List of labels\n", "# keys = batch['meta']['key'] # List of keys\n", " \n", "# # Encode labels to integer indices\n", "# encoded_labels = label_encoder.transform(labels)\n", "# encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device) # Shape: [batch_size]\n", " \n", "# # Convert keys to a boolean mask for training samples (keys not in val_keys)\n", "# mask = torch.tensor([key in val_keys for key in keys], dtype=torch.bool).to(device)\n", " \n", "# if mask.sum() == 0:\n", "# continue # No training samples in this batch\n", " \n", "# # Select training samples\n", "# images_train = images[mask] # Shape: [num_train_samples, 1, 16, 144, 320]\n", "# labels_train = encoded_labels[mask] # Shape: [num_train_samples]\n", "# # Forward pass\n", "# outputs = model(images_train) # Shape: [num_train_samples, num_classes]\n", "\n", "# # print(outputs, labels_train)\n", "# # Compute loss\n", "# loss = criterion(outputs, labels_train)\n", " \n", "# # Backward pass and optimization\n", "# optimizer.zero_grad()\n", "# loss.backward()\n", "# optimizer.step()\n", " \n", "# # Accumulate loss\n", "# running_train_loss += loss.item() * images_train.size(0)\n", "\n", "# # print(predicted, labels_train)\n", "# # Calculate accuracy\n", "# _, predicted = torch.max(outputs, 1)\n", "# correct_train += (predicted == labels_train).sum().item()\n", "# total_train += labels_train.size(0)\n", " \n", "# step = step + 1\n", "# if step % 100 == 0:\n", "# step_accuracy = 100 * correct_step / total_step if total_step > 0 else 0.0\n", "# print(f\" Step {step}: Accuracy: {correct_train/total_train:.2f}%\")\n", "# correct_step = 0\n", "# total_step = 0\n", "\n", "# epoch_train_loss = running_train_loss / total_train if total_train > 0 else 0.0\n", "# train_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0\n", " \n", "# # Validation Phase\n", "# model.eval()\n", "# running_val_loss = 0.0\n", "# correct_val = 0\n", "# total_val = 0\n", " \n", "# with torch.no_grad():\n", "# for batch in tqdm(train_dl, desc=f\"Epoch {epoch+1}/{num_epochs} - Validation\"):\n", "# images = batch['image'].to(device)\n", "# labels = batch['meta']['trial_type']\n", "# keys = batch['meta']['key']\n", " \n", "# # Encode labels to integer indices\n", "# encoded_labels = label_encoder.transform(labels)\n", "# encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)\n", " \n", "# # Convert keys to a boolean mask for validation samples (keys in val_keys)\n", "# mask = torch.tensor([key not in val_keys for key in keys], dtype=torch.bool).to(device)\n", " \n", "# if mask.sum() == 0:\n", "# continue # No validation samples in this batch\n", " \n", "# # Select validation samples\n", "# images_val = images[mask]\n", "# labels_val = encoded_labels[mask]\n", " \n", "# # Forward pass\n", "# outputs = model(images_val)\n", " \n", "# # Compute loss\n", "# loss = criterion(outputs, labels_val)\n", " \n", "# # Accumulate loss\n", "# running_val_loss += loss.item() * images_val.size(0)\n", " \n", "# # Calculate accuracy\n", "# _, predicted = torch.max(outputs, 1)\n", "# correct_val += (predicted == labels_val).sum().item()\n", "# total_val += labels_val.size(0)\n", " \n", "# epoch_val_loss = running_val_loss / total_val if total_val > 0 else 0.0\n", "# val_accuracy = 100 * correct_val / total_val if total_val > 0 else 0.0\n", " \n", "# print(f\"Epoch [{epoch+1}/{num_epochs}] \"\n", "# f\"- Training Loss: {epoch_train_loss:.4f}, Training Accuracy: {train_accuracy:.2f}% \"\n", "# f\"- Validation Loss: {epoch_val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%\")\n" ] }, { "cell_type": "code", "execution_count": 9, "id": "5ca168be-fdf2-4387-9bff-5a52b0a9b7de", "metadata": {}, "outputs": [], "source": [ "for data in train_dl:\n", " pass\n", " break" ] }, { "cell_type": "code", "execution_count": 10, "id": "15aa9f03-cd4c-48b7-99db-d69d1f6a8b7c", "metadata": {}, "outputs": [], "source": [ "def flatten_meta(meta_dict):\n", " \"\"\"\n", " Flatten the meta dictionary by:\n", " - Replacing single-item lists with the item itself.\n", " - Converting tensors to scalar numbers.\n", " \"\"\"\n", " flattened = {}\n", " for key, value in meta_dict.items():\n", " if isinstance(value, list):\n", " if len(value) == 1:\n", " flattened[key] = value[0] # Replace list with its single item\n", " else:\n", " flattened[key] = value # Keep as is if multiple items\n", " elif isinstance(value, torch.Tensor):\n", " # Convert tensor to scalar\n", " if value.numel() == 1:\n", " flattened[key] = value.item()\n", " else:\n", " flattened[key] = value.tolist() # Convert multi-element tensor to list\n", " else:\n", " flattened[key] = value # Keep the value as is\n", " return flattened" ] }, { "cell_type": "code", "execution_count": 25, "id": "26f0e52d-fd22-44c1-b086-c70b140fdf07", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "e1a9e0719ed040c7abbf2ca1704a5950", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/120000 [00:00 0 else 0.0\n", " train_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0\n", " \n", " # Validation Phase\n", " model.eval()\n", " running_val_loss = 0.0\n", " correct_val = 0\n", " total_val = 0\n", " \n", " with torch.no_grad():\n", " for batch in tqdm(test_dl, desc=f\"Epoch {epoch+1}/{num_epochs} - Validation\"):\n", " images = batch['image'].to(device)\n", " labels = batch['meta']['trial_type']\n", " keys = batch['meta']['key']\n", " \n", " # Encode labels to integer indices\n", " encoded_labels = label_encoder.transform(labels)\n", " encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)\n", " \n", " # # Convert keys to a boolean mask for validation samples (keys in val_keys)\n", " # mask = torch.tensor([key not in val_keys for key in keys], dtype=torch.bool).to(device)\n", " \n", " # if mask.sum() == 0:\n", " # continue # No validation samples in this batch\n", " \n", " # # Select validation samples\n", " images_val = images#[mask]\n", " labels_val = encoded_labels#[mask]\n", " \n", " # Forward pass\n", " outputs = model(images_val)\n", " \n", " # Compute loss\n", " loss = criterion(outputs, labels_val)\n", " \n", " # Accumulate loss\n", " running_val_loss += loss.item() * images_val.size(0)\n", " \n", " # Calculate accuracy\n", " _, predicted = torch.max(outputs, 1)\n", " correct_val += (predicted == labels_val).sum().item()\n", " total_val += labels_val.size(0)\n", " \n", " epoch_val_loss = running_val_loss / total_val if total_val > 0 else 0.0\n", " val_accuracy = 100 * correct_val / total_val if total_val > 0 else 0.0\n", " \n", " print(f\"Epoch [{epoch+1}/{num_epochs}] \"\n", " f\"- Training Loss: {epoch_train_loss:.4f}, Training Accuracy: {train_accuracy:.2f}% \"\n", " f\"- Validation Loss: {epoch_val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%\")\n" ] }, { "cell_type": "markdown", "id": "ab15aca0-148e-435f-b8f2-7a708b61a6d9", "metadata": {}, "source": [ "# hcp_flat" ] }, { "cell_type": "code", "execution_count": null, "id": "de1a5b87-fa69-44e1-bdb9-bd9e257485c6", "metadata": { "tags": [] }, "outputs": [], "source": [ "# from mae_utils.flat import load_hcp_flat_mask\n", "# from mae_utils.flat import create_hcp_flat\n", "# import mae_utils.visualize as vis\n", "\n", "# if utils.is_interactive(): # Use less samples per epoch for debugging\n", "# probe_num_samples_per_epoch = 100000\n", "# test_num_samples_per_epoch = 100000\n", "# num_epochs = 10\n", "\n", "\n", "# # Load ckpt\n", "# if not os.path.exists(outdir) or not os.path.isdir(outdir):\n", "# assert True, (f\"\\nCheckpoint folder {outdir} does not exist.\\n\")\n", "# else:\n", "# checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n", "\n", "# # Find the latest ckpt to load\n", "# epoch_numbers = []\n", "# for file in checkpoint_files:\n", "# try:\n", "# epoch_number = int(file.split('epoch')[-1].split('.')[0])\n", "# epoch_numbers.append(epoch_number)\n", "# except ValueError:\n", "# continue\n", "# latest_epoch = max(epoch_numbers)\n", "# checkpoint_name = f\"epoch{latest_epoch}.pth\"\n", " \n", "# ### Or provide the specific checkpoint you want to load\n", "# # checkpoint_name = \"epoch10.pth\" #\"epoch15.pth\"\n", "\n", "# # Load the checkpoint\n", "# checkpoint_path = os.path.join(outdir, checkpoint_name)\n", "# state = torch.load(checkpoint_path)\n", "\n", "# model = mae_vit_large_fmri(\n", "# patch_size=16,\n", "# decoder_embed_dim=decoder_embed_dim,\n", "# t_patch_size=t_patch_size,\n", "# pred_t_dim=pred_t_dim,\n", "# decoder_depth=4,\n", "# cls_embed=cls_embed,\n", "# norm_pix_loss=norm_pix_loss,\n", "# no_qkv_bias=no_qkv_bias,\n", "# sep_pos_embed=sep_pos_embed,\n", "# trunc_init=trunc_init,\n", "# img_mask=state[\"model_state_dict\"]['img_mask']\n", "# )\n", "\n", "# model.load_state_dict(state[\"model_state_dict\"], strict=True) #model_state_dict\n", "# print(f\"\\nLoaded checkpoint {checkpoint_name} from {outdir}\\n\")\n", "\n", "# model.eval()\n", "# model.requires_grad_(False)\n", "# model.to(device)\n", "# pass" ] }, { "cell_type": "code", "execution_count": null, "id": "c3461199-e805-4e9c-8c91-894e83cf8bc3", "metadata": { "tags": [] }, "outputs": [], "source": [ "# import argparse\n", "# import json\n", "# import os\n", "# import pickle\n", "# from pathlib import Path\n", "\n", "# import pandas as pd\n", "# import numpy as np\n", "# from sklearn.decomposition import PCA\n", "# from sklearn.linear_model import LogisticRegressionCV\n", "# from sklearn.model_selection import train_test_split\n", "# from sklearn.preprocessing import LabelEncoder\n", "\n", "# target = \"trial_type\"\n", "# print(f\"Target: {target}\")\n", "\n", "# train_features = pd.read_parquet(f\"{outdir}/{parquet_folder}/HCP/train.parquet\")\n", "# test_features = pd.read_parquet(f\"{outdir}/{parquet_folder}/HCP_/test.parquet\")\n", "\n", "# print(f\"train: {train_features.shape}, test: {test_features.shape}\")\n", "# print(f\"test: {test_features.shape}\")\n", "\n", "# X_train = np.stack(train_features[\"feature\"])\n", "# X_test = np.stack(test_features[\"feature\"])\n", "# print(f\"X_train: {X_train.shape}, X_test: {X_test.shape}\")\n", "# print(f\"X_test: {X_test.shape}\")\n", "\n", "\n", "# if target == \"task\":\n", "# labels_train = train_features[\"task\"].str.rstrip(\"1234\").values\n", "# labels_test = test_features[\"task\"].str.rstrip(\"1234\").values\n", "# elif target == \"trial_type\":\n", "# labels_train = train_features[\"trial_type\"].values\n", "# labels_test = test_features[\"trial_type\"].values\n", "\n", "# label_enc = LabelEncoder()\n", "# y_train = label_enc.fit_transform(labels_train)\n", "# y_test = label_enc.transform(labels_test)\n", "\n", "# print(f\"classes ({len(label_enc.classes_)}): {label_enc.classes_}\")\n", "# print(\n", "# f\"\\ny_train: {y_train.shape} {y_train[:20]}\\n\"\n", "# f\"y_test: {y_test.shape} {y_test[:20]}\"\n", "# )\n", "# del train_features, test_features\n", "\n", "# train_ind, val_ind = train_test_split(\n", "# np.arange(len(X_train)), train_size=0.9, random_state=42\n", "# )\n", "# print(\n", "# f\"\\ntrain_ind: {len(train_ind)} {train_ind[:10]}\\n\"\n", "# f\"val_ind: {len(val_ind)} {val_ind[:10]}\"\n", "# )\n", "# X_train, X_val = X_train[train_ind], X_train[val_ind]\n", "# y_train, y_val = y_train[train_ind], y_train[val_ind]\n", "\n", "# print(\"Fitting PCA projection\")\n", "# pca = PCA(n_components=384, whiten=True, svd_solver=\"randomized\")\n", "# pca.fit(X_train)\n", "\n", "# X_train = pca.transform(X_train)\n", "# X_val = pca.transform(X_val)\n", "# X_test = pca.transform(X_test)\n", "\n", "# print(\"Fitting logistic regression\")\n", "# clf = LogisticRegressionCV()\n", "# clf.fit(X_train, y_train)\n", "\n", "# train_acc = clf.score(X_train, y_train)\n", "# val_acc = clf.score(X_val, y_val)\n", "# test_acc = clf.score(X_test, y_test)\n", "\n", "# result = {\n", "# \"target\": target,\n", "# \"train_acc\": train_acc,\n", "# \"val_acc\": val_acc,\n", "# \"test_acc\": test_acc,\n", "# }\n", "\n", "# with open(f\"{outdir}/{parquet_folder}/HCP/downstream.json\", 'w') as out_json:\n", "# json.dump(result, out_json)\n", "\n", "# print(f\"Done:\\n{json.dumps(result)}\")" ] } ], "metadata": { "kernelspec": { "display_name": "foundation_env", "language": "python", "name": "foundation_env" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.10" } }, "nbformat": 4, "nbformat_minor": 5 }