{ "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 = 517997\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", "import pandas as pd\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", "import h5py\n", "from typing import List, Dict, Any, Tuple\n", "from sklearn.preprocessing import StandardScaler\n", "import argparse\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 = 5\n", "batch_size = 128\n", "# target = 'sex' # This can be 'trial_type' 'age' 'sex'\n", "\n", "print(\"PID of this process =\",os.getpid())" ] }, { "cell_type": "code", "execution_count": 15, "id": "c2539698-5566-48cf-8022-90cb70d94f9f", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "model_name_suffix: testing\n", "--hcp_flat_path=/weka/proj-medarc/shared/HCP-Flat --target=sex --model_suffix=testing --batch_size=128 --max_lr=3e-4 --num_epochs=40 --no-save_ckpt --no-wandb_log --num_workers=10 --weight_decay=1e-5\n", "The autoreload extension is already loaded. To reload it, use:\n", " %reload_ext autoreload\n" ] } ], "source": [ "# if running this interactively, can specify jupyter_args here for argparser to use\n", "if utils.is_interactive():\n", " model_name_suffix = \"testing\"\n", " print(\"model_name_suffix:\", model_name_suffix)\n", " \n", " # global_batch_size and batch_size should already be defined in the 2nd cell block\n", " jupyter_args = f\"--hcp_flat_path=/weka/proj-medarc/shared/HCP-Flat \\\n", " --target=sex \\\n", " --model_suffix={model_name_suffix} \\\n", " --batch_size={batch_size} \\\n", " --max_lr=3e-4 --num_epochs=40 --no-save_ckpt --no-wandb_log --num_workers=10 \\\n", " --weight_decay=1e-5\"\n", " # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow\n", "\n", " print(jupyter_args)\n", " jupyter_args = jupyter_args.split()\n", " \n", " from IPython.display import clear_output # function to clear print outputs in cell\n", " %load_ext autoreload \n", " # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions\n", " %autoreload 2 " ] }, { "cell_type": "code", "execution_count": 16, "id": "e3b9ee45-5409-4a03-b1a1-f59af40ccc79", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "------ ARGS ------- \n", " Namespace(model_suffix='testing', hcp_flat_path='/weka/proj-medarc/shared/HCP-Flat', batch_size=128, wandb_log=False, num_epochs=40, lr_scheduler_type='cycle', save_ckpt=False, seed=42, max_lr=0.0003, target='sex', num_workers=10, weight_decay=1e-05)\n" ] } ], "source": [ "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n", "parser.add_argument(\n", " \"--model_suffix\", type=str, default=\"Testing_flat\",\n", " help=\"name of model, used for ckpt saving and wandb logging (if enabled)\",\n", ")\n", "parser.add_argument(\n", " \"--hcp_flat_path\", type=str, default=os.getcwd(),\n", " help=\"Path to where NSD data is stored / where to download it to\",\n", ")\n", "parser.add_argument(\n", " \"--batch_size\", type=int, default=128,\n", " help=\"Batch size can be increased by 10x if only training retreival submodule and not diffusion prior\",\n", ")\n", "parser.add_argument(\n", " \"--wandb_log\",action=argparse.BooleanOptionalAction,default=False,\n", " help=\"whether to log to wandb\",\n", ")\n", "parser.add_argument(\n", " \"--num_epochs\",type=int,default=150,\n", " help=\"number of epochs of training\",\n", ")\n", "parser.add_argument(\n", " \"--lr_scheduler_type\",type=str,default='cycle',choices=['cycle','linear'],\n", ")\n", "parser.add_argument(\n", " \"--save_ckpt\",action=argparse.BooleanOptionalAction,default=True,\n", ")\n", "parser.add_argument(\n", " \"--seed\",type=int,default=42,\n", ")\n", "parser.add_argument(\n", " \"--max_lr\",type=float,default=3e-4,\n", ")\n", "parser.add_argument(\n", " \"--target\",type=str,default='trial_type',choices=['trial_type','sex','age'],\n", ")\n", "parser.add_argument(\n", " \"--num_workers\",type=int,default=10,\n", ")\n", "parser.add_argument(\n", " \"--weight_decay\",type=float,default=1e-5,\n", ")\n", "\n", "if utils.is_interactive():\n", " args = parser.parse_args(jupyter_args)\n", "else:\n", " args = parser.parse_args()\n", "\n", "print(f\"------ ARGS ------- \\n {args}\")\n", "\n", "# create global variables without the args prefix\n", "for attribute_name in vars(args).keys():\n", " globals()[attribute_name] = getattr(args, attribute_name)\n", " \n", "# seed all random functions\n", "utils.seed_everything(seed)\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "fc00235a-a9e1-4dc4-8fd3-359c1d91bfc7", "metadata": {}, "outputs": [], "source": [ "#### UNCOMMENT THIS TO SAVE THE HCP-FLAT IN HDF5 FORMAT\n", "\n", "\n", "# 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 = 26\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", "# )\n", "\n", "# 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\n", "\n", "# import h5py\n", "# meta_array = np.array([], dtype=object)\n", "# # Open an HDF5 file in write mode\n", "# with h5py.File('train_hcp_raw_flatmaps.hdf5', 'w') as h5f:\n", "# flatmaps_dset = None\n", " \n", "# total_samples = 0\n", "\n", "# for i, batch in tqdm(enumerate(train_dl), total = 120000):\n", "# images = batch['image'][0]\n", "# meta = batch['meta']\n", "# batch_size = images.shape[0]\n", "# meta_serializable = meta.copy()\n", " \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[1:]\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\n", " \n", "# print(f\"Processed {total_samples} samples\")\n", "# np.save('metadata_test_HCP_raw_flatmaps.npy', meta_array)\n", "\n", "\n", "# import h5py\n", "# meta_array = np.array([], dtype=object)\n", "# # Open an HDF5 file in write mode\n", "# with h5py.File('test_hcp_raw_flatmaps.hdf5', 'w') as h5f:\n", "# flatmaps_dset = None\n", " \n", "# total_samples = 0\n", "\n", "# for i, batch in tqdm(enumerate(test_dl), total = 12000):\n", "# images = batch['image'][0]\n", "# meta = batch['meta']\n", "# batch_size = images.shape[0]\n", "# meta_serializable = meta.copy()\n", " \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[1:]\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\n", " \n", "# print(f\"Processed {total_samples} samples\")\n", "# np.save('metadata_train_HCP_raw_flatmaps.npy', meta_array)" ] }, { "cell_type": "markdown", "id": "98629867-6d64-45ef-823d-5f4fdb279ec7", "metadata": {}, "source": [ "### Data" ] }, { "cell_type": "code", "execution_count": 5, "id": "6e66e2c0-6a73-42d7-8414-3b8acafdb0f4", "metadata": {}, "outputs": [], "source": [ "f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp_raw_flatmaps.hdf5', 'r')\n", "flatmaps_train = f_train['flatmaps']\n", "\n", "f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp_raw_flatmaps.hdf5', 'r')\n", "flatmaps_test = f_test['flatmaps']\n", "\n", "metadata_train = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_train_HCP_raw_flatmaps.npy', allow_pickle=True)\n", "metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP_raw_flatmaps.npy', allow_pickle=True)" ] }, { "cell_type": "code", "execution_count": 6, "id": "9788d371-c6ce-44d9-8ebc-d0a13d202378", "metadata": {}, "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.array(flatmaps_train[0:5000])\n", "# # flatten the flatmaps\n", "# X_train = X_train.reshape(X_train.shape[0], -1)\n", "# X_test = np.array(flatmaps_test[0:1000])\n", "# X_test = X_test.reshape(X_test.shape[0], -1)\n", "\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", "# labels_train = [json.loads(string)['trial_type'] for string in metadata_train[0:5000]]\n", "# labels_test = [json.loads(string)['trial_type'] for string in metadata_test[0:1000]]\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", "# print(f\"Done:\\n{json.dumps(result)}\")\n", "# with open(f\"{outdir}/{parquet_folder}/HCP/downstream.json\", 'w') as out_json:\n", "# json.dump(result, out_json)" ] }, { "cell_type": "markdown", "id": "a368c4fa-f2f3-4997-8ea4-ae6542efc32c", "metadata": {}, "source": [ "### Create the dataloader" ] }, { "cell_type": "code", "execution_count": 7, "id": "b52b668f-c143-40fe-b3d4-c21c7a406d0d", "metadata": {}, "outputs": [], "source": [ "from torch.utils.data import Dataset, DataLoader\n", "\n", "class HCPFlatDataset(Dataset):\n", " def __init__(self, flatmaps, metadata):\n", " self.flatmaps = flatmaps\n", " self.metadata = metadata\n", "\n", " def __len__(self):\n", " return len(self.metadata)\n", "\n", " def __getitem__(self, idx):\n", " return self.flatmaps[idx], json.loads(self.metadata[idx])\n", "\n", "# Loading to cpu for faster training, this can take several minutes. Remove this [:] if you want to move one at the time.\n", "train_dataset = HCPFlatDataset(flatmaps_train[:], metadata_train)\n", "train_dl = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)\n", "\n", "test_dataset = HCPFlatDataset(flatmaps_test[:], metadata_test)\n", "test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)\n" ] }, { "cell_type": "markdown", "id": "247a9b16-3c06-4c40-8c81-5b6169ef166e", "metadata": {}, "source": [ "### Load subject information" ] }, { "cell_type": "code", "execution_count": 17, "id": "a7d8760b-2e62-4ba5-ae34-b42376ad0eaf", "metadata": {}, "outputs": [], "source": [ "# open the file containing subject information\n", "if target == \"age\" or target == \"sex\":\n", " subject_information_HCP_path = os.path.join(hcp_flat_path, \"subjects_data_restricted.csv\")\n", " try:\n", " subject_information_HCP = pd.read_csv(subject_information_HCP_path)\n", " except:\n", " try:\n", " subject_information_HCP = pd.read_csv('./unrestricted_clane9_4_23_2024_13_28_14.csv') \n", " except:\n", " assert False, \"Subject information file not found\"\n", "\n", " ###### This is for unrestricted\n", " # age_related_columns = [\n", " # 'Age', 'PicSeq_AgeAdj', 'CardSort_AgeAdj', 'Flanker_AgeAdj',\n", " # 'ReadEng_AgeAdj', 'PicVocab_AgeAdj', 'ProcSpeed_AgeAdj',\n", " # 'CogFluidComp_AgeAdj', 'CogEarlyComp_AgeAdj', 'CogTotalComp_AgeAdj',\n", " # 'CogCrystalComp_AgeAdj', 'Endurance_AgeAdj', 'Dexterity_AgeAdj',\n", " # 'Strength_AgeAdj', 'Odor_AgeAdj', 'Taste_AgeAdj'\n", " # ]\n", " \n", " # sex_related_columns = [\n", " # 'Gender'\n", " # ]\n", "\n", " ###### This is for restricted\n", " gender_related_columns = [\n", " 'Gender'\n", " ]\n", "\n", " age_related_columns = [\n", " 'Age_in_Yrs',\n", " 'Menstrual_AgeBegan',\n", " 'Menstrual_AgeIrreg',\n", " 'Menstrual_AgeStop',\n", " 'SSAGA_Alc_Age_1st_Use',\n", " 'SSAGA_TB_Age_1st_Cig',\n", " 'SSAGA_Mj_Age_1st_Use',\n", " 'Endurance_AgeAdj',\n", " 'Dexterity_AgeAdj',\n", " 'Strength_AgeAdj',\n", " 'PicSeq_AgeAdj',\n", " 'CardSort_AgeAdj',\n", " 'Flanker_AgeAdj',\n", " 'ReadEng_AgeAdj',\n", " 'PicVocab_AgeAdj',\n", " 'ProcSpeed_AgeAdj',\n", " 'Odor_AgeAdj',\n", " 'Taste_AgeAdj'\n", " ]\n", "\n", " # # show the first few rows of the subject information\n", " # subject_information_HCP[age_related_columns + sex_related_columns].head()\n", "\n", " # Handle missing values (e.g., impute with mean)\n", " mean_age = subject_information_HCP['Age_in_Yrs'].mean()\n", " \n", " # Initialize the scaler\n", " scaler = StandardScaler()\n", " \n", " # Perform z-score normalization\n", " subject_information_HCP['Age_in_Yrs_z'] = scaler.fit_transform(subject_information_HCP[['Age_in_Yrs']])\n", "\n", "\n", " \n", "def get_label_unrestricted(subject_id: List[str], target: str, method_for_age: str = 'mean') -> List:\n", " \"\"\"\n", " Get the label for the given subject id and target.\n", "\n", " For sex 0 is F and 1 is M\n", " \"\"\"\n", "\n", " # convert to list of ints\n", " subject_id = [int(x) for x in subject_id]\n", "\n", " if target == \"age\":\n", " age_array = []\n", " for subject in subject_id:\n", " c_age = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Age'].values\n", " # if the subject is not in the subject information file trigger an error\n", " if len(c_age) == 0:\n", " assert False, f\"Subject {subject} not found in subject information file\"\n", " if len(c_age) > 1:\n", " print(f\"Warning: Multiple entries for subject {subject}\")\n", "\n", " c_age = c_age[0].split('-')\n", " if len(c_age) < 2:\n", " c_age = c_age[0].split('+')\n", " age_array.append(int(c_age[0]))\n", " else:\n", " if method_for_age == 'mean':\n", " age_array.append(np.mean([int(x) for x in c_age]))\n", " elif method_for_age == 'min':\n", " age_array.append(np.min([int(x) for x in c_age]))\n", " elif method_for_age == 'max':\n", " age_array.append(np.max([int(x) for x in c_age]))\n", " else:\n", " assert False, f\"Method {method_for_age} not recognized\"\n", "\n", " return np.array(age_array) \n", " \n", " elif target == 'sex':\n", " sex_array = []\n", " for subject in subject_id:\n", " c_sex = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Gender'].values\n", " # if the subject is not in the subject information file trigger an error\n", " if len(c_sex) == 0:\n", " assert False, f\"Subject {subject} not found in subject information file\"\n", " if len(c_sex) > 1:\n", " print(f\"Warning: Multiple entries for subject {subject}\")\n", " sex_array.append(int(c_sex[0] == 'M'))\n", " return sex_array\n", "\n", "def get_label_restricted(subject_id: List[str], target: str, normalized: bool = True) -> List:\n", " \"\"\"\n", " Get the label for the given subject id and target.\n", "\n", " For sex 0 is F and 1 is M\n", " \"\"\"\n", "\n", " # convert to list of ints\n", " subject_id = [int(x) for x in subject_id]\n", "\n", " if target == \"age\":\n", " age_array = []\n", " for subject in subject_id:\n", " c_age = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Age_in_Yrs' if not normalized else 'Age_in_Yrs_z'].values\n", " # if the subject is not in the subject information file trigger an error\n", " if len(c_age) == 0:\n", " assert False, f\"Subject {subject} not found in subject information file\"\n", " if len(c_age) > 1:\n", " print(f\"Warning: Multiple entries for subject {subject}\")\n", "\n", " age_array.append(np.int8(c_age[0]))\n", "\n", " return np.array(age_array) \n", " \n", " elif target == 'sex':\n", " sex_array = []\n", " for subject in subject_id:\n", " c_sex = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Gender'].values\n", " # if the subject is not in the subject information file trigger an error\n", " if len(c_sex) == 0:\n", " assert False, f\"Subject {subject} not found in subject information file\"\n", " if len(c_sex) > 1:\n", " print(f\"Warning: Multiple entries for subject {subject}\")\n", " sex_array.append(int(c_sex[0] == 'M'))\n", " return sex_array" ] }, { "cell_type": "code", "execution_count": 18, "id": "40f04009-13f1-473d-9808-181a46d3e97e", "metadata": {}, "outputs": [], "source": [ "from sklearn.preprocessing import LabelEncoder\n", "\n", "if target == \"trial_type\":\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": 10, "id": "592b3cd3-6232-4cff-8d90-9034067b3ab9", "metadata": {}, "outputs": [], "source": [ "# for sample in tqdm(train_dl):\n", "# x = sample[0]\n", "# subject_id = sample[1]['sub']\n", "\n", "# # benchmark time\n", "# start = time.time()\n", "# y = get_label(subject_id, 'age')\n", "# end = time.time()\n", "# print(f\"Time taken: {end - start}\")\n", "# print(x.shape, y, subject_id, torch.Tensor(y).shape)\n", "# break" ] }, { "cell_type": "markdown", "id": "bb2d7e4f-eb34-498b-bbc6-127eb4a7bdc8", "metadata": {}, "source": [ "### Create pytorch model" ] }, { "cell_type": "code", "execution_count": 19, "id": "1a2c6527-fa75-42ba-be43-3d753ed8100c", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Input dimension: 737280\n" ] } ], "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[0][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": 20, "id": "b8983a0b-c810-4424-8c31-32b9921ebd9c", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "total_steps 34800\n" ] } ], "source": [ "# Initialize the model\n", "\n", "if target == \"trial_type\":\n", " model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n", " criterion = nn.CrossEntropyLoss()\n", "\n", "elif target == \"age\":\n", " model = LinearClassifier(input_dim=input_dim, num_classes=1)\n", " criterion = nn.MSELoss()\n", "\n", "elif target == \"sex\":\n", " model = LinearClassifier(input_dim=input_dim, num_classes=1)\n", " criterion = nn.BCEWithLogitsLoss()\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", "# import schedulefree\n", "# optimizer = schedulefree.AdamWScheduleFree(model.parameters(), lr=learning_rate, weight_decay=weight_decay)\n", "optimizer = torch.optim.AdamW(model.parameters(), lr=max_lr, weight_decay=weight_decay)\n", "\n", "num_iterations_per_epoch = math.ceil(flatmaps_train.shape[0]/batch_size)\n", "\n", "if lr_scheduler_type == 'linear':\n", " lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n", " optimizer,\n", " total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)),\n", " last_epoch=-1\n", " )\n", "elif lr_scheduler_type == 'cycle':\n", " total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))\n", " print(\"total_steps\", total_steps)\n", " lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(\n", " optimizer, \n", " max_lr=max_lr,\n", " total_steps=total_steps,\n", " final_div_factor=1000,\n", " last_epoch=-1, pct_start=2/num_epochs\n", " )\n" ] }, { "cell_type": "markdown", "id": "6b205b2c-730c-4d9e-89b7-c28031de78a7", "metadata": {}, "source": [ "### Wandb logging" ] }, { "cell_type": "code", "execution_count": 21, "id": "63ce8f26-a1b9-4e29-badc-7b9b30d0d3b3", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Running in interactive notebook. Disabling W&B and ckpt saving.\n" ] } ], "source": [ "import wandb\n", "import uuid\n", "\n", "myuuid = uuid.uuid4()\n", "str(myuuid)\n", "if utils.is_interactive():\n", " print(\"Running in interactive notebook. Disabling W&B and ckpt saving.\")\n", " wandb_log = False\n", " save_ckpt = False\n", "\n", "if wandb_log:\n", " wandb_project = 'fMRI-foundation-model'\n", " wandb_config = {\n", " \"model_name\": f\"HCPflat_raw_{target}\",\n", " \"batch_size\": batch_size,\n", " \"weight_decay\": weight_decay,\n", " \"num_epochs\": num_epochs,\n", " \"seed\": seed,\n", " \"lr_scheduler_type\": lr_scheduler_type,\n", " \"save_ckpt\": save_ckpt,\n", " \"seed\": seed,\n", " \"max_lr\": max_lr,\n", " \"target\": target,\n", " \"num_workers\": num_workers,\n", " \"weight_decay\": weight_decay\n", " }\n", " print(\"wandb_config:\\n\", wandb_config)\n", " random_id = random.randint(0, 100000)\n", " wandb_id = \"HCPflat_raw\" + f\"_{model_suffix}_{target}_{myuuid}\"\n", " print(\"wandb_id:\", wandb_id)\n", " wandb.init(\n", " id=wandb_id,\n", " project=wandb_project,\n", " name=\"HCPflat_raw\"+ f\"_{model_suffix}_{target}\",\n", " config=wandb_config,\n", " resume=\"allow\",\n", " )" ] }, { "cell_type": "markdown", "id": "b7d87e5c-94fc-4528-b2eb-65ef59c79634", "metadata": {}, "source": [ "### Training loop" ] }, { "cell_type": "code", "execution_count": null, "id": "f43535f4-0f6e-4655-b4f0-926416685358", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "f509e13d66a44d78bb06b0b0c4008187", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Epoch 1/40 - Training: 0%| | 0/870 [00:00 threshold).float().squeeze()\n", " correct_train += (predicted == labels).sum().item()\n", "\n", " total_train += labels.size(0)\n", " step += 1\n", "\n", " # Print intermediate metrics every 100 steps\n", " if step % 100 == 0:\n", " if target in [\"trial_type\", \"sex\"]:\n", " current_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0\n", " print(f\"Step [{step}/{len(train_dl)}] - Training Loss: {loss.item():.4f} - Training Accuracy: {current_accuracy:.2f}%\")\n", " elif target == \"age\":\n", " current_mse = mse_age_train / total_train if total_train > 0 else 0.0\n", " print(f\"Step [{step}/{len(train_dl)}] - Training Loss: {loss.item():.4f} - Training MSE: {current_mse:.4f}\")\n", "\n", " if lr_scheduler_type is not None:\n", " lr_scheduler.step()\n", "\n", " # Calculate epoch-level metrics\n", " epoch_train_loss = running_train_loss / total_train if total_train > 0 else 0.0\n", "\n", " if target in [\"trial_type\", \"sex\"]:\n", " train_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0\n", " elif target == \"age\":\n", " train_mse = mse_age_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", " mse_age_val = 0.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[0].to(device).float() # Removed unsqueeze(1) unless specifically needed\n", "\n", " # Prepare labels based on target type\n", " if target == \"trial_type\":\n", " labels = batch[1]['trial_type'] # List of labels\n", " labels = label_encoder.transform(labels)\n", " labels = torch.tensor(labels, dtype=torch.long).to(device) # Shape: [batch_size]\n", " elif target == \"age\":\n", " labels = get_label_restricted(batch[1]['sub'], 'age')\n", " labels = torch.tensor(labels, dtype=torch.float).to(device) # Shape: [batch_size]\n", " elif target == \"sex\":\n", " labels = get_label_restricted(batch[1]['sub'], 'sex')\n", " labels = torch.tensor(labels, dtype=torch.float).to(device) # Shape: [batch_size]\n", "\n", " # labels = labels.unsqueeze(1)\n", " \n", " # Forward pass\n", " outputs = model(images)\n", "\n", " # Compute loss\n", " if target in [\"trial_type\", \"sex\"]:\n", " loss = criterion(outputs.squeeze(), labels.squeeze())\n", " elif target == \"age\":\n", " loss = criterion(outputs.squeeze(), labels.squeeze())\n", "\n", " # Accumulate loss\n", " running_val_loss += loss.item() * images.size(0)\n", "\n", " # Calculate and accumulate metrics\n", " if target == \"trial_type\":\n", " _, predicted = torch.max(outputs, 1)\n", " correct_val += (predicted == labels).sum().item()\n", " elif target == \"age\":\n", " mse_age_val += (torch.sum((outputs.squeeze() - labels) ** 2).item()) / outputs.shape[0]\n", " elif target == \"sex\":\n", " threshold = 0.5\n", " predicted = (torch.sigmoid(outputs) > threshold).float().squeeze()\n", " correct_val += (predicted == labels).sum().item()\n", "\n", " total_val += labels.size(0)\n", "\n", " # Calculate epoch-level validation metrics\n", " epoch_val_loss = running_val_loss / total_val if total_val > 0 else 0.0\n", "\n", " if target in [\"trial_type\", \"sex\"]:\n", " val_accuracy = 100 * correct_val / total_val if total_val > 0 else 0.0\n", " elif target == \"age\":\n", " val_mse = mse_age_val / total_val if total_val > 0 else 0.0\n", "\n", " # Print epoch-level metrics\n", " if target in [\"trial_type\", \"sex\"]:\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", " elif target == \"age\":\n", " print(f\"Epoch [{epoch+1}/{num_epochs}] \"\n", " f\"- Training Loss: {epoch_train_loss:.4f}, Training MSE: {train_mse:.4f} \"\n", " f\"- Validation Loss: {epoch_val_loss:.4f}, Validation MSE: {val_mse:.4f}\")\n", "\n", " # Log metrics with wandb\n", " if wandb_log:\n", " log_dict = {\n", " \"epoch_train_loss\": epoch_train_loss,\n", " \"epoch_val_loss\": epoch_val_loss,\n", " }\n", " if target in [\"trial_type\", \"sex\"]:\n", " log_dict.update({\n", " f\"train_accuracy_{target}\": train_accuracy,\n", " f\"val_accuracy_{target}\": val_accuracy,\n", " })\n", " elif target == \"age\":\n", " log_dict.update({\n", " f\"train_mse_{target}\": train_mse,\n", " f\"val_mse_{target}\": val_mse,\n", " })\n", " wandb.log(log_dict)\n", "\n", "# Save checkpoint if required\n", "if save_ckpt:\n", " outdir = os.path.abspath(f'checkpoints/{\"HCPflat_raw\"+ f\"_{model_suffix}_{target}\"}_{random_id}')\n", " os.makedirs(outdir, exist_ok=True)\n", " print(\"Saving checkpoint to:\", outdir)\n", " # Save model state\n", " torch.save(model.state_dict(), os.path.join(outdir, \"model.pth\"))\n", " # Save configuration\n", " with open(os.path.join(outdir, \"config.yaml\"), 'w') as f:\n", " yaml.dump(wandb_config, f)\n", " print(f\"Model and config saved to {outdir}\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "cf28a3bd-5376-4d73-ab2f-d87e64677b1b", "metadata": {}, "outputs": [], "source": [ "# loss = criterion(outputs, labels)" ] }, { "cell_type": "code", "execution_count": null, "id": "dc7cf575-84aa-4b11-a704-482390442dbf", "metadata": {}, "outputs": [], "source": [ "# outputs.shape, labels.squeeze().shape" ] }, { "cell_type": "code", "execution_count": null, "id": "8b97d80a-521a-4968-906a-d542d559aec8", "metadata": {}, "outputs": [], "source": [ "# loss = criterion(outputs, labels)" ] }, { "cell_type": "code", "execution_count": null, "id": "3181f512-c7ea-4aa0-b963-1cb629162060", "metadata": {}, "outputs": [], "source": [ "# if target == 'trial_type':\n", "# key = 'trial_type'\n", "# elif target == 'sex' or target == 'age':\n", "# key = 'sub'\n", "\n", "# y_train = [json.loads(metadata_train[i])[key] for i in range(0,2000)]\n", "# y_val = [json.loads(metadata_train[i])[key] for i in range(10000,11000)]\n", "# y_test = [json.loads(metadata_test[i])[key] for i in range(0,1000)]" ] }, { "cell_type": "code", "execution_count": null, "id": "14a96458-3946-4c3b-bc44-cf7614d04311", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "cc5b4da6-d06a-4cf0-beff-e41c00bea53e", "metadata": {}, "outputs": [], "source": [ "# X_train = flatmaps_train[0:2000]\n", "# X_val = flatmaps_train[10000:11000]\n", "# X_test = flatmaps_test[0:1000]\n", "\n", "# y_test = get_label_restricted(y_test, target = 'sex')\n", "# y_train = get_label_restricted(y_train, target = 'sex')\n", "# y_val = get_label_restricted(y_val, target = 'sex')\n", "\n", "# # y_train = label_encoder.transform(y_train)\n", "# # y_val = label_encoder.transform(y_val)\n", "# # y_test = label_encoder.transform(y_test)" ] }, { "cell_type": "code", "execution_count": null, "id": "fce01070-918b-4ee4-8ae6-ab0437f4f7ef", "metadata": {}, "outputs": [], "source": [ "# X_train, X_val, X_test = X_train.reshape(X_train.shape[0],-1), X_val.reshape(X_val.shape[0],-1), X_test.reshape(X_test.shape[0],-1)" ] }, { "cell_type": "code", "execution_count": null, "id": "f2405a1f-4537-4c93-b945-c6d6e8be3ea9", "metadata": {}, "outputs": [], "source": [ "# X_train.shape" ] }, { "cell_type": "code", "execution_count": null, "id": "ea9d0a30-2e06-4db0-a66a-02c944f17809", "metadata": {}, "outputs": [], "source": [ "# import numpy as np\n", "# import matplotlib.pyplot as plt\n", "# from sklearn.preprocessing import StandardScaler\n", "# from sklearn.decomposition import PCA\n", "# from sklearn.linear_model import LogisticRegressionCV\n", "# from sklearn.metrics import accuracy_score\n", "\n", "# # Supongamos que ya tienes tus datos divididos:\n", "# # X_train, y_train, X_val, y_val, X_test, y_test\n", "\n", "# # 1. Estandarizar los Datos\n", "# print(\"Estandarizando los datos...\")\n", "# scaler = StandardScaler()\n", "# X_train_scaled = scaler.fit_transform(X_train)\n", "# X_val_scaled = scaler.transform(X_val)\n", "# X_test_scaled = scaler.transform(X_test)\n", "\n", "# # 2. Aplicar PCA\n", "# print(\"Aplicando PCA...\")\n", "# # Decidir el número de componentes. Por ejemplo, mantener el 95% de la varianza.\n", "# pca = PCA(n_components=0.95, svd_solver='full') # 'full' para compatibilidad\n", "# X_train_pca = pca.fit_transform(X_train_scaled)\n", "# X_val_pca = pca.transform(X_val_scaled)\n", "# X_test_pca = pca.transform(X_test_scaled)\n", "\n", "# print(f\"Número de componentes seleccionados: {pca.n_components_}\")\n", "\n", "# # Opcional: Visualizar la varianza explicada\n", "# cumulative_variance = np.cumsum(pca.explained_variance_ratio_)\n", "# plt.figure(figsize=(8, 5))\n", "# plt.plot(range(1, len(cumulative_variance) + 1), cumulative_variance, marker='o', linestyle='--')\n", "# plt.xlabel('Número de Componentes')\n", "# plt.ylabel('Varianza Acumulada')\n", "# plt.title('Varianza Explicada por PCA')\n", "# plt.grid(True)\n", "# plt.show()\n", "\n", "# # 3. Entrenar el Modelo de Regresión Logística con Validación Cruzada\n", "# print(\"Entrenando el modelo de Regresión Logística con PCA...\")\n", "# clf = LogisticRegressionCV(max_iter=100, cv=5, scoring='accuracy', n_jobs=-1)\n", "# clf.fit(X_train_pca, y_train)\n", "\n", "# # 4. Evaluar el Modelo\n", "# print(\"Calculando precisión...\")\n", "\n", "# # Precisión en entrenamiento\n", "# y_train_pred = clf.predict(X_train_pca)\n", "# train_acc = accuracy_score(y_train, y_train_pred)\n", "\n", "# # Precisión en validación\n", "# y_val_pred = clf.predict(X_val_pca)\n", "# val_acc = accuracy_score(y_val, y_val_pred)\n", "\n", "# # Precisión en prueba\n", "# y_test_pred = clf.predict(X_test_pca)\n", "# test_acc = accuracy_score(y_test, y_test_pred)\n", "\n", "# print(f\"Precisión en entrenamiento: {train_acc:.4f}\")\n", "# print(f\"Precisión en validación: {val_acc:.4f}\")\n", "# print(f\"Precisión en prueba: {test_acc:.4f}\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5d588872-a0c7-490d-9c6f-5d14449351ac", "metadata": {}, "outputs": [], "source": [ "# X_train_scaled.shape" ] }, { "cell_type": "code", "execution_count": null, "id": "5fa286de-adc3-4813-93a3-4e9dcaf52763", "metadata": {}, "outputs": [], "source": [ "# from sklearn.linear_model import LogisticRegressionCV, Ridge\n", "# print(\"fitting\")\n", "# clf = LogisticRegressionCV(max_iter=100)\n", "# clf.fit(X_train, y_train)\n", "# print(\"calculating accuracy\")\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)" ] }, { "cell_type": "code", "execution_count": null, "id": "c60691a4-fc51-4d99-9755-e3a82ce52162", "metadata": {}, "outputs": [], "source": [ "# print(train_acc, val_acc, test_acc)" ] }, { "cell_type": "code", "execution_count": null, "id": "abb8ed7c-219d-4433-9375-750f8f98d5d0", "metadata": {}, "outputs": [], "source": [ "### SEX\n", "# Sklearn No pca just 1k examples: 1.0 0.534 0.5066666666666667\n", "# Sklearn Pca 1800 features, 2k examples 1.0000 0.5130 0.4590\n", "# All data pytorch 0.93 no_val 0.55\n", "\n", "\n", "### TRIAL TYPE\n", "# Sklearn No pca just 1k examples: 1.0 0.61 0.63\n", "# Sklearn Pca 500 features, 2k examples 1.0000 ~0.73 ~0.73\n", "# All data pytorch 0.9911 no_val 0.94" ] }, { "cell_type": "code", "execution_count": null, "id": "88987f84-472c-4864-b90d-b15fcab8b095", "metadata": {}, "outputs": [], "source": [ "# a = model.linear.weight[0][10:20]\n", "# a" ] }, { "cell_type": "code", "execution_count": null, "id": "43ce1770-a25c-4986-89cb-5d8fad3fc50a", "metadata": {}, "outputs": [], "source": [ "# loss = criterion(outputs, labels.unsqueeze(1))\n", "# loss" ] } ], "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 }