Add files using upload-large-folder tool

.gitattributes

@@ -5035,3 +5035,9 @@ fMRI-foundation-model/src/wandb/run-20241024_030627-HCPflat_large_gsrFalse__HCP_
 fMRI-foundation-model/src/wandb/run-20241024_160745-HCPflat_large_gsrFalse__HCP_FT_f1d2455a-d8d7-4d9d-9f91-58748aebb427/run-HCPflat_large_gsrFalse__HCP_FT_f1d2455a-d8d7-4d9d-9f91-58748aebb427.wandb filter=lfs diff=lfs merge=lfs -text
 fMRI-foundation-model/src/wandb/run-20241127_125303-HCPflat_raw_beta_trial_type_a0e1e642-966f-441b-9bc7-974dce26cba1/run-HCPflat_raw_beta_trial_type_a0e1e642-966f-441b-9bc7-974dce26cba1.wandb filter=lfs diff=lfs merge=lfs -text
 fMRI-foundation-model/src/wandb/run-20241127_023946-NSDflat_large_gsrFalse__beta_age_HCPFT_260e5584-8c2b-4e13-a5ed-11dcdc2a522f/run-NSDflat_large_gsrFalse__beta_age_HCPFT_260e5584-8c2b-4e13-a5ed-11dcdc2a522f.wandb filter=lfs diff=lfs merge=lfs -text
+fMRI-foundation-model/src/wandb/run-20241126_143710-HCPflat_large_gsrFalse__HCP_FT_79bf330c-a53f-43d5-86dd-b4bb676b9b78/run-HCPflat_large_gsrFalse__HCP_FT_79bf330c-a53f-43d5-86dd-b4bb676b9b78.wandb filter=lfs diff=lfs merge=lfs -text
+fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/run-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532.wandb filter=lfs diff=lfs merge=lfs -text
+fMRI-foundation-model/src/wandb/run-20241127_013141-HCPflat_raw_beta_age_18a4fb68-2904-438d-a472-1c5e5f991d72/run-HCPflat_raw_beta_age_18a4fb68-2904-438d-a472-1c5e5f991d72.wandb filter=lfs diff=lfs merge=lfs -text
+fMRI-foundation-model/src/wandb/run-20241024_213834-HCPflat_large_gsrFalse__HCP_FT_d07870e8-3c53-420c-b734-f32a1c5c3e5a/run-HCPflat_large_gsrFalse__HCP_FT_d07870e8-3c53-420c-b734-f32a1c5c3e5a.wandb filter=lfs diff=lfs merge=lfs -text
+fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/run-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427.wandb filter=lfs diff=lfs merge=lfs -text
+fMRI-foundation-model/src/wandb/run-20241023_041223-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3/run-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3.wandb filter=lfs diff=lfs merge=lfs -text

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/files/code/_session_history.ipynb

@@ -0,0 +1,1664 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "406d87ac",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "import h5py\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",
+    "# batch_size = 128\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "utils.seed_everything(seed)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "3f6365eb",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "import h5py\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",
+    "# batch_size = 128\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "utils.seed_everything(seed)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "b96ed0fa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "2344601f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "33cf89e4",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "import h5py\n",
+    "from mae_utils import flat_models\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",
+    "# batch_size = 128\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "utils.seed_everything(seed)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "bc2281a4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "acfaeaad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    latest_checkpoint = \"epoch99.pth\"\n",
+    "else:\n",
+    "    latest_checkpoint = sys.argv[2] \n",
+    "print(f\"latest_checkpoint: {latest_checkpoint}\")\n",
+    "\n",
+    "# Load the checkpoint\n",
+    "checkpoint_path = os.path.join(outdir, latest_checkpoint)\n",
+    "\n",
+    "state = torch.load(checkpoint_path)\n",
+    "model.load_state_dict(state[\"model_state_dict\"], strict=False)\n",
+    "model.to(device)\n",
+    "model.eval()\n",
+    "\n",
+    "print(f\"\\nLoaded checkpoint {latest_checkpoint} from {outdir}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "8ffbe1b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')\n",
+    "flatmaps_train = f_train['flatmaps']\n",
+    "\n",
+    "f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.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.npy', allow_pickle=True)\n",
+    "metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "767e8c90",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "mae_model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "fa59d7d5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    latest_checkpoint = \"epoch99.pth\"\n",
+    "else:\n",
+    "    latest_checkpoint = sys.argv[2] \n",
+    "print(f\"latest_checkpoint: {latest_checkpoint}\")\n",
+    "\n",
+    "# Load the checkpoint\n",
+    "checkpoint_path = os.path.join(outdir, latest_checkpoint)\n",
+    "\n",
+    "state = torch.load(checkpoint_path)\n",
+    "mae_model.load_state_dict(state[\"model_state_dict\"], strict=False)\n",
+    "mae_model.to(device)\n",
+    "\n",
+    "print(f\"\\nLoaded checkpoint {latest_checkpoint} from {outdir}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "cea19e70",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')\n",
+    "flatmaps_train = f_train['flatmaps']\n",
+    "\n",
+    "f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.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.npy', allow_pickle=True)\n",
+    "metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "f1f85737",
+   "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=num_workers)\n",
+    "\n",
+    "test_dataset = HCPFlatDataset(flatmaps_test, metadata_test)\n",
+    "test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "f80c3176",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for i in train_dl:\n",
+    "    break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "eb364065",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for i in test_dl:\n",
+    "    break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "d7581eea",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         ...,\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]]], dtype=torch.float16),\n",
+      " {'key': ['sub-102109_mod-tfMRI_task-WM_mag-3T_dir-LR',\n",
+      "   'sub-731140_mod-tfMRI_task-WM_mag-3T_dir-RL',\n",
+      "   'sub-149539_mod-tfMRI_task-MOTOR_mag-3T_dir-RL',\n",
+      "   'sub-376247_mod-tfMRI_task-WM_mag-3T_dir-RL',\n",
+      "   'sub-164939_mod-tfMRI_task-EMOTION_mag-3T_dir-LR',\n",
+      "   'sub-198653_mod-tfMRI_task-SOCIAL_mag-3T_dir-RL',\n",
+      "   'sub-210011_mod-tfMRI_task-WM_mag-3T_dir-RL',\n",
+      "   'sub-356948_mod-tfMRI_task-EMOTION_mag-3T_dir-LR'],\n",
+      "  'sub': ['102109',\n",
+      "   '731140',\n",
+      "   '149539',\n",
+      "   '376247',\n",
+      "   '164939',\n",
+      "   '198653',\n",
+      "   '210011',\n",
+      "   '356948'],\n",
+      "  'mod': ['tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI'],\n",
+      "  'task': ['WM', 'WM', 'MOTOR', 'WM', 'EMOTION', 'SOCIAL', 'WM', 'EMOTION'],\n",
+      "  'mag': ['3T', '3T', '3T', '3T', '3T', '3T', '3T', '3T'],\n",
+      "  'dir': ['LR', 'RL', 'RL', 'RL', 'LR', 'RL', 'RL', 'LR'],\n",
+      "  'start': tensor([15, 15, 19, 15, 19, 15, 15, 19]),\n",
+      "  'trial_type': ['2bk_tools',\n",
+      "   '2bk_body',\n",
+      "   'lh',\n",
+      "   '2bk_body',\n",
+      "   'neut',\n",
+      "   'mental',\n",
+      "   '2bk_body',\n",
+      "   'neut']}]"
+     ]
+    }
+   ],
+   "source": [
+    "i"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "ae8a1da7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 16, 144, 320])"
+     ]
+    }
+   ],
+   "source": [
+    "i[0].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "ce3f32de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mae_model(torch.randn(1,16,144,320)).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "ff571604",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "global_pool"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "348c52bc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if os.getenv('global_pool') == \"False\":\n",
+    "    global_pool = False\n",
+    "else:\n",
+    "    global_pool = True\n",
+    "print(f\"global_pool = {global_pool}\")\n",
+    "\n",
+    "try:\n",
+    "    gsr\n",
+    "except:\n",
+    "    gsr = True\n",
+    "    print(\"set gsr to True\")\n",
+    "print(f\"gsr = {gsr}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "82ccddb7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mae_model(torch.randn(1,16,144,320),global_pool=global_pool, forward_features = True).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "856a0186",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mae_model(torch.randn(1,1,16,144,320),global_pool=global_pool, forward_features = True).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "33e1c8cd",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 1024])"
+     ]
+    }
+   ],
+   "source": [
+    "mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "c56d3fc0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 1024])"
+     ]
+    }
+   ],
+   "source": [
+    "mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "d6395f12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:].sum()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "2dca71b2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "list(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:]).sum()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "a25dd030",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1024"
+     ]
+    }
+   ],
+   "source": [
+    "sum(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "edd2edd2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "np.int64(1024)"
+     ]
+    }
+   ],
+   "source": [
+    "np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "7c934b0f",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])\n",
+    "print(f\"Input dimension: {input_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "bd0fcfa7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FullModel(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        self.mae_model = mae_model\n",
+    "        \n",
+    "        \n",
+    "    def forward(self, x, gsr):\n",
+    "        x = self.mae_model(x, global_pool=global_pool, forward_features = True)\n",
+    "        x = self.lc_model(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "25b06ed9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LinearClassifier(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        \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",
+    "input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])\n",
+    "print(f\"Input dimension: {input_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "97f9bbc3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the model\n",
+    "lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n",
+    "\n",
+    "model = FullModel(lc_model, mae_model)\n",
+    "\n",
+    "# Move the model to the GPU\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",
+    "num_epochs = 20  # Adjust as needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "313b529e",
+   "metadata": {},
+   "outputs": [],
+   "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": 33,
+   "id": "c5f58f30",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')\n",
+    "flatmaps_train = f_train['flatmaps']\n",
+    "\n",
+    "f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.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.npy', allow_pickle=True)\n",
+    "metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "id": "928abadb",
+   "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=num_workers)\n",
+    "\n",
+    "test_dataset = HCPFlatDataset(flatmaps_test, metadata_test)\n",
+    "test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "ec69248a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "mae_model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "4e5045c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    latest_checkpoint = \"epoch99.pth\"\n",
+    "else:\n",
+    "    latest_checkpoint = sys.argv[2] \n",
+    "print(f\"latest_checkpoint: {latest_checkpoint}\")\n",
+    "\n",
+    "# Load the checkpoint\n",
+    "checkpoint_path = os.path.join(outdir, latest_checkpoint)\n",
+    "\n",
+    "state = torch.load(checkpoint_path)\n",
+    "mae_model.load_state_dict(state[\"model_state_dict\"], strict=False)\n",
+    "mae_model.to(device)\n",
+    "\n",
+    "print(f\"\\nLoaded checkpoint {latest_checkpoint} from {outdir}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "0173f847",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FullModel(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        self.mae_model = mae_model\n",
+    "        \n",
+    "        \n",
+    "    def forward(self, x, gsr):\n",
+    "        x = self.mae_model(x, global_pool=global_pool, forward_features = True)\n",
+    "        x = self.lc_model(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 38,
+   "id": "551a5976",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LinearClassifier(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        \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",
+    "input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])\n",
+    "print(f\"Input dimension: {input_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "id": "a21df922",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the model\n",
+    "lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n",
+    "\n",
+    "model = FullModel(lc_model, mae_model)\n",
+    "\n",
+    "# Move the model to the GPU\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",
+    "num_epochs = 20  # Adjust as needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 40,
+   "id": "7b262408",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])\n",
+    "print(f\"Input dimension: {input_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "id": "7bff8f73",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FullModel(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        self.mae_model = mae_model\n",
+    "        \n",
+    "        \n",
+    "    def forward(self, x, gsr):\n",
+    "        x = self.mae_model(x, global_pool=global_pool, forward_features = True)\n",
+    "        x = self.lc_model(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "id": "697fc651",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the model\n",
+    "lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n",
+    "\n",
+    "model = FullModel(lc_model, mae_model)\n",
+    "\n",
+    "# Move the model to the GPU\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",
+    "num_epochs = 20  # Adjust as needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 43,
+   "id": "99570156",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FullModel(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(FullModel, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        self.mae_model = mae_model\n",
+    "        \n",
+    "        \n",
+    "    def forward(self, x, gsr):\n",
+    "        x = self.mae_model(x, global_pool=global_pool, forward_features = True)\n",
+    "        x = self.lc_model(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 44,
+   "id": "8e5c0e3e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the model\n",
+    "lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n",
+    "\n",
+    "model = FullModel(lc_model, mae_model)\n",
+    "\n",
+    "# Move the model to the GPU\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",
+    "num_epochs = 20  # Adjust as needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 45,
+   "id": "dddffb31",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.18.3"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://stability.wandb.io/ckadirt/fMRI-foundation-model/runs/HCPflat_raw_83810' target=\"_blank\">HCPflat_raw</a></strong> to <a href='https://stability.wandb.io/ckadirt/fMRI-foundation-model' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://stability.wandb.io/ckadirt/fMRI-foundation-model' target=\"_blank\">https://stability.wandb.io/ckadirt/fMRI-foundation-model</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://stability.wandb.io/ckadirt/fMRI-foundation-model/runs/HCPflat_raw_83810' target=\"_blank\">https://stability.wandb.io/ckadirt/fMRI-foundation-model/runs/HCPflat_raw_83810</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import wandb\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    print(\"Running in interactive notebook. Disabling W&B and ckpt saving.\")\n",
+    "    wandb_log = True #False\n",
+    "    save_ckpt = True #False\n",
+    "\n",
+    "if wandb_log:\n",
+    "    wandb_project = 'fMRI-foundation-model'\n",
+    "    wandb_config = {\n",
+    "        \"model_name\": \"HCPflat_raw\",\n",
+    "        \"batch_size\": batch_size,\n",
+    "        \"learning_rate\": learning_rate,\n",
+    "        \"weight_decay\": weight_decay,\n",
+    "        \"num_epochs\": num_epochs,\n",
+    "        \"seed\": seed,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\", wandb_config)\n",
+    "    random_id = random.randint(0, 100000)\n",
+    "    print(\"wandb_id:\", \"HCPflat_raw\" + f\"_{random_id}\")\n",
+    "    wandb.init(\n",
+    "        id=\"HCPflat_raw\" + f\"_{random_id}\",\n",
+    "        project=wandb_project,\n",
+    "        name=\"HCPflat_raw\",\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 46,
+   "id": "f8a61d67",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import wandb\n",
+    "\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\": model_name+'_HCP_FT',\n",
+    "        \"batch_size\": batch_size,\n",
+    "        \"learning_rate\": learning_rate,\n",
+    "        \"weight_decay\": weight_decay,\n",
+    "        \"num_epochs\": num_epochs,\n",
+    "        \"seed\": seed,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\", wandb_config)\n",
+    "    random_id = random.randint(0, 100000)\n",
+    "    print(\"wandb_id:\", \"HCPflat_raw\" + f\"_{random_id}\")\n",
+    "    wandb.init(\n",
+    "        id=model_name+'_HCP_FT' + f\"_{random_id}\",\n",
+    "        project=wandb_project,\n",
+    "        name=model_name+'_HCP_FT',\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
+   "id": "71ac3c4f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import wandb\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    print(\"Running in interactive notebook. Disabling W&B and ckpt saving.\")\n",
+    "    wandb_log = True\n",
+    "    save_ckpt = True\n",
+    "\n",
+    "if wandb_log:\n",
+    "    wandb_project = 'fMRI-foundation-model'\n",
+    "    wandb_config = {\n",
+    "        \"model_name\": model_name+'_HCP_FT',\n",
+    "        \"batch_size\": batch_size,\n",
+    "        \"learning_rate\": learning_rate,\n",
+    "        \"weight_decay\": weight_decay,\n",
+    "        \"num_epochs\": num_epochs,\n",
+    "        \"seed\": seed,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\", wandb_config)\n",
+    "    random_id = random.randint(0, 100000)\n",
+    "    print(\"wandb_id:\", \"HCPflat_raw\" + f\"_{random_id}\")\n",
+    "    wandb.init(\n",
+    "        id=model_name+'_HCP_FT' + f\"_{random_id}\",\n",
+    "        project=wandb_project,\n",
+    "        name=model_name+'_HCP_FT',\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "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
+}

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/files/config.yaml

@@ -0,0 +1,49 @@
+_wandb:
+    value:
+        cli_version: 0.18.3
+        m: []
+        python_version: 3.11.10
+        session_history: code/_session_history.ipynb
+        t:
+            "1":
+                - 1
+                - 5
+                - 41
+                - 49
+                - 53
+                - 55
+                - 63
+            "2":
+                - 1
+                - 5
+                - 41
+                - 49
+                - 53
+                - 55
+                - 63
+            "3":
+                - 2
+                - 13
+                - 14
+                - 16
+                - 23
+                - 55
+            "4": 3.11.10
+            "5": 0.18.3
+            "8":
+                - 1
+                - 5
+            "12": 0.18.3
+            "13": linux-x86_64
+batch_size:
+    value: 8
+learning_rate:
+    value: 0.0001
+model_name:
+    value: HCPflat_raw
+num_epochs:
+    value: 20
+seed:
+    value: 42
+weight_decay:
+    value: 1e-05

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/files/output.log

@@ -0,0 +1,5 @@
+Running in interactive notebook. Disabling W&B and ckpt saving.
+Running in interactive notebook. Disabling W&B and ckpt saving.
+wandb_config:
+ {'model_name': 'HCPflat_large_gsrFalse__HCP_FT', 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 1e-05, 'num_epochs': 20, 'seed': 42}
+wandb_id: HCPflat_raw_14592

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/files/wandb-metadata.json

@@ -0,0 +1,144 @@
+{
+  "os":  "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+  "python":  "3.11.10",
+  "startedAt":  "2024-10-23T03:22:15.095152Z",
+  "program":  "ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.ipynb",
+  "git":  {
+    "remote":  "https://github.com/MedARC-AI/fMRI-foundation-model",
+    "commit":  "b1ba684ae7a5cc4155cc046b0abe613de09bf700"
+  },
+  "email":  "torrico.villanueva.cesar.kadir@gmail.com",
+  "root":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+  "host":  "ip-10-0-160-143",
+  "username":  "ckadirt",
+  "executable":  "/admin/home-ckadirt/foundation_env/bin/python",
+  "cpu_count":  96,
+  "cpu_count_logical":  192,
+  "gpu":  "[NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3]",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "249555763200",
+      "used":  "184990625792"
+    }
+  },
+  "memory":  {
+    "total":  "2147443429376"
+  },
+  "cpu":  {
+    "count":  96,
+    "countLogical":  192
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    }
+  ],
+  "slurm":  {
+    "cluster_name":  "sagemaker2",
+    "conf":  "/opt/slurm/etc/slurm.conf",
+    "cpu_bind":  "quiet,mask_cpu:0x00000000000000FFC000000000000000000000FFC0000000",
+    "cpu_bind_list":  "0x00000000000000FFC000000000000000000000FFC0000000",
+    "cpu_bind_type":  "mask_cpu:",
+    "cpu_bind_verbose":  "quiet",
+    "cpus_on_node":  "20",
+    "gpus":  "1",
+    "gpus_on_node":  "1",
+    "gtids":  "0",
+    "job_account":  "fmri",
+    "job_cpus_per_node":  "20",
+    "job_end_time":  "1729702756",
+    "job_gid":  "1879800513",
+    "job_group":  "Domain Users",
+    "job_id":  "528040",
+    "job_name":  "bash",
+    "job_nodelist":  "ip-10-0-160-143",
+    "job_num_nodes":  "1",
+    "job_partition":  "p5",
+    "job_qos":  "idle",
+    "job_start_time":  "1729648756",
+    "job_uid":  "1879804696",
+    "job_user":  "ckadirt",
+    "jobid":  "528040",
+    "launch_node_ipaddr":  "172.17.12.61",
+    "localid":  "0",
+    "mpi_type":  "pmix_v3",
+    "nnodes":  "1",
+    "nodeid":  "0",
+    "nodelist":  "ip-10-0-160-143",
+    "nprocs":  "1",
+    "ntasks":  "1",
+    "pmix_mapping_serv":  "(vector,(0,1,1))",
+    "pmixp_abort_agent_port":  "34923",
+    "prio_process":  "0",
+    "procid":  "0",
+    "pty_port":  "45733",
+    "pty_win_col":  "199",
+    "pty_win_row":  "17",
+    "script_context":  "prolog_task",
+    "srun_comm_host":  "172.17.12.61",
+    "srun_comm_port":  "39353",
+    "step_gpus":  "3",
+    "step_id":  "0",
+    "step_launcher_port":  "39353",
+    "step_nodelist":  "ip-10-0-160-143",
+    "step_num_nodes":  "1",
+    "step_num_tasks":  "1",
+    "step_tasks_per_node":  "1",
+    "stepid":  "0",
+    "submit_dir":  "/weka/proj-fmri",
+    "submit_host":  "ip-172-17-12-61",
+    "task_pid":  "1032669",
+    "tasks_per_node":  "1",
+    "topology_addr":  "ip-10-0-160-143",
+    "topology_addr_pattern":  "node",
+    "umask":  "0022",
+    "working_cluster":  "sagemaker2:ip-172-17-63-161:6817:9984:109"
+  },
+  "cudaVersion":  "12.2"
+}

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/files/wandb-summary.json

@@ -0,0 +1 @@
+{"_wandb":{"runtime":1}}

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/logs/debug-core.log

@@ -0,0 +1,12 @@
+{"time":"2024-10-23T03:22:14.402317432Z","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmps2dvy9x1/port-1034100.txt","pid":1034100,"debug":false,"disable-analytics":false}
+{"time":"2024-10-23T03:22:14.402858522Z","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false}
+{"time":"2024-10-23T03:22:14.407755671Z","level":"INFO","msg":"Will exit if parent process dies.","ppid":1034100}
+{"time":"2024-10-23T03:22:14.40773793Z","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":41495,"Zone":""}}
+{"time":"2024-10-23T03:22:14.557294663Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:58122"}
+{"time":"2024-10-23T03:22:15.098249959Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_raw_83810","id":"127.0.0.1:58122"}
+{"time":"2024-10-23T03:22:15.159214287Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_raw_83810","id":"127.0.0.1:58122"}
+{"time":"2024-10-23T03:24:02.210685377Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"HCPflat_raw_83810","id":"127.0.0.1:58122"}
+{"time":"2024-10-23T03:24:02.210981773Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"HCPflat_raw_83810","id":"127.0.0.1:58122"}
+{"time":"2024-10-23T03:24:02.366411684Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_large_gsrFalse__HCP_FT_14592","id":"127.0.0.1:58122"}
+{"time":"2024-10-23T03:24:02.542190524Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_large_gsrFalse__HCP_FT_14592","id":"127.0.0.1:58122"}
+{"time":"2024-10-23T03:36:47.830235978Z","level":"INFO","msg":"Parent process exited, terminating service process."}

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/logs/debug-internal.log

@@ -0,0 +1,25 @@
+{"time":"2024-10-23T03:22:15.107200519Z","level":"INFO","msg":"using version","core version":"0.18.3"}
+{"time":"2024-10-23T03:22:15.107222919Z","level":"INFO","msg":"created symlink","path":"/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/logs/debug-core.log"}
+{"time":"2024-10-23T03:22:15.12212287Z","level":"ERROR","msg":"dialing: google: could not find default credentials. See https://cloud.google.com/docs/authentication/external/set-up-adc for more information"}
+{"time":"2024-10-23T03:22:15.159184506Z","level":"INFO","msg":"created new stream","id":"HCPflat_raw_83810"}
+{"time":"2024-10-23T03:22:15.159208527Z","level":"INFO","msg":"stream: started","id":"HCPflat_raw_83810"}
+{"time":"2024-10-23T03:22:15.159221567Z","level":"INFO","msg":"writer: Do: started","stream_id":{"value":"HCPflat_raw_83810"}}
+{"time":"2024-10-23T03:22:15.159246878Z","level":"INFO","msg":"sender: started","stream_id":{"value":"HCPflat_raw_83810"}}
+{"time":"2024-10-23T03:22:15.159265338Z","level":"INFO","msg":"handler: started","stream_id":{"value":"HCPflat_raw_83810"}}
+{"time":"2024-10-23T03:22:15.639713975Z","level":"INFO","msg":"wandb-core","!BADKEY":null}
+{"time":"2024-10-23T03:22:15.65038308Z","level":"INFO","msg":"Starting system monitor"}
+{"time":"2024-10-23T03:22:15.650435701Z","level":"WARN","msg":"handleCodeSave: program relative path is empty"}
+{"time":"2024-10-23T03:22:15.652277478Z","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}
+{"time":"2024-10-23T03:22:16.318300603Z","level":"INFO","msg":"Pausing system monitor"}
+{"time":"2024-10-23T03:23:43.467389909Z","level":"INFO","msg":"Resuming system monitor"}
+{"time":"2024-10-23T03:23:43.474183646Z","level":"INFO","msg":"Pausing system monitor"}
+{"time":"2024-10-23T03:23:54.344053787Z","level":"INFO","msg":"Resuming system monitor"}
+{"time":"2024-10-23T03:23:55.025154766Z","level":"INFO","msg":"Stopping system monitor"}
+{"time":"2024-10-23T03:23:55.038944374Z","level":"INFO","msg":"Stopped system monitor"}
+{"time":"2024-10-23T03:24:00.028953872Z","level":"ERROR","msg":"monitor: gpu: timeout waiting for process to exit"}
+{"time":"2024-10-23T03:24:00.893610579Z","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2024-10-23T03:24:02.210783329Z","level":"INFO","msg":"stream: closing","id":"HCPflat_raw_83810"}
+{"time":"2024-10-23T03:24:02.21081877Z","level":"INFO","msg":"handler: closed","stream_id":{"value":"HCPflat_raw_83810"}}
+{"time":"2024-10-23T03:24:02.21083897Z","level":"INFO","msg":"writer: Close: closed","stream_id":{"value":"HCPflat_raw_83810"}}
+{"time":"2024-10-23T03:24:02.21085142Z","level":"INFO","msg":"sender: closed","stream_id":{"value":"HCPflat_raw_83810"}}
+{"time":"2024-10-23T03:24:02.210972013Z","level":"INFO","msg":"stream: closed","id":"HCPflat_raw_83810"}

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/logs/debug.log

@@ -0,0 +1,59 @@
+2024-10-23 03:22:15,080 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Current SDK version is 0.18.3
+2024-10-23 03:22:15,081 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Configure stats pid to 1034100
+2024-10-23 03:22:15,081 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-10-23 03:22:15,081 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Loading settings from /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/settings
+2024-10-23 03:22:15,081 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Loading settings from environment variables: {}
+2024-10-23 03:22:15,081 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Applying setup settings: {'mode': None, '_disable_service': None}
+2024-10-23 03:22:15,081 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Inferring run settings from compute environment: {'program': '<python with no main file>'}
+2024-10-23 03:22:15,081 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Applying login settings: {}
+2024-10-23 03:22:15,081 INFO    MainThread:1034100 [wandb_init.py:_log_setup():532] Logging user logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/logs/debug.log
+2024-10-23 03:22:15,082 INFO    MainThread:1034100 [wandb_init.py:_log_setup():533] Logging internal logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/logs/debug-internal.log
+2024-10-23 03:22:15,082 INFO    MainThread:1034100 [wandb_init.py:_jupyter_setup():478] configuring jupyter hooks <wandb.sdk.wandb_init._WandbInit object at 0x7f0957b13950>
+2024-10-23 03:22:15,083 INFO    MainThread:1034100 [wandb_init.py:init():617] calling init triggers
+2024-10-23 03:22:15,083 INFO    MainThread:1034100 [wandb_init.py:init():624] wandb.init called with sweep_config: {}
+config: {'model_name': 'HCPflat_raw', 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 1e-05, 'num_epochs': 20, 'seed': 42}
+2024-10-23 03:22:15,083 INFO    MainThread:1034100 [wandb_init.py:init():667] starting backend
+2024-10-23 03:22:15,083 INFO    MainThread:1034100 [wandb_init.py:init():671] sending inform_init request
+2024-10-23 03:22:15,093 INFO    MainThread:1034100 [backend.py:_multiprocessing_setup():104] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-10-23 03:22:15,094 INFO    MainThread:1034100 [wandb_init.py:init():684] backend started and connected
+2024-10-23 03:22:15,118 INFO    MainThread:1034100 [wandb_run.py:_label_probe_notebook():1346] probe notebook
+2024-10-23 03:22:15,124 INFO    MainThread:1034100 [wandb_run.py:_label_probe_notebook():1356] Unable to probe notebook: 'NoneType' object has no attribute 'get'
+2024-10-23 03:22:15,124 INFO    MainThread:1034100 [wandb_init.py:init():779] updated telemetry
+2024-10-23 03:22:15,219 INFO    MainThread:1034100 [wandb_init.py:init():812] communicating run to backend with 90.0 second timeout
+2024-10-23 03:22:15,624 INFO    MainThread:1034100 [wandb_init.py:init():863] starting run threads in backend
+2024-10-23 03:22:16,273 INFO    MainThread:1034100 [wandb_run.py:_console_start():2465] atexit reg
+2024-10-23 03:22:16,273 INFO    MainThread:1034100 [wandb_run.py:_redirect():2313] redirect: wrap_raw
+2024-10-23 03:22:16,273 INFO    MainThread:1034100 [wandb_run.py:_redirect():2378] Wrapping output streams.
+2024-10-23 03:22:16,273 INFO    MainThread:1034100 [wandb_run.py:_redirect():2403] Redirects installed.
+2024-10-23 03:22:16,289 INFO    MainThread:1034100 [wandb_init.py:init():907] run started, returning control to user process
+2024-10-23 03:22:16,295 INFO    MainThread:1034100 [jupyter.py:_save_ipynb():398] looking for notebook: ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.ipynb
+2024-10-23 03:22:16,296 INFO    MainThread:1034100 [wandb_init.py:_pause_backend():443] pausing backend
+2024-10-23 03:23:43,467 INFO    MainThread:1034100 [wandb_init.py:_resume_backend():448] resuming backend
+2024-10-23 03:23:43,469 INFO    MainThread:1034100 [jupyter.py:_save_ipynb():398] looking for notebook: ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.ipynb
+2024-10-23 03:23:43,474 INFO    MainThread:1034100 [wandb_init.py:_pause_backend():443] pausing backend
+2024-10-23 03:23:54,343 INFO    MainThread:1034100 [wandb_init.py:_resume_backend():448] resuming backend
+2024-10-23 03:23:54,855 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Current SDK version is 0.18.3
+2024-10-23 03:23:54,855 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Configure stats pid to 1034100
+2024-10-23 03:23:54,856 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-10-23 03:23:54,856 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Loading settings from /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/settings
+2024-10-23 03:23:54,856 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Loading settings from environment variables: {}
+2024-10-23 03:23:54,856 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Applying setup settings: {'mode': None, '_disable_service': None}
+2024-10-23 03:23:54,856 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Inferring run settings from compute environment: {'program': '<python with no main file>'}
+2024-10-23 03:23:54,856 INFO    MainThread:1034100 [wandb_setup.py:_flush():79] Applying login settings: {}
+2024-10-23 03:23:54,860 INFO    MainThread:1034100 [wandb_init.py:_log_setup():532] Logging user logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_032354-HCPflat_large_gsrFalse__HCP_FT_14592/logs/debug.log
+2024-10-23 03:23:54,865 INFO    MainThread:1034100 [wandb_init.py:_log_setup():533] Logging internal logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_032354-HCPflat_large_gsrFalse__HCP_FT_14592/logs/debug-internal.log
+2024-10-23 03:23:54,865 INFO    MainThread:1034100 [wandb_init.py:init():617] calling init triggers
+2024-10-23 03:23:54,865 INFO    MainThread:1034100 [wandb_init.py:init():624] wandb.init called with sweep_config: {}
+config: {'model_name': 'HCPflat_large_gsrFalse__HCP_FT', 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 1e-05, 'num_epochs': 20, 'seed': 42}
+2024-10-23 03:23:54,865 INFO    MainThread:1034100 [wandb_init.py:init():642] re-initializing run, found existing run on stack: HCPflat_raw_83810
+2024-10-23 03:23:54,871 INFO    MainThread:1034100 [wandb_run.py:_finish():2164] finishing run ckadirt/fMRI-foundation-model/HCPflat_raw_83810
+2024-10-23 03:23:54,912 INFO    MainThread:1034100 [jupyter.py:save_history():488] saving 47 cells to _session_history.ipynb
+2024-10-23 03:23:54,913 INFO    MainThread:1034100 [wandb_run.py:_config_callback():1394] config_cb ('_wandb', 'session_history') code/_session_history.ipynb None
+2024-10-23 03:23:55,013 INFO    MainThread:1034100 [jupyter.py:_save_ipynb():398] looking for notebook: ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.ipynb
+2024-10-23 03:23:55,017 INFO    MainThread:1034100 [wandb_init.py:_jupyter_teardown():460] cleaning up jupyter logic
+2024-10-23 03:23:55,017 INFO    MainThread:1034100 [wandb_run.py:_atexit_cleanup():2428] got exitcode: 0
+2024-10-23 03:23:55,024 INFO    MainThread:1034100 [wandb_run.py:_restore():2410] restore
+2024-10-23 03:23:55,024 INFO    MainThread:1034100 [wandb_run.py:_restore():2416] restore done
+2024-10-23 03:24:02,162 INFO    MainThread:1034100 [wandb_run.py:_footer_history_summary_info():4049] rendering history
+2024-10-23 03:24:02,162 INFO    MainThread:1034100 [wandb_run.py:_footer_history_summary_info():4081] rendering summary
+2024-10-23 03:24:02,204 INFO    MainThread:1034100 [wandb_run.py:_footer_sync_info():4008] logging synced files

fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810/tmp/code/_session_history.ipynb

@@ -0,0 +1,1664 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "406d87ac",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "import h5py\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",
+    "# batch_size = 128\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "utils.seed_everything(seed)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "3f6365eb",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "import h5py\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",
+    "# batch_size = 128\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "utils.seed_everything(seed)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "b96ed0fa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "2344601f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "33cf89e4",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "import h5py\n",
+    "from mae_utils import flat_models\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",
+    "# batch_size = 128\n",
+    "\n",
+    "print(\"PID of this process =\",os.getpid())\n",
+    "utils.seed_everything(seed)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "bc2281a4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "acfaeaad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    latest_checkpoint = \"epoch99.pth\"\n",
+    "else:\n",
+    "    latest_checkpoint = sys.argv[2] \n",
+    "print(f\"latest_checkpoint: {latest_checkpoint}\")\n",
+    "\n",
+    "# Load the checkpoint\n",
+    "checkpoint_path = os.path.join(outdir, latest_checkpoint)\n",
+    "\n",
+    "state = torch.load(checkpoint_path)\n",
+    "model.load_state_dict(state[\"model_state_dict\"], strict=False)\n",
+    "model.to(device)\n",
+    "model.eval()\n",
+    "\n",
+    "print(f\"\\nLoaded checkpoint {latest_checkpoint} from {outdir}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "8ffbe1b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')\n",
+    "flatmaps_train = f_train['flatmaps']\n",
+    "\n",
+    "f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.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.npy', allow_pickle=True)\n",
+    "metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "767e8c90",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "mae_model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "fa59d7d5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    latest_checkpoint = \"epoch99.pth\"\n",
+    "else:\n",
+    "    latest_checkpoint = sys.argv[2] \n",
+    "print(f\"latest_checkpoint: {latest_checkpoint}\")\n",
+    "\n",
+    "# Load the checkpoint\n",
+    "checkpoint_path = os.path.join(outdir, latest_checkpoint)\n",
+    "\n",
+    "state = torch.load(checkpoint_path)\n",
+    "mae_model.load_state_dict(state[\"model_state_dict\"], strict=False)\n",
+    "mae_model.to(device)\n",
+    "\n",
+    "print(f\"\\nLoaded checkpoint {latest_checkpoint} from {outdir}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "cea19e70",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')\n",
+    "flatmaps_train = f_train['flatmaps']\n",
+    "\n",
+    "f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.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.npy', allow_pickle=True)\n",
+    "metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "f1f85737",
+   "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=num_workers)\n",
+    "\n",
+    "test_dataset = HCPFlatDataset(flatmaps_test, metadata_test)\n",
+    "test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "f80c3176",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for i in train_dl:\n",
+    "    break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "eb364065",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for i in test_dl:\n",
+    "    break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "d7581eea",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         ...,\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
+      " \n",
+      " \n",
+      "         [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          ...,\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]],\n",
+      " \n",
+      "          [[0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           ...,\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.],\n",
+      "           [0., 0., 0.,  ..., 0., 0., 0.]]]], dtype=torch.float16),\n",
+      " {'key': ['sub-102109_mod-tfMRI_task-WM_mag-3T_dir-LR',\n",
+      "   'sub-731140_mod-tfMRI_task-WM_mag-3T_dir-RL',\n",
+      "   'sub-149539_mod-tfMRI_task-MOTOR_mag-3T_dir-RL',\n",
+      "   'sub-376247_mod-tfMRI_task-WM_mag-3T_dir-RL',\n",
+      "   'sub-164939_mod-tfMRI_task-EMOTION_mag-3T_dir-LR',\n",
+      "   'sub-198653_mod-tfMRI_task-SOCIAL_mag-3T_dir-RL',\n",
+      "   'sub-210011_mod-tfMRI_task-WM_mag-3T_dir-RL',\n",
+      "   'sub-356948_mod-tfMRI_task-EMOTION_mag-3T_dir-LR'],\n",
+      "  'sub': ['102109',\n",
+      "   '731140',\n",
+      "   '149539',\n",
+      "   '376247',\n",
+      "   '164939',\n",
+      "   '198653',\n",
+      "   '210011',\n",
+      "   '356948'],\n",
+      "  'mod': ['tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI',\n",
+      "   'tfMRI'],\n",
+      "  'task': ['WM', 'WM', 'MOTOR', 'WM', 'EMOTION', 'SOCIAL', 'WM', 'EMOTION'],\n",
+      "  'mag': ['3T', '3T', '3T', '3T', '3T', '3T', '3T', '3T'],\n",
+      "  'dir': ['LR', 'RL', 'RL', 'RL', 'LR', 'RL', 'RL', 'LR'],\n",
+      "  'start': tensor([15, 15, 19, 15, 19, 15, 15, 19]),\n",
+      "  'trial_type': ['2bk_tools',\n",
+      "   '2bk_body',\n",
+      "   'lh',\n",
+      "   '2bk_body',\n",
+      "   'neut',\n",
+      "   'mental',\n",
+      "   '2bk_body',\n",
+      "   'neut']}]"
+     ]
+    }
+   ],
+   "source": [
+    "i"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "ae8a1da7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 16, 144, 320])"
+     ]
+    }
+   ],
+   "source": [
+    "i[0].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "ce3f32de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mae_model(torch.randn(1,16,144,320)).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "ff571604",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "global_pool"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "348c52bc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if os.getenv('global_pool') == \"False\":\n",
+    "    global_pool = False\n",
+    "else:\n",
+    "    global_pool = True\n",
+    "print(f\"global_pool = {global_pool}\")\n",
+    "\n",
+    "try:\n",
+    "    gsr\n",
+    "except:\n",
+    "    gsr = True\n",
+    "    print(\"set gsr to True\")\n",
+    "print(f\"gsr = {gsr}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "82ccddb7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mae_model(torch.randn(1,16,144,320),global_pool=global_pool, forward_features = True).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "856a0186",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mae_model(torch.randn(1,1,16,144,320),global_pool=global_pool, forward_features = True).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "33e1c8cd",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 1024])"
+     ]
+    }
+   ],
+   "source": [
+    "mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "c56d3fc0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 1024])"
+     ]
+    }
+   ],
+   "source": [
+    "mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "d6395f12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:].sum()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "2dca71b2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "list(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:]).sum()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "a25dd030",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1024"
+     ]
+    }
+   ],
+   "source": [
+    "sum(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "edd2edd2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "np.int64(1024)"
+     ]
+    }
+   ],
+   "source": [
+    "np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "7c934b0f",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])\n",
+    "print(f\"Input dimension: {input_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "bd0fcfa7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FullModel(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        self.mae_model = mae_model\n",
+    "        \n",
+    "        \n",
+    "    def forward(self, x, gsr):\n",
+    "        x = self.mae_model(x, global_pool=global_pool, forward_features = True)\n",
+    "        x = self.lc_model(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "25b06ed9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LinearClassifier(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        \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",
+    "input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])\n",
+    "print(f\"Input dimension: {input_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "97f9bbc3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the model\n",
+    "lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n",
+    "\n",
+    "model = FullModel(lc_model, mae_model)\n",
+    "\n",
+    "# Move the model to the GPU\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",
+    "num_epochs = 20  # Adjust as needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "313b529e",
+   "metadata": {},
+   "outputs": [],
+   "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": 33,
+   "id": "c5f58f30",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')\n",
+    "flatmaps_train = f_train['flatmaps']\n",
+    "\n",
+    "f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.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.npy', allow_pickle=True)\n",
+    "metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "id": "928abadb",
+   "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=num_workers)\n",
+    "\n",
+    "test_dataset = HCPFlatDataset(flatmaps_test, metadata_test)\n",
+    "test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "ec69248a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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",
+    "flat_mask = load_hcp_flat_mask(hcp_flat_path)\n",
+    "\n",
+    "mae_model = flat_models.mae_vit_large_fmri(\n",
+    "    patch_size=patch_size,\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",
+    "    pct_masks_to_decode=pct_masks_to_decode,\n",
+    "    img_mask=flat_mask,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "4e5045c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    latest_checkpoint = \"epoch99.pth\"\n",
+    "else:\n",
+    "    latest_checkpoint = sys.argv[2] \n",
+    "print(f\"latest_checkpoint: {latest_checkpoint}\")\n",
+    "\n",
+    "# Load the checkpoint\n",
+    "checkpoint_path = os.path.join(outdir, latest_checkpoint)\n",
+    "\n",
+    "state = torch.load(checkpoint_path)\n",
+    "mae_model.load_state_dict(state[\"model_state_dict\"], strict=False)\n",
+    "mae_model.to(device)\n",
+    "\n",
+    "print(f\"\\nLoaded checkpoint {latest_checkpoint} from {outdir}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "0173f847",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FullModel(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        self.mae_model = mae_model\n",
+    "        \n",
+    "        \n",
+    "    def forward(self, x, gsr):\n",
+    "        x = self.mae_model(x, global_pool=global_pool, forward_features = True)\n",
+    "        x = self.lc_model(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 38,
+   "id": "551a5976",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LinearClassifier(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        \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",
+    "input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])\n",
+    "print(f\"Input dimension: {input_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "id": "a21df922",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the model\n",
+    "lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n",
+    "\n",
+    "model = FullModel(lc_model, mae_model)\n",
+    "\n",
+    "# Move the model to the GPU\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",
+    "num_epochs = 20  # Adjust as needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 40,
+   "id": "7b262408",
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])\n",
+    "print(f\"Input dimension: {input_dim}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "id": "7bff8f73",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FullModel(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(LinearClassifier, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        self.mae_model = mae_model\n",
+    "        \n",
+    "        \n",
+    "    def forward(self, x, gsr):\n",
+    "        x = self.mae_model(x, global_pool=global_pool, forward_features = True)\n",
+    "        x = self.lc_model(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "id": "697fc651",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the model\n",
+    "lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n",
+    "\n",
+    "model = FullModel(lc_model, mae_model)\n",
+    "\n",
+    "# Move the model to the GPU\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",
+    "num_epochs = 20  # Adjust as needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 43,
+   "id": "99570156",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FullModel(nn.Module):\n",
+    "    def __init__(self, lc_model, mae_model):\n",
+    "        super(FullModel, self).__init__()\n",
+    "        self.lc_model = lc_model\n",
+    "        self.mae_model = mae_model\n",
+    "        \n",
+    "        \n",
+    "    def forward(self, x, gsr):\n",
+    "        x = self.mae_model(x, global_pool=global_pool, forward_features = True)\n",
+    "        x = self.lc_model(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 44,
+   "id": "8e5c0e3e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize the model\n",
+    "lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)\n",
+    "\n",
+    "model = FullModel(lc_model, mae_model)\n",
+    "\n",
+    "# Move the model to the GPU\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",
+    "num_epochs = 20  # Adjust as needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 45,
+   "id": "dddffb31",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.18.3"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_032214-HCPflat_raw_83810</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://stability.wandb.io/ckadirt/fMRI-foundation-model/runs/HCPflat_raw_83810' target=\"_blank\">HCPflat_raw</a></strong> to <a href='https://stability.wandb.io/ckadirt/fMRI-foundation-model' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://stability.wandb.io/ckadirt/fMRI-foundation-model' target=\"_blank\">https://stability.wandb.io/ckadirt/fMRI-foundation-model</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://stability.wandb.io/ckadirt/fMRI-foundation-model/runs/HCPflat_raw_83810' target=\"_blank\">https://stability.wandb.io/ckadirt/fMRI-foundation-model/runs/HCPflat_raw_83810</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import wandb\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    print(\"Running in interactive notebook. Disabling W&B and ckpt saving.\")\n",
+    "    wandb_log = True #False\n",
+    "    save_ckpt = True #False\n",
+    "\n",
+    "if wandb_log:\n",
+    "    wandb_project = 'fMRI-foundation-model'\n",
+    "    wandb_config = {\n",
+    "        \"model_name\": \"HCPflat_raw\",\n",
+    "        \"batch_size\": batch_size,\n",
+    "        \"learning_rate\": learning_rate,\n",
+    "        \"weight_decay\": weight_decay,\n",
+    "        \"num_epochs\": num_epochs,\n",
+    "        \"seed\": seed,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\", wandb_config)\n",
+    "    random_id = random.randint(0, 100000)\n",
+    "    print(\"wandb_id:\", \"HCPflat_raw\" + f\"_{random_id}\")\n",
+    "    wandb.init(\n",
+    "        id=\"HCPflat_raw\" + f\"_{random_id}\",\n",
+    "        project=wandb_project,\n",
+    "        name=\"HCPflat_raw\",\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 46,
+   "id": "f8a61d67",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import wandb\n",
+    "\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\": model_name+'_HCP_FT',\n",
+    "        \"batch_size\": batch_size,\n",
+    "        \"learning_rate\": learning_rate,\n",
+    "        \"weight_decay\": weight_decay,\n",
+    "        \"num_epochs\": num_epochs,\n",
+    "        \"seed\": seed,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\", wandb_config)\n",
+    "    random_id = random.randint(0, 100000)\n",
+    "    print(\"wandb_id:\", \"HCPflat_raw\" + f\"_{random_id}\")\n",
+    "    wandb.init(\n",
+    "        id=model_name+'_HCP_FT' + f\"_{random_id}\",\n",
+    "        project=wandb_project,\n",
+    "        name=model_name+'_HCP_FT',\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
+   "id": "71ac3c4f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import wandb\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    print(\"Running in interactive notebook. Disabling W&B and ckpt saving.\")\n",
+    "    wandb_log = True\n",
+    "    save_ckpt = True\n",
+    "\n",
+    "if wandb_log:\n",
+    "    wandb_project = 'fMRI-foundation-model'\n",
+    "    wandb_config = {\n",
+    "        \"model_name\": model_name+'_HCP_FT',\n",
+    "        \"batch_size\": batch_size,\n",
+    "        \"learning_rate\": learning_rate,\n",
+    "        \"weight_decay\": weight_decay,\n",
+    "        \"num_epochs\": num_epochs,\n",
+    "        \"seed\": seed,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\", wandb_config)\n",
+    "    random_id = random.randint(0, 100000)\n",
+    "    print(\"wandb_id:\", \"HCPflat_raw\" + f\"_{random_id}\")\n",
+    "    wandb.init(\n",
+    "        id=model_name+'_HCP_FT' + f\"_{random_id}\",\n",
+    "        project=wandb_project,\n",
+    "        name=model_name+'_HCP_FT',\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "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
+}

fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/files/code/src/HCP_downstream_finetune.py

@@ -0,0 +1,587 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# Import packages and setup gpu configuration.
+# This code block shouldnt need to be adjusted!
+import os
+import sys
+import json
+import yaml
+import numpy as np
+import copy
+import math
+import time
+import random
+from tqdm.auto import tqdm
+import webdataset as wds
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import utils
+from mae_utils.flat_models import *
+import h5py
+from mae_utils import flat_models
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+# following fixes a Conv3D CUDNN_NOT_SUPPORTED error
+torch.backends.cudnn.benchmark = True
+
+# ## MODEL TO LOAD ##
+if utils.is_interactive():
+    model_name = "HCPflat_large_gsrFalse_"
+else:
+    model_name = sys.argv[1]
+    
+
+# outdir = os.path.abspath(f'checkpoints/{model_name}')
+outdir = os.path.abspath(f'checkpoints/{model_name}')
+
+print("outdir", outdir)
+# Load previous config.yaml if available
+if os.path.exists(f"{outdir}/config.yaml"):
+    config = yaml.load(open(f"{outdir}/config.yaml", 'r'), Loader=yaml.FullLoader)
+    print(f"Loaded config.yaml from ckpt folder {outdir}")
+    # create global variables from the config
+    print("\n__CONFIG__")
+    for attribute_name in config.keys():
+        print(f"{attribute_name} = {config[attribute_name]}")
+        globals()[attribute_name] = config[f'{attribute_name}']
+    print("\n")
+
+world_size = os.getenv('WORLD_SIZE')
+if world_size is None: 
+    world_size = 1
+else:
+    world_size = int(world_size)
+print(f"WORLD_SIZE={world_size}")
+
+if utils.is_interactive():
+    # Following allows you to change functions in models.py or utils.py and 
+    # have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    get_ipython().run_line_magic('autoreload', '2')
+
+batch_size = probe_batch_size
+num_epochs = probe_num_epochs
+
+data_type = torch.float32 # change depending on your mixed_precision
+global_batch_size = batch_size * world_size
+
+device = torch.device('cuda')
+
+hcp_flat_path = "/weka/proj-medarc/shared/HCP-Flat"
+# seed = 42
+# num_frames = 16
+# gsr = False
+# num_workers = 10
+# batch_size = 128
+
+print("PID of this process =",os.getpid())
+utils.seed_everything(seed)
+
+
+# In[2]:
+
+
+if os.getenv('global_pool') == "False":
+    global_pool = False
+else:
+    global_pool = True
+print(f"global_pool = {global_pool}")
+
+try:
+    gsr
+except:
+    gsr = True
+    print("set gsr to True")
+print(f"gsr = {gsr}")
+
+
+# In[3]:
+
+
+#### UNCOMMENT THIS TO SAVE THE HCP-FLAT IN HDF5 FORMAT
+
+
+# from torch.utils.data import default_collate
+# from mae_utils.flat import load_hcp_flat_mask
+# from mae_utils.flat import create_hcp_flat
+# from mae_utils.flat import batch_unmask
+# import mae_utils.visualize as vis
+
+
+# batch_size = 26
+# print(f"changed batch_size to {batch_size}")
+
+# ## Test ##
+# datasets_to_include = "HCP"
+# assert "HCP" in datasets_to_include
+# test_dataset = create_hcp_flat(root=hcp_flat_path, 
+#                 clip_mode="event", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'test')
+# test_dl = wds.WebLoader(
+#     test_dataset.batched(batch_size, partial=False, collation_fn=default_collate),
+#     batch_size=None,
+#     shuffle=False,
+#     num_workers=num_workers,
+#     pin_memory=True,
+# )
+
+# ## Train ##
+# assert "HCP" in datasets_to_include
+# train_dataset = create_hcp_flat(root=hcp_flat_path, 
+#                 clip_mode="event", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'train')
+# train_dl = wds.WebLoader(
+#     train_dataset.batched(batch_size, partial=False, collation_fn=default_collate),
+#     batch_size=None,
+#     shuffle=False,
+#     num_workers=num_workers,
+#     pin_memory=True,
+# )
+
+# def flatten_meta(meta_dict):
+#     """
+#     Flatten the meta dictionary by:
+#     - Replacing single-item lists with the item itself.
+#     - Converting tensors to scalar numbers.
+#     """
+#     flattened = {}
+#     for key, value in meta_dict.items():
+#         if isinstance(value, list):
+#             if len(value) == 1:
+#                 flattened[key] = value[0]  # Replace list with its single item
+#             else:
+#                 flattened[key] = value  # Keep as is if multiple items
+#         elif isinstance(value, torch.Tensor):
+#             # Convert tensor to scalar
+#             if value.numel() == 1:
+#                 flattened[key] = value.item()
+#             else:
+#                 flattened[key] = value.tolist()  # Convert multi-element tensor to list
+#         else:
+#             flattened[key] = value  # Keep the value as is
+#     return flattened
+
+# import h5py
+# meta_array = np.array([], dtype=object)
+# # Open an HDF5 file in write mode
+# with h5py.File('train_hcp.hdf5', 'w') as h5f:
+#     flatmaps_dset = None
+    
+#     total_samples = 0
+
+#     for i, batch in tqdm(enumerate(train_dl), total = 120000):
+#         images = batch['image'][0]
+#         meta = batch['meta']
+#         batch_size = images.shape[0]
+#         meta_serializable = meta.copy()
+        
+        
+#         # Step 2: Serialize the dictionary to a JSON string
+#         meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)
+#         meta_array = np.append(meta_array, meta_str)
+#         if flatmaps_dset is None:
+#             # Initialize datasets with unlimited (None) maxshape along the first axis
+#             flatmaps_shape = (0,) + images.shape[1:]
+#             flatmaps_maxshape = (None,) + images.shape[1:]
+
+#             flatmaps_dset = h5f.create_dataset(
+#                 'flatmaps',
+#                 shape=flatmaps_shape,
+#                 maxshape=flatmaps_maxshape,
+#                 dtype=np.float16,
+#                 chunks=True  # Enable chunking for efficient resizing
+#             )
+
+#         # Resize datasets to accommodate new data
+#         flatmaps_dset.resize(total_samples + batch_size, axis=0)
+
+#         # Write data to the datasets
+#         flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)
+
+#         total_samples += batch_size
+        
+#     print(f"Processed {total_samples} samples")
+# np.save('metadata_test_HCP.npy', meta_array)
+
+
+# import h5py
+# meta_array = np.array([], dtype=object)
+# # Open an HDF5 file in write mode
+# with h5py.File('test_hcp.hdf5', 'w') as h5f:
+#     flatmaps_dset = None
+    
+#     total_samples = 0
+
+#     for i, batch in tqdm(enumerate(test_dl), total = 12000):
+#         images = batch['image'][0]
+#         meta = batch['meta']
+#         batch_size = images.shape[0]
+#         meta_serializable = meta.copy()
+        
+        
+#         # Step 2: Serialize the dictionary to a JSON string
+#         meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)
+#         meta_array = np.append(meta_array, meta_str)
+#         if flatmaps_dset is None:
+#             # Initialize datasets with unlimited (None) maxshape along the first axis
+#             flatmaps_shape = (0,) + images.shape[1:]
+#             flatmaps_maxshape = (None,) + images.shape[1:]
+
+#             flatmaps_dset = h5f.create_dataset(
+#                 'flatmaps',
+#                 shape=flatmaps_shape,
+#                 maxshape=flatmaps_maxshape,
+#                 dtype=np.float16,
+#                 chunks=True  # Enable chunking for efficient resizing
+#             )
+
+#         # Resize datasets to accommodate new data
+#         flatmaps_dset.resize(total_samples + batch_size, axis=0)
+
+#         # Write data to the datasets
+#         flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)
+
+#         total_samples += batch_size
+        
+#     print(f"Processed {total_samples} samples")
+# np.save('metadata_train_HCP.npy', meta_array)
+
+
+# ### Preparing data
+
+# In[4]:
+
+
+from sklearn.preprocessing import LabelEncoder
+
+INCLUDE_CONDS = {
+    "fear",
+    "neut",
+    "math",
+    "story",
+    "lf",
+    "lh",
+    "rf",
+    "rh",
+    "t",
+    "match",
+    "relation",
+    "mental",
+    "rnd",
+    "0bk_body",
+    "2bk_body",
+    "0bk_faces",
+    "2bk_faces",
+    "0bk_places",
+    "2bk_places",
+    "0bk_tools",
+    "2bk_tools",
+}
+
+# test_data = []
+
+# # Iterate over the DataLoader with a progress bar
+# for sample in tqdm(train_dl, desc="Processing samples"):
+#     x = sample['image']
+#     y = sample['meta']['trial_type']
+#     key = sample['meta']['key']
+#     print(x.shape, y, key)
+#     break
+# Initialize the label encoder
+label_encoder = LabelEncoder()
+label_encoder.fit(sorted(INCLUDE_CONDS))  # Ensure consistent ordering
+
+num_classes = len(label_encoder.classes_)
+print(f"Number of classes: {num_classes}")
+
+
+# In[5]:
+
+
+f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')
+flatmaps_train = f_train['flatmaps']
+
+f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.hdf5', 'r')
+flatmaps_test = f_test['flatmaps']
+
+metadata_train = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_train_HCP.npy', allow_pickle=True)
+metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)
+
+
+# In[6]:
+
+
+from torch.utils.data import Dataset, DataLoader
+
+class HCPFlatDataset(Dataset):
+    def __init__(self, flatmaps, metadata):
+        self.flatmaps = flatmaps
+        self.metadata = metadata
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, idx):
+        return self.flatmaps[idx], json.loads(self.metadata[idx])
+print("Moving datasets to ram")
+# Loading to cpu for faster training, this can take several minutes.  Remove this [:] if you want to move one at the time.
+train_dataset = HCPFlatDataset(flatmaps_train, metadata_train)
+train_dl = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+
+test_dataset = HCPFlatDataset(flatmaps_test, metadata_test)
+test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+print("Datasets ready")
+
+
+# ### Creating and loading Model
+
+# In[7]:
+
+
+from mae_utils.flat import load_hcp_flat_mask
+from mae_utils.flat import create_hcp_flat
+from mae_utils.flat import batch_unmask
+import mae_utils.visualize as vis
+
+flat_mask = load_hcp_flat_mask(hcp_flat_path)
+
+mae_model = flat_models.mae_vit_large_fmri(
+    patch_size=patch_size,
+    decoder_embed_dim=decoder_embed_dim,
+    t_patch_size=t_patch_size,
+    pred_t_dim=pred_t_dim,
+    decoder_depth=4,
+    cls_embed=cls_embed,
+    norm_pix_loss=norm_pix_loss,
+    no_qkv_bias=no_qkv_bias,
+    sep_pos_embed=sep_pos_embed,
+    trunc_init=trunc_init,
+    pct_masks_to_decode=pct_masks_to_decode,
+    img_mask=flat_mask,
+)
+
+
+# In[8]:
+
+
+checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]
+
+if utils.is_interactive():
+    latest_checkpoint = "epoch99.pth"
+else:
+    latest_checkpoint = sys.argv[2] 
+print(f"latest_checkpoint: {latest_checkpoint}")
+
+# Load the checkpoint
+checkpoint_path = os.path.join(outdir, latest_checkpoint)
+
+state = torch.load(checkpoint_path)
+mae_model.load_state_dict(state["model_state_dict"], strict=False)
+mae_model.to(device)
+
+print(f"\nLoaded checkpoint {latest_checkpoint} from {outdir}\n")
+
+
+# In[9]:
+
+
+class LinearClassifier(nn.Module):
+    def __init__(self, input_dim, num_classes):
+        super(LinearClassifier, self).__init__()
+        self.linear = nn.Linear(input_dim, num_classes)
+        
+    def forward(self, x):
+        # Flatten the input except for the batch dimension
+        x = x.view(x.size(0), -1)
+        out = self.linear(x)
+        return out  # Raw logits
+
+# Determine the input dimension from a single sample
+# Assuming images are of shape [1, 16, 144, 320]
+input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])
+print(f"Input dimension: {input_dim}")
+
+
+# In[10]:
+
+
+class FullModel(nn.Module):
+    def __init__(self, lc_model, mae_model):
+        super(FullModel, self).__init__()
+        self.lc_model = lc_model
+        self.mae_model = mae_model
+        
+        
+    def forward(self, x, gsr):
+        x = self.mae_model(x, global_pool=global_pool, forward_features = True)
+        x = self.lc_model(x)
+        return x
+
+
+# In[11]:
+
+
+# Initialize the model
+lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)
+
+model = FullModel(lc_model, mae_model)
+
+# Move the model to the GPU
+model.to(device)
+
+# Define loss function
+criterion = nn.CrossEntropyLoss()
+
+# Define optimizer with L2 regularization (weight_decay)
+learning_rate = 1e-4
+weight_decay = 1e-5  # Adjust based on your needs
+optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+num_epochs = 20  # Adjust as needed
+
+
+# ### Data
+
+# In[12]:
+
+
+import wandb
+
+if utils.is_interactive():
+    print("Running in interactive notebook. Disabling W&B and ckpt saving.")
+    wandb_log = True
+    save_ckpt = True
+
+if wandb_log:
+    wandb_project = 'fMRI-foundation-model'
+    wandb_config = {
+        "model_name": model_name+'_HCP_FT',
+        "batch_size": batch_size,
+        "learning_rate": learning_rate,
+        "weight_decay": weight_decay,
+        "num_epochs": num_epochs,
+        "seed": seed,
+    }
+    print("wandb_config:\n", wandb_config)
+    random_id = random.randint(0, 100000)
+    print("wandb_id:", "HCPflat_raw" + f"_{random_id}")
+    wandb.init(
+        id=model_name+'_HCP_FT' + f"_{random_id}",
+        project=wandb_project,
+        name=model_name+'_HCP_FT',
+        config=wandb_config,
+        resume="allow",
+    )
+
+
+# In[13]:
+
+
+for epoch in range(num_epochs):
+    running_train_loss = 0.0
+    correct_train = 0
+    total_train = 0
+    step = 0
+
+    # with torch.amp.autocast(device_type='cuda'):
+    # Training Phase
+    model.train()
+    for  batch in tqdm(train_dl, desc=f"Epoch {epoch+1}/{num_epochs} - Training"):
+        optimizer.zero_grad()
+        images = batch[0].to(device).float().unsqueeze(1) #fix this  # Shape: [batch_size, 1, 16, 144, 320]
+        labels = batch[1]['trial_type']  # List of labels
+        
+        encoded_labels = label_encoder.transform(labels)
+        encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)  # Shape: [batch_size]
+        
+        # Forward pass
+        outputs = model(images, gsr=gsr)  # Shape: [num_train_samples, num_classes]
+        
+        # Compute loss
+        loss = criterion(outputs, encoded_labels)
+        
+        # Backward pass and optimization
+        loss.backward()
+        optimizer.step()
+        
+        # Accumulate loss
+        running_train_loss += loss.item() * images.size(0)
+
+        
+        # Calculate accuracy
+        _, predicted = torch.max(outputs, 1)
+        
+        correct_train += (predicted == encoded_labels).sum().item()
+        total_train += encoded_labels.size(0)
+        
+        step = step + 1
+        if step % 100 == 0:
+            print(f"Step [{step}/{len(train_dl)}] - Training Loss: {loss.item():.4f} - Training Accuracy: {100 * correct_train / total_train:.2f}%")
+        # thth
+
+    epoch_train_loss = running_train_loss / total_train if total_train > 0 else 0.0
+    train_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0
+    
+    # Validation Phase
+    model.eval()
+    running_val_loss = 0.0
+    correct_val = 0
+    total_val = 0
+    
+    with torch.no_grad():
+        for batch in tqdm(test_dl, desc=f"Epoch {epoch+1}/{num_epochs} - Validation"):
+            
+            images = batch[0].to(device).float().unsqueeze(1) #fix this
+            labels = batch[1]['trial_type']
+            
+            # Encode labels to integer indices
+            encoded_labels = label_encoder.transform(labels)
+            encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)
+            
+        
+            # Forward pass
+            outputs = model(images, gsr=gsr)
+            
+            # Compute loss
+            loss = criterion(outputs, encoded_labels)
+            
+            # Accumulate loss
+            running_val_loss += loss.item() * images.size(0)
+            
+            # Calculate accuracy
+            _, predicted = torch.max(outputs, 1)
+            correct_val += (predicted == encoded_labels).sum().item()
+            total_val += encoded_labels.size(0)
+
+            
+    
+    epoch_val_loss = running_val_loss / total_val if total_val > 0 else 0.0
+    val_accuracy = 100 * correct_val / total_val if total_val > 0 else 0.0
+    
+    print(f"Epoch [{epoch+1}/{num_epochs}] "
+        f"- Training Loss: {epoch_train_loss:.4f}, Training Accuracy: {train_accuracy:.2f}% "
+        f"- Validation Loss: {epoch_val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%")
+    
+    if wandb_log:
+        wandb.log({
+            "epoch_train_loss": epoch_train_loss,
+            "epoch_val_loss": epoch_val_loss,
+            "train_accuracy": train_accuracy,
+            "val_accuracy": val_accuracy,
+        })
+    if save_ckpt:
+        outdir = os.path.abspath(f'checkpoints/{model_name+"HCP_FT"}')
+        os.makedirs(outdir, exist_ok=True)
+        print("outdir", outdir)
+        # Save model and config
+        torch.save(model.state_dict(), f"{outdir}/model.pth")
+        with open(f"{outdir}/config.yaml", 'w') as f:
+            yaml.dump(wandb_config, f)
+        print(f"Saved model and config to {outdir}")
+    
+

fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/files/output.log

@@ -0,0 +1,2 @@
+Epoch 1/20 - Training:   1%|          | 106/13913 [00:56<1:19:52,  2.88it/s]
+Step [100/13913] - Training Loss: 2.6343 - Training Accuracy: 10.25%

fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/files/requirements.txt

@@ -0,0 +1,198 @@
+protobuf==5.28.2
+imageio==2.35.1
+MarkupSafe==3.0.0
+regex==2024.9.11
+matplotlib==3.9.2
+notebook==7.2.2
+debugpy==1.8.6
+aiosignal==1.3.1
+jupyter_core==5.7.2
+torchaudio==2.4.1+cu121
+python-json-logger==2.0.7
+six==1.16.0
+scikit-image==0.24.0
+types-python-dateutil==2.9.0.20241003
+PyYAML==6.0.2
+httpcore==1.0.6
+clip==1.0
+babel==2.16.0
+webcolors==24.8.0
+omegaconf==2.3.0
+webencodings==0.5.1
+kiwisolver==1.4.7
+uri-template==1.3.0
+diffusers==0.23.0
+idna==3.10
+fsspec==2024.9.0
+parso==0.8.4
+setuptools==65.5.0
+tornado==6.4.1
+webdataset==0.2.100
+decord==0.6.0
+nvidia-curand-cu12==10.3.2.106
+ipykernel==6.29.5
+jupyter==1.1.1
+pexpect==4.9.0
+kornia_rs==0.1.5
+iopath==0.1.10
+async-lru==2.0.4
+future==1.0.0
+torchvision==0.19.1+cu121
+botocore==1.34.162
+cycler==0.12.1
+tzdata==2024.2
+jupyter_server_terminals==0.5.3
+click==8.1.7
+einops==0.8.0
+pyzmq==26.2.0
+jupyter_client==8.6.3
+nbconvert==7.16.4
+scikit-learn==1.5.2
+executing==2.1.0
+asttokens==2.4.1
+docker-pycreds==0.4.0
+matplotlib-inline==0.1.7
+overrides==7.7.0
+websocket-client==1.8.0
+nbformat==5.10.4
+elbow==0.1.1
+contourpy==1.3.0
+nvidia-cudnn-cu12==9.1.0.70
+transformers==4.44.2
+gitdb==4.0.11
+jupyterlab_nvdashboard==0.11.0
+lazy_loader==0.4
+jsonpointer==3.0.0
+notebook_shim==0.2.4
+nvidia-nccl-cu12==2.20.5
+ffmpeg-python==0.2.0
+triton==3.0.0
+mistune==3.0.2
+python-dateutil==2.9.0.post0
+beautifulsoup4==4.12.3
+nbclient==0.10.0
+h5py==3.12.1
+ftfy==6.2.3
+zipp==3.20.2
+ptyprocess==0.7.0
+huggingface-hub==0.25.1
+pytz==2024.2
+jupyterlab_pygments==0.3.0
+nvidia-cublas-cu12==12.1.3.1
+pandocfilters==1.5.1
+Jinja2==3.1.4
+arrow==1.3.0
+rpds-py==0.20.0
+jupyter_server==2.14.2
+simplejson==3.19.3
+networkx==3.3
+packaging==24.1
+traitlets==5.14.3
+pandas==2.2.3
+xformers==0.0.22.post7
+lightning-utilities==0.11.7
+tifffile==2024.9.20
+nvidia-cuda-cupti-cu12==12.1.105
+mpmath==1.3.0
+GitPython==3.1.43
+scipy==1.14.1
+jsonschema==4.23.0
+prompt_toolkit==3.0.48
+s3transfer==0.10.2
+multidict==6.1.0
+bleach==6.1.0
+sentry-sdk==2.15.0
+nibabel==5.2.1
+accelerate==1.0.0
+pyarrow==17.0.0
+threadpoolctl==3.5.0
+attrs==24.2.0
+rfc3986-validator==0.1.1
+nvidia-cuda-runtime-cu12==12.1.105
+ipywidgets==8.1.5
+frozenlist==1.4.1
+pycparser==2.22
+jupyterlab_server==2.27.3
+nvidia-cuda-nvrtc-cu12==12.1.105
+yarl==1.13.1
+setproctitle==1.3.3
+isoduration==20.11.0
+Pygments==2.18.0
+jedi==0.19.1
+boto3==1.34.57
+tokenizers==0.19.1
+referencing==0.35.1
+rfc3339-validator==0.1.4
+pillow==10.4.0
+jupyterlab==4.2.5
+stack-data==0.6.3
+h11==0.14.0
+anyio==4.6.0
+nilearn==0.10.4
+nvidia-cusolver-cu12==11.4.5.107
+tinycss2==1.3.0
+defusedxml==0.7.1
+argon2-cffi-bindings==21.2.0
+soupsieve==2.6
+nest-asyncio==1.6.0
+torchmetrics==1.3.0.post0
+tqdm==4.66.5
+cffi==1.17.1
+charset-normalizer==3.3.2
+jsonschema-specifications==2023.12.1
+decorator==5.1.1
+open_clip_torch==2.26.1
+jupyter-events==0.10.0
+smart-open==7.0.5
+antlr4-python3-runtime==4.9.3
+prometheus_client==0.21.0
+kornia==0.7.3
+typing_extensions==4.12.2
+sniffio==1.3.1
+joblib==1.4.2
+comm==0.2.2
+aiohappyeyeballs==2.4.3
+numpy==2.1.2
+braceexpand==0.1.7
+certifi==2024.8.30
+psutil==6.0.0
+pyparsing==3.1.4
+pure_eval==0.2.3
+nvidia-cusparse-cu12==12.1.0.106
+wandb==0.18.3
+urllib3==2.2.3
+smmap==5.0.1
+platformdirs==4.3.6
+torch==2.4.1+cu121
+requests==2.32.3
+json5==0.9.25
+nvidia-nvjitlink-cu12==12.6.77
+jupyterlab_widgets==3.0.13
+lxml==5.3.0
+httpx==0.27.2
+opencv-python==4.6.0.66
+portalocker==2.10.1
+pytorch-lightning==2.0.1
+sympy==1.13.3
+wcwidth==0.2.13
+jmespath==1.0.1
+fqdn==1.5.1
+pynvml==11.5.3
+pip==24.0
+wrapt==1.16.0
+aiohttp==3.10.9
+filelock==3.16.1
+fonttools==4.54.1
+fastjsonschema==2.20.0
+jupyter-console==6.6.3
+widgetsnbextension==4.0.13
+timm==1.0.9
+nvidia-cufft-cu12==11.0.2.54
+ipython==8.28.0
+nvidia-nvtx-cu12==12.1.105
+jupyter-lsp==2.2.5
+safetensors==0.4.5
+terminado==0.18.1
+argon2-cffi==23.1.0
+Send2Trash==1.8.3
+importlib_metadata==8.5.0

fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/files/wandb-metadata.json

@@ -0,0 +1,131 @@
+{
+  "os":  "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+  "python":  "3.11.10",
+  "startedAt":  "2024-10-23T04:08:30.389906Z",
+  "args":  [
+    "NSDflat_large_gsrFalse_",
+    "epoch99.pth"
+  ],
+  "program":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py",
+  "codePath":  "src/HCP_downstream_finetune.py",
+  "git":  {
+    "remote":  "https://github.com/MedARC-AI/fMRI-foundation-model",
+    "commit":  "b1ba684ae7a5cc4155cc046b0abe613de09bf700"
+  },
+  "email":  "torrico.villanueva.cesar.kadir@gmail.com",
+  "root":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+  "host":  "ip-10-0-139-117",
+  "username":  "ckadirt",
+  "executable":  "/admin/home-ckadirt/foundation_env/bin/python",
+  "codePathLocal":  "HCP_downstream_finetune.py",
+  "cpu_count":  96,
+  "cpu_count_logical":  192,
+  "gpu":  "[NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3]",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "249555763200",
+      "used":  "181767704576"
+    }
+  },
+  "memory":  {
+    "total":  "2147443380224"
+  },
+  "cpu":  {
+    "count":  96,
+    "countLogical":  192
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    }
+  ],
+  "slurm":  {
+    "cluster_name":  "sagemaker2",
+    "conf":  "/opt/slurm/etc/slurm.conf",
+    "cpus_on_node":  "20",
+    "gpus_on_node":  "1",
+    "gpus_per_task":  "1",
+    "gtids":  "0",
+    "job_account":  "fmri",
+    "job_cpus_per_node":  "20",
+    "job_end_time":  "1729699688",
+    "job_gid":  "1879800513",
+    "job_gpus":  "4",
+    "job_id":  "528150",
+    "job_name":  "finetuneHCP",
+    "job_nodelist":  "ip-10-0-139-117",
+    "job_num_nodes":  "1",
+    "job_partition":  "p5",
+    "job_qos":  "idle",
+    "job_start_time":  "1729656488",
+    "job_uid":  "1879804696",
+    "job_user":  "ckadirt",
+    "jobid":  "528150",
+    "localid":  "0",
+    "mem_per_cpu":  "11500",
+    "nnodes":  "1",
+    "node_aliases":  "(null)",
+    "nodeid":  "0",
+    "nodelist":  "ip-10-0-139-117",
+    "nprocs":  "1",
+    "ntasks":  "1",
+    "ntasks_per_node":  "1",
+    "prio_process":  "0",
+    "procid":  "0",
+    "script_context":  "prolog_task",
+    "submit_dir":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+    "submit_host":  "ip-172-17-12-61",
+    "task_pid":  "3329662",
+    "tasks_per_node":  "1",
+    "topology_addr":  "ip-10-0-139-117",
+    "topology_addr_pattern":  "node",
+    "working_cluster":  "sagemaker2:ip-172-17-63-161:6817:9984:109"
+  },
+  "cudaVersion":  "12.2"
+}

fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/logs/debug-core.log

@@ -0,0 +1,14 @@
+{"time":"2024-10-23T04:08:29.759432555Z","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmpv0zbfmyf/port-3329596.txt","pid":3329596,"debug":false,"disable-analytics":false}
+{"time":"2024-10-23T04:08:29.759434395Z","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmpdgtkqv8m/port-3329693.txt","pid":3329693,"debug":false,"disable-analytics":false}
+{"time":"2024-10-23T04:08:29.75965549Z","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false}
+{"time":"2024-10-23T04:08:29.75967947Z","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false}
+{"time":"2024-10-23T04:08:29.762594934Z","level":"INFO","msg":"Will exit if parent process dies.","ppid":3329693}
+{"time":"2024-10-23T04:08:29.762592144Z","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":43169,"Zone":""}}
+{"time":"2024-10-23T04:08:29.764024661Z","level":"INFO","msg":"Will exit if parent process dies.","ppid":3329596}
+{"time":"2024-10-23T04:08:29.764057981Z","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":45635,"Zone":""}}
+{"time":"2024-10-23T04:08:29.893047528Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:44072"}
+{"time":"2024-10-23T04:08:29.893048808Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:60086"}
+{"time":"2024-10-23T04:08:30.348153132Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_large_gsrFalse__HCP_FT_83810","id":"127.0.0.1:60086"}
+{"time":"2024-10-23T04:08:30.390297655Z","level":"INFO","msg":"handleInformInit: received","streamId":"NSDflat_large_gsrFalse__HCP_FT_83810","id":"127.0.0.1:44072"}
+{"time":"2024-10-23T04:08:30.392846482Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_large_gsrFalse__HCP_FT_83810","id":"127.0.0.1:60086"}
+{"time":"2024-10-23T04:08:30.434661309Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"NSDflat_large_gsrFalse__HCP_FT_83810","id":"127.0.0.1:44072"}

fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/logs/debug-internal.log

@@ -0,0 +1,11 @@
+{"time":"2024-10-23T04:08:30.398164322Z","level":"INFO","msg":"using version","core version":"0.18.3"}
+{"time":"2024-10-23T04:08:30.398178952Z","level":"INFO","msg":"created symlink","path":"/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/logs/debug-core.log"}
+{"time":"2024-10-23T04:08:30.400590046Z","level":"ERROR","msg":"dialing: google: could not find default credentials. See https://cloud.google.com/docs/authentication/external/set-up-adc for more information"}
+{"time":"2024-10-23T04:08:30.434629119Z","level":"INFO","msg":"created new stream","id":"NSDflat_large_gsrFalse__HCP_FT_83810"}
+{"time":"2024-10-23T04:08:30.434655859Z","level":"INFO","msg":"stream: started","id":"NSDflat_large_gsrFalse__HCP_FT_83810"}
+{"time":"2024-10-23T04:08:30.434669539Z","level":"INFO","msg":"writer: Do: started","stream_id":{"value":"NSDflat_large_gsrFalse__HCP_FT_83810"}}
+{"time":"2024-10-23T04:08:30.43470082Z","level":"INFO","msg":"handler: started","stream_id":{"value":"NSDflat_large_gsrFalse__HCP_FT_83810"}}
+{"time":"2024-10-23T04:08:30.4347016Z","level":"INFO","msg":"sender: started","stream_id":{"value":"NSDflat_large_gsrFalse__HCP_FT_83810"}}
+{"time":"2024-10-23T04:08:30.880621614Z","level":"INFO","msg":"wandb-core","!BADKEY":null}
+{"time":"2024-10-23T04:08:30.884818882Z","level":"INFO","msg":"Starting system monitor"}
+{"time":"2024-10-23T04:08:30.912744151Z","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}

fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/logs/debug.log

@@ -0,0 +1,26 @@
+2024-10-23 04:08:30,377 INFO    MainThread:3329693 [wandb_setup.py:_flush():79] Current SDK version is 0.18.3
+2024-10-23 04:08:30,377 INFO    MainThread:3329693 [wandb_setup.py:_flush():79] Configure stats pid to 3329693
+2024-10-23 04:08:30,377 INFO    MainThread:3329693 [wandb_setup.py:_flush():79] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-10-23 04:08:30,377 INFO    MainThread:3329693 [wandb_setup.py:_flush():79] Loading settings from /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/settings
+2024-10-23 04:08:30,377 INFO    MainThread:3329693 [wandb_setup.py:_flush():79] Loading settings from environment variables: {}
+2024-10-23 04:08:30,377 INFO    MainThread:3329693 [wandb_setup.py:_flush():79] Applying setup settings: {'mode': None, '_disable_service': None}
+2024-10-23 04:08:30,377 INFO    MainThread:3329693 [wandb_setup.py:_flush():79] Inferring run settings from compute environment: {'program_relpath': 'src/HCP_downstream_finetune.py', 'program_abspath': '/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py', 'program': '/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py'}
+2024-10-23 04:08:30,377 INFO    MainThread:3329693 [wandb_setup.py:_flush():79] Applying login settings: {}
+2024-10-23 04:08:30,379 INFO    MainThread:3329693 [wandb_init.py:_log_setup():532] Logging user logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/logs/debug.log
+2024-10-23 04:08:30,380 INFO    MainThread:3329693 [wandb_init.py:_log_setup():533] Logging internal logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_040830-NSDflat_large_gsrFalse__HCP_FT_83810/logs/debug-internal.log
+2024-10-23 04:08:30,380 INFO    MainThread:3329693 [wandb_init.py:init():617] calling init triggers
+2024-10-23 04:08:30,380 INFO    MainThread:3329693 [wandb_init.py:init():624] wandb.init called with sweep_config: {}
+config: {'model_name': 'NSDflat_large_gsrFalse__HCP_FT', 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 1e-05, 'num_epochs': 20, 'seed': 42}
+2024-10-23 04:08:30,380 INFO    MainThread:3329693 [wandb_init.py:init():667] starting backend
+2024-10-23 04:08:30,380 INFO    MainThread:3329693 [wandb_init.py:init():671] sending inform_init request
+2024-10-23 04:08:30,389 INFO    MainThread:3329693 [backend.py:_multiprocessing_setup():104] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-10-23 04:08:30,389 INFO    MainThread:3329693 [wandb_init.py:init():684] backend started and connected
+2024-10-23 04:08:30,409 INFO    MainThread:3329693 [wandb_init.py:init():779] updated telemetry
+2024-10-23 04:08:30,435 INFO    MainThread:3329693 [wandb_init.py:init():812] communicating run to backend with 90.0 second timeout
+2024-10-23 04:08:30,824 INFO    MainThread:3329693 [wandb_init.py:init():855] run resumed
+2024-10-23 04:08:30,866 INFO    MainThread:3329693 [wandb_init.py:init():863] starting run threads in backend
+2024-10-23 04:08:31,304 INFO    MainThread:3329693 [wandb_run.py:_console_start():2465] atexit reg
+2024-10-23 04:08:31,304 INFO    MainThread:3329693 [wandb_run.py:_redirect():2313] redirect: wrap_raw
+2024-10-23 04:08:31,304 INFO    MainThread:3329693 [wandb_run.py:_redirect():2378] Wrapping output streams.
+2024-10-23 04:08:31,304 INFO    MainThread:3329693 [wandb_run.py:_redirect():2403] Redirects installed.
+2024-10-23 04:08:31,310 INFO    MainThread:3329693 [wandb_init.py:init():907] run started, returning control to user process

fMRI-foundation-model/src/wandb/run-20241023_041122-HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f/files/requirements.txt

@@ -0,0 +1,198 @@
+protobuf==5.28.2
+imageio==2.35.1
+MarkupSafe==3.0.0
+regex==2024.9.11
+matplotlib==3.9.2
+notebook==7.2.2
+debugpy==1.8.6
+aiosignal==1.3.1
+jupyter_core==5.7.2
+torchaudio==2.4.1+cu121
+python-json-logger==2.0.7
+six==1.16.0
+scikit-image==0.24.0
+types-python-dateutil==2.9.0.20241003
+PyYAML==6.0.2
+httpcore==1.0.6
+clip==1.0
+babel==2.16.0
+webcolors==24.8.0
+omegaconf==2.3.0
+webencodings==0.5.1
+kiwisolver==1.4.7
+uri-template==1.3.0
+diffusers==0.23.0
+idna==3.10
+fsspec==2024.9.0
+parso==0.8.4
+setuptools==65.5.0
+tornado==6.4.1
+webdataset==0.2.100
+decord==0.6.0
+nvidia-curand-cu12==10.3.2.106
+ipykernel==6.29.5
+jupyter==1.1.1
+pexpect==4.9.0
+kornia_rs==0.1.5
+iopath==0.1.10
+async-lru==2.0.4
+future==1.0.0
+torchvision==0.19.1+cu121
+botocore==1.34.162
+cycler==0.12.1
+tzdata==2024.2
+jupyter_server_terminals==0.5.3
+click==8.1.7
+einops==0.8.0
+pyzmq==26.2.0
+jupyter_client==8.6.3
+nbconvert==7.16.4
+scikit-learn==1.5.2
+executing==2.1.0
+asttokens==2.4.1
+docker-pycreds==0.4.0
+matplotlib-inline==0.1.7
+overrides==7.7.0
+websocket-client==1.8.0
+nbformat==5.10.4
+elbow==0.1.1
+contourpy==1.3.0
+nvidia-cudnn-cu12==9.1.0.70
+transformers==4.44.2
+gitdb==4.0.11
+jupyterlab_nvdashboard==0.11.0
+lazy_loader==0.4
+jsonpointer==3.0.0
+notebook_shim==0.2.4
+nvidia-nccl-cu12==2.20.5
+ffmpeg-python==0.2.0
+triton==3.0.0
+mistune==3.0.2
+python-dateutil==2.9.0.post0
+beautifulsoup4==4.12.3
+nbclient==0.10.0
+h5py==3.12.1
+ftfy==6.2.3
+zipp==3.20.2
+ptyprocess==0.7.0
+huggingface-hub==0.25.1
+pytz==2024.2
+jupyterlab_pygments==0.3.0
+nvidia-cublas-cu12==12.1.3.1
+pandocfilters==1.5.1
+Jinja2==3.1.4
+arrow==1.3.0
+rpds-py==0.20.0
+jupyter_server==2.14.2
+simplejson==3.19.3
+networkx==3.3
+packaging==24.1
+traitlets==5.14.3
+pandas==2.2.3
+xformers==0.0.22.post7
+lightning-utilities==0.11.7
+tifffile==2024.9.20
+nvidia-cuda-cupti-cu12==12.1.105
+mpmath==1.3.0
+GitPython==3.1.43
+scipy==1.14.1
+jsonschema==4.23.0
+prompt_toolkit==3.0.48
+s3transfer==0.10.2
+multidict==6.1.0
+bleach==6.1.0
+sentry-sdk==2.15.0
+nibabel==5.2.1
+accelerate==1.0.0
+pyarrow==17.0.0
+threadpoolctl==3.5.0
+attrs==24.2.0
+rfc3986-validator==0.1.1
+nvidia-cuda-runtime-cu12==12.1.105
+ipywidgets==8.1.5
+frozenlist==1.4.1
+pycparser==2.22
+jupyterlab_server==2.27.3
+nvidia-cuda-nvrtc-cu12==12.1.105
+yarl==1.13.1
+setproctitle==1.3.3
+isoduration==20.11.0
+Pygments==2.18.0
+jedi==0.19.1
+boto3==1.34.57
+tokenizers==0.19.1
+referencing==0.35.1
+rfc3339-validator==0.1.4
+pillow==10.4.0
+jupyterlab==4.2.5
+stack-data==0.6.3
+h11==0.14.0
+anyio==4.6.0
+nilearn==0.10.4
+nvidia-cusolver-cu12==11.4.5.107
+tinycss2==1.3.0
+defusedxml==0.7.1
+argon2-cffi-bindings==21.2.0
+soupsieve==2.6
+nest-asyncio==1.6.0
+torchmetrics==1.3.0.post0
+tqdm==4.66.5
+cffi==1.17.1
+charset-normalizer==3.3.2
+jsonschema-specifications==2023.12.1
+decorator==5.1.1
+open_clip_torch==2.26.1
+jupyter-events==0.10.0
+smart-open==7.0.5
+antlr4-python3-runtime==4.9.3
+prometheus_client==0.21.0
+kornia==0.7.3
+typing_extensions==4.12.2
+sniffio==1.3.1
+joblib==1.4.2
+comm==0.2.2
+aiohappyeyeballs==2.4.3
+numpy==2.1.2
+braceexpand==0.1.7
+certifi==2024.8.30
+psutil==6.0.0
+pyparsing==3.1.4
+pure_eval==0.2.3
+nvidia-cusparse-cu12==12.1.0.106
+wandb==0.18.3
+urllib3==2.2.3
+smmap==5.0.1
+platformdirs==4.3.6
+torch==2.4.1+cu121
+requests==2.32.3
+json5==0.9.25
+nvidia-nvjitlink-cu12==12.6.77
+jupyterlab_widgets==3.0.13
+lxml==5.3.0
+httpx==0.27.2
+opencv-python==4.6.0.66
+portalocker==2.10.1
+pytorch-lightning==2.0.1
+sympy==1.13.3
+wcwidth==0.2.13
+jmespath==1.0.1
+fqdn==1.5.1
+pynvml==11.5.3
+pip==24.0
+wrapt==1.16.0
+aiohttp==3.10.9
+filelock==3.16.1
+fonttools==4.54.1
+fastjsonschema==2.20.0
+jupyter-console==6.6.3
+widgetsnbextension==4.0.13
+timm==1.0.9
+nvidia-cufft-cu12==11.0.2.54
+ipython==8.28.0
+nvidia-nvtx-cu12==12.1.105
+jupyter-lsp==2.2.5
+safetensors==0.4.5
+terminado==0.18.1
+argon2-cffi==23.1.0
+Send2Trash==1.8.3
+importlib_metadata==8.5.0

fMRI-foundation-model/src/wandb/run-20241023_041122-HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f/files/wandb-metadata.json

@@ -0,0 +1,144 @@
+{
+  "os":  "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+  "python":  "3.11.10",
+  "startedAt":  "2024-10-23T04:11:25.180317Z",
+  "program":  "ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.ipynb",
+  "git":  {
+    "remote":  "https://github.com/MedARC-AI/fMRI-foundation-model",
+    "commit":  "b1ba684ae7a5cc4155cc046b0abe613de09bf700"
+  },
+  "email":  "torrico.villanueva.cesar.kadir@gmail.com",
+  "root":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+  "host":  "ip-10-0-160-143",
+  "username":  "ckadirt",
+  "executable":  "/admin/home-ckadirt/foundation_env/bin/python",
+  "cpu_count":  96,
+  "cpu_count_logical":  192,
+  "gpu":  "[NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3]",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "249555763200",
+      "used":  "185027031040"
+    }
+  },
+  "memory":  {
+    "total":  "2147443429376"
+  },
+  "cpu":  {
+    "count":  96,
+    "countLogical":  192
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    }
+  ],
+  "slurm":  {
+    "cluster_name":  "sagemaker2",
+    "conf":  "/opt/slurm/etc/slurm.conf",
+    "cpu_bind":  "quiet,mask_cpu:0x00000000000000FFC000000000000000000000FFC0000000",
+    "cpu_bind_list":  "0x00000000000000FFC000000000000000000000FFC0000000",
+    "cpu_bind_type":  "mask_cpu:",
+    "cpu_bind_verbose":  "quiet",
+    "cpus_on_node":  "20",
+    "gpus":  "1",
+    "gpus_on_node":  "1",
+    "gtids":  "0",
+    "job_account":  "fmri",
+    "job_cpus_per_node":  "20",
+    "job_end_time":  "1729702756",
+    "job_gid":  "1879800513",
+    "job_group":  "Domain Users",
+    "job_id":  "528040",
+    "job_name":  "bash",
+    "job_nodelist":  "ip-10-0-160-143",
+    "job_num_nodes":  "1",
+    "job_partition":  "p5",
+    "job_qos":  "idle",
+    "job_start_time":  "1729648756",
+    "job_uid":  "1879804696",
+    "job_user":  "ckadirt",
+    "jobid":  "528040",
+    "launch_node_ipaddr":  "172.17.12.61",
+    "localid":  "0",
+    "mpi_type":  "pmix_v3",
+    "nnodes":  "1",
+    "nodeid":  "0",
+    "nodelist":  "ip-10-0-160-143",
+    "nprocs":  "1",
+    "ntasks":  "1",
+    "pmix_mapping_serv":  "(vector,(0,1,1))",
+    "pmixp_abort_agent_port":  "34923",
+    "prio_process":  "0",
+    "procid":  "0",
+    "pty_port":  "45733",
+    "pty_win_col":  "199",
+    "pty_win_row":  "17",
+    "script_context":  "prolog_task",
+    "srun_comm_host":  "172.17.12.61",
+    "srun_comm_port":  "39353",
+    "step_gpus":  "3",
+    "step_id":  "0",
+    "step_launcher_port":  "39353",
+    "step_nodelist":  "ip-10-0-160-143",
+    "step_num_nodes":  "1",
+    "step_num_tasks":  "1",
+    "step_tasks_per_node":  "1",
+    "stepid":  "0",
+    "submit_dir":  "/weka/proj-fmri",
+    "submit_host":  "ip-172-17-12-61",
+    "task_pid":  "1032669",
+    "tasks_per_node":  "1",
+    "topology_addr":  "ip-10-0-160-143",
+    "topology_addr_pattern":  "node",
+    "umask":  "0022",
+    "working_cluster":  "sagemaker2:ip-172-17-63-161:6817:9984:109"
+  },
+  "cudaVersion":  "12.2"
+}

fMRI-foundation-model/src/wandb/run-20241023_041122-HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f/logs/debug-core.log

@@ -0,0 +1,12 @@
+{"time":"2024-10-23T04:00:15.891404575Z","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmp623bhuv_/port-1087257.txt","pid":1087257,"debug":false,"disable-analytics":false}
+{"time":"2024-10-23T04:00:15.891695701Z","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false}
+{"time":"2024-10-23T04:00:15.894505677Z","level":"INFO","msg":"Will exit if parent process dies.","ppid":1087257}
+{"time":"2024-10-23T04:00:15.894459316Z","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":40341,"Zone":""}}
+{"time":"2024-10-23T04:00:16.078034023Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:58780"}
+{"time":"2024-10-23T04:00:16.909578833Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_large_gsrFalse__HCP_FT_83810","id":"127.0.0.1:58780"}
+{"time":"2024-10-23T04:00:17.021711364Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_large_gsrFalse__HCP_FT_83810","id":"127.0.0.1:58780"}
+{"time":"2024-10-23T04:11:25.031512552Z","level":"INFO","msg":"handleInformFinish: finish message received","streamId":"HCPflat_large_gsrFalse__HCP_FT_83810","id":"127.0.0.1:58780"}
+{"time":"2024-10-23T04:11:25.032442271Z","level":"INFO","msg":"handleInformFinish: stream closed","streamId":"HCPflat_large_gsrFalse__HCP_FT_83810","id":"127.0.0.1:58780"}
+{"time":"2024-10-23T04:11:25.179399522Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f","id":"127.0.0.1:58780"}
+{"time":"2024-10-23T04:11:25.24570696Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f","id":"127.0.0.1:58780"}
+{"time":"2024-10-23T04:19:30.473127709Z","level":"INFO","msg":"Parent process exited, terminating service process."}

fMRI-foundation-model/src/wandb/run-20241023_041122-HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f/logs/debug-internal.log

@@ -0,0 +1,12 @@
+{"time":"2024-10-23T04:11:25.205972524Z","level":"INFO","msg":"using version","core version":"0.18.3"}
+{"time":"2024-10-23T04:11:25.205987165Z","level":"INFO","msg":"created symlink","path":"/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041122-HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f/logs/debug-core.log"}
+{"time":"2024-10-23T04:11:25.206996565Z","level":"ERROR","msg":"dialing: google: could not find default credentials. See https://cloud.google.com/docs/authentication/external/set-up-adc for more information"}
+{"time":"2024-10-23T04:11:25.245647669Z","level":"INFO","msg":"created new stream","id":"HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f"}
+{"time":"2024-10-23T04:11:25.24569935Z","level":"INFO","msg":"stream: started","id":"HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f"}
+{"time":"2024-10-23T04:11:25.24572727Z","level":"INFO","msg":"handler: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f"}}
+{"time":"2024-10-23T04:11:25.24571807Z","level":"INFO","msg":"sender: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f"}}
+{"time":"2024-10-23T04:11:25.24571519Z","level":"INFO","msg":"writer: Do: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f"}}
+{"time":"2024-10-23T04:11:25.87412901Z","level":"INFO","msg":"wandb-core","!BADKEY":null}
+{"time":"2024-10-23T04:11:25.884775353Z","level":"INFO","msg":"Starting system monitor"}
+{"time":"2024-10-23T04:11:25.884793743Z","level":"WARN","msg":"handleCodeSave: program relative path is empty"}
+{"time":"2024-10-23T04:11:25.886547389Z","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}

fMRI-foundation-model/src/wandb/run-20241023_041122-HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f/logs/debug.log

@@ -0,0 +1,39 @@
+2024-10-23 04:11:22,311 INFO    MainThread:1087257 [wandb_setup.py:_flush():79] Current SDK version is 0.18.3
+2024-10-23 04:11:22,311 INFO    MainThread:1087257 [wandb_setup.py:_flush():79] Configure stats pid to 1087257
+2024-10-23 04:11:22,311 INFO    MainThread:1087257 [wandb_setup.py:_flush():79] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-10-23 04:11:22,311 INFO    MainThread:1087257 [wandb_setup.py:_flush():79] Loading settings from /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/settings
+2024-10-23 04:11:22,311 INFO    MainThread:1087257 [wandb_setup.py:_flush():79] Loading settings from environment variables: {}
+2024-10-23 04:11:22,312 INFO    MainThread:1087257 [wandb_setup.py:_flush():79] Applying setup settings: {'mode': None, '_disable_service': None}
+2024-10-23 04:11:22,312 INFO    MainThread:1087257 [wandb_setup.py:_flush():79] Inferring run settings from compute environment: {'program': '<python with no main file>'}
+2024-10-23 04:11:22,312 INFO    MainThread:1087257 [wandb_setup.py:_flush():79] Applying login settings: {}
+2024-10-23 04:11:22,312 INFO    MainThread:1087257 [wandb_init.py:_log_setup():532] Logging user logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041122-HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f/logs/debug.log
+2024-10-23 04:11:22,313 INFO    MainThread:1087257 [wandb_init.py:_log_setup():533] Logging internal logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041122-HCPflat_large_gsrFalse__HCP_FT_7ee35929-85c0-47da-91ab-dac2776c444f/logs/debug-internal.log
+2024-10-23 04:11:22,313 INFO    MainThread:1087257 [wandb_init.py:init():617] calling init triggers
+2024-10-23 04:11:22,313 INFO    MainThread:1087257 [wandb_init.py:init():624] wandb.init called with sweep_config: {}
+config: {'model_name': 'HCPflat_large_gsrFalse__HCP_FT', 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 1e-05, 'num_epochs': 20, 'seed': 42}
+2024-10-23 04:11:22,313 INFO    MainThread:1087257 [wandb_init.py:init():642] re-initializing run, found existing run on stack: HCPflat_large_gsrFalse__HCP_FT_83810
+2024-10-23 04:11:22,315 INFO    MainThread:1087257 [wandb_run.py:_finish():2164] finishing run ckadirt/fMRI-foundation-model/HCPflat_large_gsrFalse__HCP_FT_83810
+2024-10-23 04:11:22,357 INFO    MainThread:1087257 [jupyter.py:save_history():488] saving 17 cells to _session_history.ipynb
+2024-10-23 04:11:22,358 INFO    MainThread:1087257 [wandb_run.py:_config_callback():1394] config_cb ('_wandb', 'session_history') code/_session_history.ipynb None
+2024-10-23 04:11:22,410 INFO    MainThread:1087257 [jupyter.py:_save_ipynb():398] looking for notebook: ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.ipynb
+2024-10-23 04:11:22,410 INFO    MainThread:1087257 [wandb_init.py:_jupyter_teardown():460] cleaning up jupyter logic
+2024-10-23 04:11:22,410 INFO    MainThread:1087257 [wandb_run.py:_atexit_cleanup():2428] got exitcode: 0
+2024-10-23 04:11:22,412 INFO    MainThread:1087257 [wandb_run.py:_restore():2410] restore
+2024-10-23 04:11:22,413 INFO    MainThread:1087257 [wandb_run.py:_restore():2416] restore done
+2024-10-23 04:11:25,001 INFO    MainThread:1087257 [wandb_run.py:_footer_history_summary_info():4049] rendering history
+2024-10-23 04:11:25,001 INFO    MainThread:1087257 [wandb_run.py:_footer_history_summary_info():4081] rendering summary
+2024-10-23 04:11:25,028 INFO    MainThread:1087257 [wandb_run.py:_footer_sync_info():4008] logging synced files
+2024-10-23 04:11:25,164 INFO    MainThread:1087257 [wandb_init.py:init():667] starting backend
+2024-10-23 04:11:25,164 INFO    MainThread:1087257 [wandb_init.py:init():671] sending inform_init request
+2024-10-23 04:11:25,179 INFO    MainThread:1087257 [backend.py:_multiprocessing_setup():104] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-10-23 04:11:25,179 INFO    MainThread:1087257 [wandb_init.py:init():684] backend started and connected
+2024-10-23 04:11:25,203 INFO    MainThread:1087257 [wandb_run.py:_label_probe_notebook():1346] probe notebook
+2024-10-23 04:11:25,204 INFO    MainThread:1087257 [wandb_run.py:_label_probe_notebook():1356] Unable to probe notebook: 'NoneType' object has no attribute 'get'
+2024-10-23 04:11:25,204 INFO    MainThread:1087257 [wandb_init.py:init():779] updated telemetry
+2024-10-23 04:11:25,319 INFO    MainThread:1087257 [wandb_init.py:init():812] communicating run to backend with 90.0 second timeout
+2024-10-23 04:11:25,855 INFO    MainThread:1087257 [wandb_init.py:init():863] starting run threads in backend
+2024-10-23 04:11:26,356 INFO    MainThread:1087257 [wandb_run.py:_console_start():2465] atexit reg
+2024-10-23 04:11:26,356 INFO    MainThread:1087257 [wandb_run.py:_redirect():2313] redirect: wrap_raw
+2024-10-23 04:11:26,357 INFO    MainThread:1087257 [wandb_run.py:_redirect():2378] Wrapping output streams.
+2024-10-23 04:11:26,357 INFO    MainThread:1087257 [wandb_run.py:_redirect():2403] Redirects installed.
+2024-10-23 04:11:26,357 INFO    MainThread:1087257 [wandb_init.py:init():907] run started, returning control to user process

fMRI-foundation-model/src/wandb/run-20241023_041223-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3/logs/debug-internal.log

@@ -0,0 +1,11 @@
+{"time":"2024-10-23T04:12:23.854291574Z","level":"INFO","msg":"using version","core version":"0.18.3"}
+{"time":"2024-10-23T04:12:23.854310795Z","level":"INFO","msg":"created symlink","path":"/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041223-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3/logs/debug-core.log"}
+{"time":"2024-10-23T04:12:23.856970424Z","level":"ERROR","msg":"dialing: google: could not find default credentials. See https://cloud.google.com/docs/authentication/external/set-up-adc for more information"}
+{"time":"2024-10-23T04:12:23.904377013Z","level":"INFO","msg":"created new stream","id":"NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3"}
+{"time":"2024-10-23T04:12:23.904404853Z","level":"INFO","msg":"stream: started","id":"NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3"}
+{"time":"2024-10-23T04:12:23.904415433Z","level":"INFO","msg":"sender: started","stream_id":{"value":"NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3"}}
+{"time":"2024-10-23T04:12:23.904437214Z","level":"INFO","msg":"writer: Do: started","stream_id":{"value":"NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3"}}
+{"time":"2024-10-23T04:12:23.904419533Z","level":"INFO","msg":"handler: started","stream_id":{"value":"NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3"}}
+{"time":"2024-10-23T04:12:24.41246663Z","level":"INFO","msg":"wandb-core","!BADKEY":null}
+{"time":"2024-10-23T04:12:24.425346089Z","level":"INFO","msg":"Starting system monitor"}
+{"time":"2024-10-23T04:12:24.491097327Z","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}

fMRI-foundation-model/src/wandb/run-20241023_041223-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3/logs/debug.log

@@ -0,0 +1,25 @@
+2024-10-23 04:12:23,820 INFO    MainThread:3333420 [wandb_setup.py:_flush():79] Current SDK version is 0.18.3
+2024-10-23 04:12:23,820 INFO    MainThread:3333420 [wandb_setup.py:_flush():79] Configure stats pid to 3333420
+2024-10-23 04:12:23,820 INFO    MainThread:3333420 [wandb_setup.py:_flush():79] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-10-23 04:12:23,820 INFO    MainThread:3333420 [wandb_setup.py:_flush():79] Loading settings from /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/settings
+2024-10-23 04:12:23,820 INFO    MainThread:3333420 [wandb_setup.py:_flush():79] Loading settings from environment variables: {}
+2024-10-23 04:12:23,820 INFO    MainThread:3333420 [wandb_setup.py:_flush():79] Applying setup settings: {'mode': None, '_disable_service': None}
+2024-10-23 04:12:23,820 INFO    MainThread:3333420 [wandb_setup.py:_flush():79] Inferring run settings from compute environment: {'program_relpath': 'src/HCP_downstream_finetune.py', 'program_abspath': '/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py', 'program': '/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py'}
+2024-10-23 04:12:23,820 INFO    MainThread:3333420 [wandb_setup.py:_flush():79] Applying login settings: {}
+2024-10-23 04:12:23,821 INFO    MainThread:3333420 [wandb_init.py:_log_setup():532] Logging user logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041223-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3/logs/debug.log
+2024-10-23 04:12:23,822 INFO    MainThread:3333420 [wandb_init.py:_log_setup():533] Logging internal logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041223-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3/logs/debug-internal.log
+2024-10-23 04:12:23,822 INFO    MainThread:3333420 [wandb_init.py:init():617] calling init triggers
+2024-10-23 04:12:23,822 INFO    MainThread:3333420 [wandb_init.py:init():624] wandb.init called with sweep_config: {}
+config: {'model_name': 'NSDflat_large_gsrFalse__HCP_FT', 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 1e-05, 'num_epochs': 20, 'seed': 42}
+2024-10-23 04:12:23,822 INFO    MainThread:3333420 [wandb_init.py:init():667] starting backend
+2024-10-23 04:12:23,822 INFO    MainThread:3333420 [wandb_init.py:init():671] sending inform_init request
+2024-10-23 04:12:23,831 INFO    MainThread:3333420 [backend.py:_multiprocessing_setup():104] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-10-23 04:12:23,831 INFO    MainThread:3333420 [wandb_init.py:init():684] backend started and connected
+2024-10-23 04:12:23,852 INFO    MainThread:3333420 [wandb_init.py:init():779] updated telemetry
+2024-10-23 04:12:23,895 INFO    MainThread:3333420 [wandb_init.py:init():812] communicating run to backend with 90.0 second timeout
+2024-10-23 04:12:24,396 INFO    MainThread:3333420 [wandb_init.py:init():863] starting run threads in backend
+2024-10-23 04:12:24,755 INFO    MainThread:3333420 [wandb_run.py:_console_start():2465] atexit reg
+2024-10-23 04:12:24,755 INFO    MainThread:3333420 [wandb_run.py:_redirect():2313] redirect: wrap_raw
+2024-10-23 04:12:24,755 INFO    MainThread:3333420 [wandb_run.py:_redirect():2378] Wrapping output streams.
+2024-10-23 04:12:24,756 INFO    MainThread:3333420 [wandb_run.py:_redirect():2403] Redirects installed.
+2024-10-23 04:12:24,758 INFO    MainThread:3333420 [wandb_init.py:init():907] run started, returning control to user process

fMRI-foundation-model/src/wandb/run-20241023_041223-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3/run-NSDflat_large_gsrFalse__HCP_FT_9df42194-aaea-4534-8c08-2364d19544b3.wandb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b94ade7deadb51db63431f075745b200d705982cf2fd53012bc0d3a272293b0
+size 3604480

fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/files/code/src/HCP_downstream_finetune.py

@@ -0,0 +1,596 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# Import packages and setup gpu configuration.
+# This code block shouldnt need to be adjusted!
+import os
+import sys
+import json
+import yaml
+import numpy as np
+import copy
+import math
+import time
+import random
+from tqdm.auto import tqdm
+import webdataset as wds
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import utils
+from mae_utils.flat_models import *
+import h5py
+from mae_utils import flat_models
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+# following fixes a Conv3D CUDNN_NOT_SUPPORTED error
+torch.backends.cudnn.benchmark = True
+
+# ## MODEL TO LOAD ##
+if utils.is_interactive():
+    model_name = "HCPflat_large_gsrFalse_"
+else:
+    model_name = sys.argv[1]
+    
+
+# outdir = os.path.abspath(f'checkpoints/{model_name}')
+outdir = os.path.abspath(f'checkpoints/{model_name}')
+
+print("outdir", outdir)
+# Load previous config.yaml if available
+if os.path.exists(f"{outdir}/config.yaml"):
+    config = yaml.load(open(f"{outdir}/config.yaml", 'r'), Loader=yaml.FullLoader)
+    print(f"Loaded config.yaml from ckpt folder {outdir}")
+    # create global variables from the config
+    print("\n__CONFIG__")
+    for attribute_name in config.keys():
+        print(f"{attribute_name} = {config[attribute_name]}")
+        globals()[attribute_name] = config[f'{attribute_name}']
+    print("\n")
+
+world_size = os.getenv('WORLD_SIZE')
+if world_size is None: 
+    world_size = 1
+else:
+    world_size = int(world_size)
+print(f"WORLD_SIZE={world_size}")
+
+if utils.is_interactive():
+    # Following allows you to change functions in models.py or utils.py and 
+    # have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    get_ipython().run_line_magic('autoreload', '2')
+
+batch_size = probe_batch_size
+num_epochs = probe_num_epochs
+
+data_type = torch.float32 # change depending on your mixed_precision
+global_batch_size = batch_size * world_size
+
+device = torch.device('cuda')
+
+hcp_flat_path = "/weka/proj-medarc/shared/HCP-Flat"
+# seed = 42
+# num_frames = 16
+# gsr = False
+# num_workers = 10
+# batch_size = 128
+
+print("PID of this process =",os.getpid())
+utils.seed_everything(seed)
+
+
+# In[2]:
+
+
+if os.getenv('global_pool') == "False":
+    global_pool = False
+else:
+    global_pool = True
+print(f"global_pool = {global_pool}")
+
+try:
+    gsr
+except:
+    gsr = True
+    print("set gsr to True")
+print(f"gsr = {gsr}")
+
+
+# In[3]:
+
+
+#### UNCOMMENT THIS TO SAVE THE HCP-FLAT IN HDF5 FORMAT
+
+
+# from torch.utils.data import default_collate
+# from mae_utils.flat import load_hcp_flat_mask
+# from mae_utils.flat import create_hcp_flat
+# from mae_utils.flat import batch_unmask
+# import mae_utils.visualize as vis
+
+
+# batch_size = 26
+# print(f"changed batch_size to {batch_size}")
+
+# ## Test ##
+# datasets_to_include = "HCP"
+# assert "HCP" in datasets_to_include
+# test_dataset = create_hcp_flat(root=hcp_flat_path, 
+#                 clip_mode="event", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'test')
+# test_dl = wds.WebLoader(
+#     test_dataset.batched(batch_size, partial=False, collation_fn=default_collate),
+#     batch_size=None,
+#     shuffle=False,
+#     num_workers=num_workers,
+#     pin_memory=True,
+# )
+
+# ## Train ##
+# assert "HCP" in datasets_to_include
+# train_dataset = create_hcp_flat(root=hcp_flat_path, 
+#                 clip_mode="event", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'train')
+# train_dl = wds.WebLoader(
+#     train_dataset.batched(batch_size, partial=False, collation_fn=default_collate),
+#     batch_size=None,
+#     shuffle=False,
+#     num_workers=num_workers,
+#     pin_memory=True,
+# )
+
+# def flatten_meta(meta_dict):
+#     """
+#     Flatten the meta dictionary by:
+#     - Replacing single-item lists with the item itself.
+#     - Converting tensors to scalar numbers.
+#     """
+#     flattened = {}
+#     for key, value in meta_dict.items():
+#         if isinstance(value, list):
+#             if len(value) == 1:
+#                 flattened[key] = value[0]  # Replace list with its single item
+#             else:
+#                 flattened[key] = value  # Keep as is if multiple items
+#         elif isinstance(value, torch.Tensor):
+#             # Convert tensor to scalar
+#             if value.numel() == 1:
+#                 flattened[key] = value.item()
+#             else:
+#                 flattened[key] = value.tolist()  # Convert multi-element tensor to list
+#         else:
+#             flattened[key] = value  # Keep the value as is
+#     return flattened
+
+# import h5py
+# meta_array = np.array([], dtype=object)
+# # Open an HDF5 file in write mode
+# with h5py.File('train_hcp.hdf5', 'w') as h5f:
+#     flatmaps_dset = None
+    
+#     total_samples = 0
+
+#     for i, batch in tqdm(enumerate(train_dl), total = 120000):
+#         images = batch['image'][0]
+#         meta = batch['meta']
+#         batch_size = images.shape[0]
+#         meta_serializable = meta.copy()
+        
+        
+#         # Step 2: Serialize the dictionary to a JSON string
+#         meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)
+#         meta_array = np.append(meta_array, meta_str)
+#         if flatmaps_dset is None:
+#             # Initialize datasets with unlimited (None) maxshape along the first axis
+#             flatmaps_shape = (0,) + images.shape[1:]
+#             flatmaps_maxshape = (None,) + images.shape[1:]
+
+#             flatmaps_dset = h5f.create_dataset(
+#                 'flatmaps',
+#                 shape=flatmaps_shape,
+#                 maxshape=flatmaps_maxshape,
+#                 dtype=np.float16,
+#                 chunks=True  # Enable chunking for efficient resizing
+#             )
+
+#         # Resize datasets to accommodate new data
+#         flatmaps_dset.resize(total_samples + batch_size, axis=0)
+
+#         # Write data to the datasets
+#         flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)
+
+#         total_samples += batch_size
+        
+#     print(f"Processed {total_samples} samples")
+# np.save('metadata_test_HCP.npy', meta_array)
+
+
+# import h5py
+# meta_array = np.array([], dtype=object)
+# # Open an HDF5 file in write mode
+# with h5py.File('test_hcp.hdf5', 'w') as h5f:
+#     flatmaps_dset = None
+    
+#     total_samples = 0
+
+#     for i, batch in tqdm(enumerate(test_dl), total = 12000):
+#         images = batch['image'][0]
+#         meta = batch['meta']
+#         batch_size = images.shape[0]
+#         meta_serializable = meta.copy()
+        
+        
+#         # Step 2: Serialize the dictionary to a JSON string
+#         meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)
+#         meta_array = np.append(meta_array, meta_str)
+#         if flatmaps_dset is None:
+#             # Initialize datasets with unlimited (None) maxshape along the first axis
+#             flatmaps_shape = (0,) + images.shape[1:]
+#             flatmaps_maxshape = (None,) + images.shape[1:]
+
+#             flatmaps_dset = h5f.create_dataset(
+#                 'flatmaps',
+#                 shape=flatmaps_shape,
+#                 maxshape=flatmaps_maxshape,
+#                 dtype=np.float16,
+#                 chunks=True  # Enable chunking for efficient resizing
+#             )
+
+#         # Resize datasets to accommodate new data
+#         flatmaps_dset.resize(total_samples + batch_size, axis=0)
+
+#         # Write data to the datasets
+#         flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)
+
+#         total_samples += batch_size
+        
+#     print(f"Processed {total_samples} samples")
+# np.save('metadata_train_HCP.npy', meta_array)
+
+
+# ### Preparing data
+
+# In[4]:
+
+
+from sklearn.preprocessing import LabelEncoder
+
+INCLUDE_CONDS = {
+    "fear",
+    "neut",
+    "math",
+    "story",
+    "lf",
+    "lh",
+    "rf",
+    "rh",
+    "t",
+    "match",
+    "relation",
+    "mental",
+    "rnd",
+    "0bk_body",
+    "2bk_body",
+    "0bk_faces",
+    "2bk_faces",
+    "0bk_places",
+    "2bk_places",
+    "0bk_tools",
+    "2bk_tools",
+}
+
+# test_data = []
+
+# # Iterate over the DataLoader with a progress bar
+# for sample in tqdm(train_dl, desc="Processing samples"):
+#     x = sample['image']
+#     y = sample['meta']['trial_type']
+#     key = sample['meta']['key']
+#     print(x.shape, y, key)
+#     break
+# Initialize the label encoder
+label_encoder = LabelEncoder()
+label_encoder.fit(sorted(INCLUDE_CONDS))  # Ensure consistent ordering
+
+num_classes = len(label_encoder.classes_)
+print(f"Number of classes: {num_classes}")
+
+
+# In[5]:
+
+
+f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')
+flatmaps_train = f_train['flatmaps']
+
+f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.hdf5', 'r')
+flatmaps_test = f_test['flatmaps']
+
+metadata_train = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_train_HCP.npy', allow_pickle=True)
+metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)
+
+
+# In[6]:
+
+
+from torch.utils.data import Dataset, DataLoader
+
+class HCPFlatDataset(Dataset):
+    def __init__(self, flatmaps, metadata):
+        self.flatmaps = flatmaps
+        self.metadata = metadata
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, idx):
+        return self.flatmaps[idx], json.loads(self.metadata[idx])
+print("Moving datasets to ram")
+# Loading to cpu for faster training, this can take several minutes.  Remove this [:] if you want to move one at the time.
+train_dataset = HCPFlatDataset(flatmaps_train, metadata_train)
+train_dl = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+
+test_dataset = HCPFlatDataset(flatmaps_test, metadata_test)
+test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+print("Datasets ready")
+
+
+# ### Creating and loading Model
+
+# In[7]:
+
+
+from mae_utils.flat import load_hcp_flat_mask
+from mae_utils.flat import create_hcp_flat
+from mae_utils.flat import batch_unmask
+import mae_utils.visualize as vis
+
+flat_mask = load_hcp_flat_mask(hcp_flat_path)
+
+mae_model = flat_models.mae_vit_large_fmri(
+    patch_size=patch_size,
+    decoder_embed_dim=decoder_embed_dim,
+    t_patch_size=t_patch_size,
+    pred_t_dim=pred_t_dim,
+    decoder_depth=4,
+    cls_embed=cls_embed,
+    norm_pix_loss=norm_pix_loss,
+    no_qkv_bias=no_qkv_bias,
+    sep_pos_embed=sep_pos_embed,
+    trunc_init=trunc_init,
+    pct_masks_to_decode=pct_masks_to_decode,
+    img_mask=flat_mask,
+)
+
+
+# In[8]:
+
+
+checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]
+
+if utils.is_interactive():
+    latest_checkpoint = "epoch99.pth"
+else:
+    latest_checkpoint = sys.argv[2] 
+print(f"latest_checkpoint: {latest_checkpoint}")
+
+# Load the checkpoint
+checkpoint_path = os.path.join(outdir, latest_checkpoint)
+
+state = torch.load(checkpoint_path)
+mae_model.load_state_dict(state["model_state_dict"], strict=False)
+mae_model.to(device)
+
+print(f"\nLoaded checkpoint {latest_checkpoint} from {outdir}\n")
+
+
+# In[9]:
+
+
+class LinearClassifier(nn.Module):
+    def __init__(self, input_dim, num_classes):
+        super(LinearClassifier, self).__init__()
+        self.linear = nn.Linear(input_dim, num_classes)
+        
+    def forward(self, x):
+        # Flatten the input except for the batch dimension
+        x = x.view(x.size(0), -1)
+        out = self.linear(x)
+        return out  # Raw logits
+
+# Determine the input dimension from a single sample
+# Assuming images are of shape [1, 16, 144, 320]
+input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])
+print(f"Input dimension: {input_dim}")
+
+
+# In[10]:
+
+
+class FullModel(nn.Module):
+    def __init__(self, lc_model, mae_model):
+        super(FullModel, self).__init__()
+        self.lc_model = lc_model
+        self.mae_model = mae_model
+        
+        
+    def forward(self, x, gsr):
+        x = self.mae_model(x, global_pool=global_pool, forward_features = True)
+        x = self.lc_model(x)
+        return x
+
+
+# In[11]:
+
+
+# Initialize the model
+lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)
+
+model = FullModel(lc_model, mae_model)
+
+# Move the model to the GPU
+model.to(device)
+
+# Define loss function
+criterion = nn.CrossEntropyLoss()
+
+# Define optimizer with L2 regularization (weight_decay)
+learning_rate = 1e-4
+weight_decay = 1e-5  # Adjust based on your needs
+optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+num_epochs = 20  # Adjust as needed
+
+
+# ### Data
+
+# In[16]:
+
+
+import uuid
+
+myuuid = uuid.uuid4()
+str(myuuid)
+
+
+# In[17]:
+
+
+import wandb
+
+if utils.is_interactive():
+    print("Running in interactive notebook. Disabling W&B and ckpt saving.")
+    wandb_log = True
+    save_ckpt = True
+
+if wandb_log:
+    wandb_project = 'fMRI-foundation-model'
+    wandb_config = {
+        "model_name": model_name+'_HCP_FT',
+        "batch_size": batch_size,
+        "learning_rate": learning_rate,
+        "weight_decay": weight_decay,
+        "num_epochs": num_epochs,
+        "seed": seed,
+    }
+    print("wandb_config:\n", wandb_config)
+    random_id = str(uuid.uuid4())
+    print("wandb_id:", "HCPflat_raw" + f"_{random_id}")
+    wandb.init(
+        id=model_name+'_HCP_FT' + f"_{random_id}",
+        project=wandb_project,
+        name=model_name+'_HCP_FT',
+        config=wandb_config,
+        resume="allow",
+    )
+
+
+# In[13]:
+
+
+for epoch in range(num_epochs):
+    running_train_loss = 0.0
+    correct_train = 0
+    total_train = 0
+    step = 0
+
+    # with torch.amp.autocast(device_type='cuda'):
+    # Training Phase
+    model.train()
+    for  batch in tqdm(train_dl, desc=f"Epoch {epoch+1}/{num_epochs} - Training"):
+        optimizer.zero_grad()
+        images = batch[0].to(device).float().unsqueeze(1) #fix this  # Shape: [batch_size, 1, 16, 144, 320]
+        labels = batch[1]['trial_type']  # List of labels
+        
+        encoded_labels = label_encoder.transform(labels)
+        encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)  # Shape: [batch_size]
+        
+        # Forward pass
+        outputs = model(images, gsr=gsr)  # Shape: [num_train_samples, num_classes]
+        
+        # Compute loss
+        loss = criterion(outputs, encoded_labels)
+        
+        # Backward pass and optimization
+        loss.backward()
+        optimizer.step()
+        
+        # Accumulate loss
+        running_train_loss += loss.item() * images.size(0)
+
+        
+        # Calculate accuracy
+        _, predicted = torch.max(outputs, 1)
+        
+        correct_train += (predicted == encoded_labels).sum().item()
+        total_train += encoded_labels.size(0)
+        
+        step = step + 1
+        if step % 100 == 0:
+            print(f"Step [{step}/{len(train_dl)}] - Training Loss: {loss.item():.4f} - Training Accuracy: {100 * correct_train / total_train:.2f}%")
+        # thth
+
+    epoch_train_loss = running_train_loss / total_train if total_train > 0 else 0.0
+    train_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0
+    
+    # Validation Phase
+    model.eval()
+    running_val_loss = 0.0
+    correct_val = 0
+    total_val = 0
+    
+    with torch.no_grad():
+        for batch in tqdm(test_dl, desc=f"Epoch {epoch+1}/{num_epochs} - Validation"):
+            
+            images = batch[0].to(device).float().unsqueeze(1) #fix this
+            labels = batch[1]['trial_type']
+            
+            # Encode labels to integer indices
+            encoded_labels = label_encoder.transform(labels)
+            encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)
+            
+        
+            # Forward pass
+            outputs = model(images, gsr=gsr)
+            
+            # Compute loss
+            loss = criterion(outputs, encoded_labels)
+            
+            # Accumulate loss
+            running_val_loss += loss.item() * images.size(0)
+            
+            # Calculate accuracy
+            _, predicted = torch.max(outputs, 1)
+            correct_val += (predicted == encoded_labels).sum().item()
+            total_val += encoded_labels.size(0)
+
+            
+    
+    epoch_val_loss = running_val_loss / total_val if total_val > 0 else 0.0
+    val_accuracy = 100 * correct_val / total_val if total_val > 0 else 0.0
+    
+    print(f"Epoch [{epoch+1}/{num_epochs}] "
+        f"- Training Loss: {epoch_train_loss:.4f}, Training Accuracy: {train_accuracy:.2f}% "
+        f"- Validation Loss: {epoch_val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%")
+    
+    if wandb_log:
+        wandb.log({
+            "epoch_train_loss": epoch_train_loss,
+            "epoch_val_loss": epoch_val_loss,
+            "train_accuracy": train_accuracy,
+            "val_accuracy": val_accuracy,
+        })
+    if save_ckpt:
+        outdir = os.path.abspath(f'checkpoints/{model_name+"HCP_FT"}')
+        os.makedirs(outdir, exist_ok=True)
+        print("outdir", outdir)
+        # Save model and config
+        torch.save(model.state_dict(), f"{outdir}/model.pth")
+        with open(f"{outdir}/config.yaml", 'w') as f:
+            yaml.dump(wandb_config, f)
+        print(f"Saved model and config to {outdir}")
+    
+

fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/files/output.log

@@ -0,0 +1,59 @@
+Epoch 1/20 - Training:  23%|██▎       | 3199/13913 [18:32<1:01:49,  2.89it/s]
+Step [100/13913] - Training Loss: 1.9649 - Training Accuracy: 59.38%
+Step [200/13913] - Training Loss: 1.0904 - Training Accuracy: 71.19%
+Step [300/13913] - Training Loss: 0.0775 - Training Accuracy: 75.50%
+Step [400/13913] - Training Loss: 0.6052 - Training Accuracy: 78.84%
+Step [500/13913] - Training Loss: 0.0226 - Training Accuracy: 80.95%
+Step [600/13913] - Training Loss: 0.2728 - Training Accuracy: 82.54%
+Step [700/13913] - Training Loss: 0.1662 - Training Accuracy: 83.70%
+Step [800/13913] - Training Loss: 0.0385 - Training Accuracy: 84.89%
+Step [900/13913] - Training Loss: 0.2377 - Training Accuracy: 85.61%
+Step [1000/13913] - Training Loss: 0.8172 - Training Accuracy: 85.83%
+Step [1100/13913] - Training Loss: 0.2276 - Training Accuracy: 86.68%
+Step [1200/13913] - Training Loss: 0.0118 - Training Accuracy: 87.36%
+Step [1300/13913] - Training Loss: 1.0419 - Training Accuracy: 87.68%
+Step [1400/13913] - Training Loss: 0.8943 - Training Accuracy: 87.96%
+Step [1500/13913] - Training Loss: 0.2801 - Training Accuracy: 88.23%
+Step [1600/13913] - Training Loss: 0.6734 - Training Accuracy: 88.58%
+Step [1700/13913] - Training Loss: 0.6202 - Training Accuracy: 88.88%
+Step [1800/13913] - Training Loss: 0.0159 - Training Accuracy: 89.12%
+Step [1900/13913] - Training Loss: 0.0682 - Training Accuracy: 89.37%
+Step [2000/13913] - Training Loss: 0.3378 - Training Accuracy: 89.52%
+Step [2100/13913] - Training Loss: 0.0509 - Training Accuracy: 89.77%
+Step [2200/13913] - Training Loss: 0.1161 - Training Accuracy: 89.99%
+Step [2300/13913] - Training Loss: 0.0025 - Training Accuracy: 90.12%
+Step [2400/13913] - Training Loss: 0.5385 - Training Accuracy: 90.25%
+Step [2500/13913] - Training Loss: 0.0003 - Training Accuracy: 90.47%
+Step [2600/13913] - Training Loss: 0.0962 - Training Accuracy: 90.52%
+Step [2700/13913] - Training Loss: 0.0480 - Training Accuracy: 90.62%
+Step [2800/13913] - Training Loss: 0.0004 - Training Accuracy: 90.71%
+Step [2900/13913] - Training Loss: 0.3049 - Training Accuracy: 90.84%
+Step [3000/13913] - Training Loss: 0.0339 - Training Accuracy: 90.95%
+Step [3100/13913] - Training Loss: 0.5572 - Training Accuracy: 91.02%
+Step [3200/13913] - Training Loss: 0.0673 - Training Accuracy: 91.11%
+Step [3300/13913] - Training Loss: 0.0018 - Training Accuracy: 91.21%
+Step [3400/13913] - Training Loss: 0.0048 - Training Accuracy: 91.28%
+Step [3500/13913] - Training Loss: 1.2120 - Training Accuracy: 91.33%
+Step [3600/13913] - Training Loss: 0.0542 - Training Accuracy: 91.37%
+Step [3700/13913] - Training Loss: 0.0016 - Training Accuracy: 91.43%
+Step [3800/13913] - Training Loss: 0.0582 - Training Accuracy: 91.52%
+Step [3900/13913] - Training Loss: 0.0960 - Training Accuracy: 91.60%
+Step [4000/13913] - Training Loss: 0.0020 - Training Accuracy: 91.68%
+Step [4100/13913] - Training Loss: 0.0043 - Training Accuracy: 91.76%
+Step [4200/13913] - Training Loss: 0.0029 - Training Accuracy: 91.77%
+Step [4300/13913] - Training Loss: 0.0107 - Training Accuracy: 91.83%
+Step [4400/13913] - Training Loss: 0.1122 - Training Accuracy: 91.86%
+Step [4500/13913] - Training Loss: 0.1595 - Training Accuracy: 91.92%
+Step [4600/13913] - Training Loss: 0.0453 - Training Accuracy: 91.97%
+Step [4700/13913] - Training Loss: 0.2770 - Training Accuracy: 92.05%
+Step [4800/13913] - Training Loss: 0.0057 - Training Accuracy: 92.09%
+Step [4900/13913] - Training Loss: 0.0120 - Training Accuracy: 92.16%
+Step [5000/13913] - Training Loss: 0.0235 - Training Accuracy: 92.24%
+Step [5100/13913] - Training Loss: 0.3907 - Training Accuracy: 92.32%
+Step [5200/13913] - Training Loss: 0.4558 - Training Accuracy: 92.34%
+Step [5300/13913] - Training Loss: 0.0051 - Training Accuracy: 92.40%
+Step [5400/13913] - Training Loss: 0.0017 - Training Accuracy: 92.46%
+Step [5500/13913] - Training Loss: 1.3554 - Training Accuracy: 92.50%
+Step [5600/13913] - Training Loss: 0.0617 - Training Accuracy: 92.55%
+Step [5700/13913] - Training Loss: 0.2618 - Training Accuracy: 92.60%
+Step [5800/13913] - Training Loss: 0.0192 - Training Accuracy: 92.60%

fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/files/requirements.txt

@@ -0,0 +1,198 @@
+protobuf==5.28.2
+imageio==2.35.1
+MarkupSafe==3.0.0
+regex==2024.9.11
+matplotlib==3.9.2
+notebook==7.2.2
+debugpy==1.8.6
+aiosignal==1.3.1
+jupyter_core==5.7.2
+torchaudio==2.4.1+cu121
+python-json-logger==2.0.7
+six==1.16.0
+scikit-image==0.24.0
+types-python-dateutil==2.9.0.20241003
+PyYAML==6.0.2
+httpcore==1.0.6
+clip==1.0
+babel==2.16.0
+webcolors==24.8.0
+omegaconf==2.3.0
+webencodings==0.5.1
+kiwisolver==1.4.7
+uri-template==1.3.0
+diffusers==0.23.0
+idna==3.10
+fsspec==2024.9.0
+parso==0.8.4
+setuptools==65.5.0
+tornado==6.4.1
+webdataset==0.2.100
+decord==0.6.0
+nvidia-curand-cu12==10.3.2.106
+ipykernel==6.29.5
+jupyter==1.1.1
+pexpect==4.9.0
+kornia_rs==0.1.5
+iopath==0.1.10
+async-lru==2.0.4
+future==1.0.0
+torchvision==0.19.1+cu121
+botocore==1.34.162
+cycler==0.12.1
+tzdata==2024.2
+jupyter_server_terminals==0.5.3
+click==8.1.7
+einops==0.8.0
+pyzmq==26.2.0
+jupyter_client==8.6.3
+nbconvert==7.16.4
+scikit-learn==1.5.2
+executing==2.1.0
+asttokens==2.4.1
+docker-pycreds==0.4.0
+matplotlib-inline==0.1.7
+overrides==7.7.0
+websocket-client==1.8.0
+nbformat==5.10.4
+elbow==0.1.1
+contourpy==1.3.0
+nvidia-cudnn-cu12==9.1.0.70
+transformers==4.44.2
+gitdb==4.0.11
+jupyterlab_nvdashboard==0.11.0
+lazy_loader==0.4
+jsonpointer==3.0.0
+notebook_shim==0.2.4
+nvidia-nccl-cu12==2.20.5
+ffmpeg-python==0.2.0
+triton==3.0.0
+mistune==3.0.2
+python-dateutil==2.9.0.post0
+beautifulsoup4==4.12.3
+nbclient==0.10.0
+h5py==3.12.1
+ftfy==6.2.3
+zipp==3.20.2
+ptyprocess==0.7.0
+huggingface-hub==0.25.1
+pytz==2024.2
+jupyterlab_pygments==0.3.0
+nvidia-cublas-cu12==12.1.3.1
+pandocfilters==1.5.1
+Jinja2==3.1.4
+arrow==1.3.0
+rpds-py==0.20.0
+jupyter_server==2.14.2
+simplejson==3.19.3
+networkx==3.3
+packaging==24.1
+traitlets==5.14.3
+pandas==2.2.3
+xformers==0.0.22.post7
+lightning-utilities==0.11.7
+tifffile==2024.9.20
+nvidia-cuda-cupti-cu12==12.1.105
+mpmath==1.3.0
+GitPython==3.1.43
+scipy==1.14.1
+jsonschema==4.23.0
+prompt_toolkit==3.0.48
+s3transfer==0.10.2
+multidict==6.1.0
+bleach==6.1.0
+sentry-sdk==2.15.0
+nibabel==5.2.1
+accelerate==1.0.0
+pyarrow==17.0.0
+threadpoolctl==3.5.0
+attrs==24.2.0
+rfc3986-validator==0.1.1
+nvidia-cuda-runtime-cu12==12.1.105
+ipywidgets==8.1.5
+frozenlist==1.4.1
+pycparser==2.22
+jupyterlab_server==2.27.3
+nvidia-cuda-nvrtc-cu12==12.1.105
+yarl==1.13.1
+setproctitle==1.3.3
+isoduration==20.11.0
+Pygments==2.18.0
+jedi==0.19.1
+boto3==1.34.57
+tokenizers==0.19.1
+referencing==0.35.1
+rfc3339-validator==0.1.4
+pillow==10.4.0
+jupyterlab==4.2.5
+stack-data==0.6.3
+h11==0.14.0
+anyio==4.6.0
+nilearn==0.10.4
+nvidia-cusolver-cu12==11.4.5.107
+tinycss2==1.3.0
+defusedxml==0.7.1
+argon2-cffi-bindings==21.2.0
+soupsieve==2.6
+nest-asyncio==1.6.0
+torchmetrics==1.3.0.post0
+tqdm==4.66.5
+cffi==1.17.1
+charset-normalizer==3.3.2
+jsonschema-specifications==2023.12.1
+decorator==5.1.1
+open_clip_torch==2.26.1
+jupyter-events==0.10.0
+smart-open==7.0.5
+antlr4-python3-runtime==4.9.3
+prometheus_client==0.21.0
+kornia==0.7.3
+typing_extensions==4.12.2
+sniffio==1.3.1
+joblib==1.4.2
+comm==0.2.2
+aiohappyeyeballs==2.4.3
+numpy==2.1.2
+braceexpand==0.1.7
+certifi==2024.8.30
+psutil==6.0.0
+pyparsing==3.1.4
+pure_eval==0.2.3
+nvidia-cusparse-cu12==12.1.0.106
+wandb==0.18.3
+urllib3==2.2.3
+smmap==5.0.1
+platformdirs==4.3.6
+torch==2.4.1+cu121
+requests==2.32.3
+json5==0.9.25
+nvidia-nvjitlink-cu12==12.6.77
+jupyterlab_widgets==3.0.13
+lxml==5.3.0
+httpx==0.27.2
+opencv-python==4.6.0.66
+portalocker==2.10.1
+pytorch-lightning==2.0.1
+sympy==1.13.3
+wcwidth==0.2.13
+jmespath==1.0.1
+fqdn==1.5.1
+pynvml==11.5.3
+pip==24.0
+wrapt==1.16.0
+aiohttp==3.10.9
+filelock==3.16.1
+fonttools==4.54.1
+fastjsonschema==2.20.0
+jupyter-console==6.6.3
+widgetsnbextension==4.0.13
+timm==1.0.9
+nvidia-cufft-cu12==11.0.2.54
+ipython==8.28.0
+nvidia-nvtx-cu12==12.1.105
+jupyter-lsp==2.2.5
+safetensors==0.4.5
+terminado==0.18.1
+argon2-cffi==23.1.0
+Send2Trash==1.8.3
+importlib_metadata==8.5.0

fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/files/wandb-metadata.json

@@ -0,0 +1,131 @@
+{
+  "os":  "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+  "python":  "3.11.10",
+  "startedAt":  "2024-10-23T04:13:25.994897Z",
+  "args":  [
+    "HCPflat_large_gsrFalse_",
+    "epoch99.pth"
+  ],
+  "program":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py",
+  "codePath":  "src/HCP_downstream_finetune.py",
+  "git":  {
+    "remote":  "https://github.com/MedARC-AI/fMRI-foundation-model",
+    "commit":  "b1ba684ae7a5cc4155cc046b0abe613de09bf700"
+  },
+  "email":  "torrico.villanueva.cesar.kadir@gmail.com",
+  "root":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+  "host":  "ip-10-0-139-117",
+  "username":  "ckadirt",
+  "executable":  "/admin/home-ckadirt/foundation_env/bin/python",
+  "codePathLocal":  "HCP_downstream_finetune.py",
+  "cpu_count":  96,
+  "cpu_count_logical":  192,
+  "gpu":  "[NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3]",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "249555763200",
+      "used":  "181820596224"
+    }
+  },
+  "memory":  {
+    "total":  "2147443380224"
+  },
+  "cpu":  {
+    "count":  96,
+    "countLogical":  192
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    }
+  ],
+  "slurm":  {
+    "cluster_name":  "sagemaker2",
+    "conf":  "/opt/slurm/etc/slurm.conf",
+    "cpus_on_node":  "20",
+    "gpus_on_node":  "1",
+    "gpus_per_task":  "1",
+    "gtids":  "0",
+    "job_account":  "fmri",
+    "job_cpus_per_node":  "20",
+    "job_end_time":  "1729699988",
+    "job_gid":  "1879800513",
+    "job_gpus":  "4",
+    "job_id":  "528152",
+    "job_name":  "finetuneHCP",
+    "job_nodelist":  "ip-10-0-139-117",
+    "job_num_nodes":  "1",
+    "job_partition":  "p5",
+    "job_qos":  "idle",
+    "job_start_time":  "1729656788",
+    "job_uid":  "1879804696",
+    "job_user":  "ckadirt",
+    "jobid":  "528152",
+    "localid":  "0",
+    "mem_per_cpu":  "11500",
+    "nnodes":  "1",
+    "node_aliases":  "(null)",
+    "nodeid":  "0",
+    "nodelist":  "ip-10-0-139-117",
+    "nprocs":  "1",
+    "ntasks":  "1",
+    "ntasks_per_node":  "1",
+    "prio_process":  "0",
+    "procid":  "0",
+    "script_context":  "prolog_task",
+    "submit_dir":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+    "submit_host":  "ip-172-17-12-61",
+    "task_pid":  "3335147",
+    "tasks_per_node":  "1",
+    "topology_addr":  "ip-10-0-139-117",
+    "topology_addr_pattern":  "node",
+    "working_cluster":  "sagemaker2:ip-172-17-63-161:6817:9984:109"
+  },
+  "cudaVersion":  "12.2"
+}

fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/logs/debug-core.log

@@ -0,0 +1,7 @@
+{"time":"2024-10-23T04:13:25.195730514Z","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmp2j4o8616/port-3335224.txt","pid":3335224,"debug":false,"disable-analytics":false}
+{"time":"2024-10-23T04:13:25.196029969Z","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false}
+{"time":"2024-10-23T04:13:25.199363081Z","level":"INFO","msg":"Will exit if parent process dies.","ppid":3335224}
+{"time":"2024-10-23T04:13:25.199355901Z","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":44407,"Zone":""}}
+{"time":"2024-10-23T04:13:25.388406125Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:59886"}
+{"time":"2024-10-23T04:13:25.995203792Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532","id":"127.0.0.1:59886"}
+{"time":"2024-10-23T04:13:26.086124707Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532","id":"127.0.0.1:59886"}

fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/logs/debug-internal.log

@@ -0,0 +1,11 @@
+{"time":"2024-10-23T04:13:26.025058265Z","level":"INFO","msg":"using version","core version":"0.18.3"}
+{"time":"2024-10-23T04:13:26.025081085Z","level":"INFO","msg":"created symlink","path":"/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/logs/debug-core.log"}
+{"time":"2024-10-23T04:13:26.027044192Z","level":"ERROR","msg":"dialing: google: could not find default credentials. See https://cloud.google.com/docs/authentication/external/set-up-adc for more information"}
+{"time":"2024-10-23T04:13:26.086056755Z","level":"INFO","msg":"created new stream","id":"HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532"}
+{"time":"2024-10-23T04:13:26.086114916Z","level":"INFO","msg":"stream: started","id":"HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532"}
+{"time":"2024-10-23T04:13:26.086145987Z","level":"INFO","msg":"sender: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532"}}
+{"time":"2024-10-23T04:13:26.086149037Z","level":"INFO","msg":"handler: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532"}}
+{"time":"2024-10-23T04:13:26.086128047Z","level":"INFO","msg":"writer: Do: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532"}}
+{"time":"2024-10-23T04:13:26.629699111Z","level":"INFO","msg":"wandb-core","!BADKEY":null}
+{"time":"2024-10-23T04:13:26.639782248Z","level":"INFO","msg":"Starting system monitor"}
+{"time":"2024-10-23T04:13:26.752242533Z","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}

fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/logs/debug.log

@@ -0,0 +1,25 @@
+2024-10-23 04:13:25,982 INFO    MainThread:3335224 [wandb_setup.py:_flush():79] Current SDK version is 0.18.3
+2024-10-23 04:13:25,982 INFO    MainThread:3335224 [wandb_setup.py:_flush():79] Configure stats pid to 3335224
+2024-10-23 04:13:25,982 INFO    MainThread:3335224 [wandb_setup.py:_flush():79] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-10-23 04:13:25,982 INFO    MainThread:3335224 [wandb_setup.py:_flush():79] Loading settings from /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/settings
+2024-10-23 04:13:25,982 INFO    MainThread:3335224 [wandb_setup.py:_flush():79] Loading settings from environment variables: {}
+2024-10-23 04:13:25,982 INFO    MainThread:3335224 [wandb_setup.py:_flush():79] Applying setup settings: {'mode': None, '_disable_service': None}
+2024-10-23 04:13:25,982 INFO    MainThread:3335224 [wandb_setup.py:_flush():79] Inferring run settings from compute environment: {'program_relpath': 'src/HCP_downstream_finetune.py', 'program_abspath': '/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py', 'program': '/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py'}
+2024-10-23 04:13:25,982 INFO    MainThread:3335224 [wandb_setup.py:_flush():79] Applying login settings: {}
+2024-10-23 04:13:25,983 INFO    MainThread:3335224 [wandb_init.py:_log_setup():532] Logging user logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/logs/debug.log
+2024-10-23 04:13:25,984 INFO    MainThread:3335224 [wandb_init.py:_log_setup():533] Logging internal logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/logs/debug-internal.log
+2024-10-23 04:13:25,984 INFO    MainThread:3335224 [wandb_init.py:init():617] calling init triggers
+2024-10-23 04:13:25,984 INFO    MainThread:3335224 [wandb_init.py:init():624] wandb.init called with sweep_config: {}
+config: {'model_name': 'HCPflat_large_gsrFalse__HCP_FT', 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 1e-05, 'num_epochs': 20, 'seed': 42}
+2024-10-23 04:13:25,984 INFO    MainThread:3335224 [wandb_init.py:init():667] starting backend
+2024-10-23 04:13:25,984 INFO    MainThread:3335224 [wandb_init.py:init():671] sending inform_init request
+2024-10-23 04:13:25,994 INFO    MainThread:3335224 [backend.py:_multiprocessing_setup():104] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-10-23 04:13:25,994 INFO    MainThread:3335224 [wandb_init.py:init():684] backend started and connected
+2024-10-23 04:13:26,013 INFO    MainThread:3335224 [wandb_init.py:init():779] updated telemetry
+2024-10-23 04:13:26,066 INFO    MainThread:3335224 [wandb_init.py:init():812] communicating run to backend with 90.0 second timeout
+2024-10-23 04:13:26,593 INFO    MainThread:3335224 [wandb_init.py:init():863] starting run threads in backend
+2024-10-23 04:13:27,015 INFO    MainThread:3335224 [wandb_run.py:_console_start():2465] atexit reg
+2024-10-23 04:13:27,015 INFO    MainThread:3335224 [wandb_run.py:_redirect():2313] redirect: wrap_raw
+2024-10-23 04:13:27,015 INFO    MainThread:3335224 [wandb_run.py:_redirect():2378] Wrapping output streams.
+2024-10-23 04:13:27,015 INFO    MainThread:3335224 [wandb_run.py:_redirect():2403] Redirects installed.
+2024-10-23 04:13:27,018 INFO    MainThread:3335224 [wandb_init.py:init():907] run started, returning control to user process

fMRI-foundation-model/src/wandb/run-20241023_041325-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532/run-HCPflat_large_gsrFalse__HCP_FT_62df940a-b42c-469f-b82e-ec86df778532.wandb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0adae22e19034c2b45367dc668038ee90f898ca7caf78be9bc50b8504115501
+size 3506176

fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/files/code/src/HCP_downstream_finetune.py

@@ -0,0 +1,597 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# Import packages and setup gpu configuration.
+# This code block shouldnt need to be adjusted!
+import os
+import sys
+import json
+import yaml
+import numpy as np
+import copy
+import math
+import time
+import random
+from tqdm.auto import tqdm
+import webdataset as wds
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import utils
+from mae_utils.flat_models import *
+import h5py
+from mae_utils import flat_models
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+# following fixes a Conv3D CUDNN_NOT_SUPPORTED error
+torch.backends.cudnn.benchmark = True
+
+# ## MODEL TO LOAD ##
+if utils.is_interactive():
+    model_name = "HCPflat_large_gsrFalse_"
+else:
+    model_name = sys.argv[1]
+    
+
+# outdir = os.path.abspath(f'checkpoints/{model_name}')
+outdir = os.path.abspath(f'checkpoints/{model_name}')
+
+print("outdir", outdir)
+# Load previous config.yaml if available
+if os.path.exists(f"{outdir}/config.yaml"):
+    config = yaml.load(open(f"{outdir}/config.yaml", 'r'), Loader=yaml.FullLoader)
+    print(f"Loaded config.yaml from ckpt folder {outdir}")
+    # create global variables from the config
+    print("\n__CONFIG__")
+    for attribute_name in config.keys():
+        print(f"{attribute_name} = {config[attribute_name]}")
+        globals()[attribute_name] = config[f'{attribute_name}']
+    print("\n")
+
+world_size = os.getenv('WORLD_SIZE')
+if world_size is None: 
+    world_size = 1
+else:
+    world_size = int(world_size)
+print(f"WORLD_SIZE={world_size}")
+
+if utils.is_interactive():
+    # Following allows you to change functions in models.py or utils.py and 
+    # have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    get_ipython().run_line_magic('autoreload', '2')
+
+batch_size = probe_batch_size
+num_epochs = probe_num_epochs
+
+data_type = torch.float32 # change depending on your mixed_precision
+global_batch_size = batch_size * world_size
+
+device = torch.device('cuda')
+
+hcp_flat_path = "/weka/proj-medarc/shared/HCP-Flat"
+# seed = 42
+# num_frames = 16
+# gsr = False
+# num_workers = 10
+# batch_size = 128
+save_ckpt = True
+wandb_log = True
+print("PID of this process =",os.getpid())
+utils.seed_everything(seed)
+
+
+# In[2]:
+
+
+if os.getenv('global_pool') == "False":
+    global_pool = False
+else:
+    global_pool = True
+print(f"global_pool = {global_pool}")
+
+try:
+    gsr
+except:
+    gsr = True
+    print("set gsr to True")
+print(f"gsr = {gsr}")
+
+
+# In[3]:
+
+
+#### UNCOMMENT THIS TO SAVE THE HCP-FLAT IN HDF5 FORMAT
+
+
+# from torch.utils.data import default_collate
+# from mae_utils.flat import load_hcp_flat_mask
+# from mae_utils.flat import create_hcp_flat
+# from mae_utils.flat import batch_unmask
+# import mae_utils.visualize as vis
+
+
+# batch_size = 26
+# print(f"changed batch_size to {batch_size}")
+
+# ## Test ##
+# datasets_to_include = "HCP"
+# assert "HCP" in datasets_to_include
+# test_dataset = create_hcp_flat(root=hcp_flat_path, 
+#                 clip_mode="event", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'test')
+# test_dl = wds.WebLoader(
+#     test_dataset.batched(batch_size, partial=False, collation_fn=default_collate),
+#     batch_size=None,
+#     shuffle=False,
+#     num_workers=num_workers,
+#     pin_memory=True,
+# )
+
+# ## Train ##
+# assert "HCP" in datasets_to_include
+# train_dataset = create_hcp_flat(root=hcp_flat_path, 
+#                 clip_mode="event", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'train')
+# train_dl = wds.WebLoader(
+#     train_dataset.batched(batch_size, partial=False, collation_fn=default_collate),
+#     batch_size=None,
+#     shuffle=False,
+#     num_workers=num_workers,
+#     pin_memory=True,
+# )
+
+# def flatten_meta(meta_dict):
+#     """
+#     Flatten the meta dictionary by:
+#     - Replacing single-item lists with the item itself.
+#     - Converting tensors to scalar numbers.
+#     """
+#     flattened = {}
+#     for key, value in meta_dict.items():
+#         if isinstance(value, list):
+#             if len(value) == 1:
+#                 flattened[key] = value[0]  # Replace list with its single item
+#             else:
+#                 flattened[key] = value  # Keep as is if multiple items
+#         elif isinstance(value, torch.Tensor):
+#             # Convert tensor to scalar
+#             if value.numel() == 1:
+#                 flattened[key] = value.item()
+#             else:
+#                 flattened[key] = value.tolist()  # Convert multi-element tensor to list
+#         else:
+#             flattened[key] = value  # Keep the value as is
+#     return flattened
+
+# import h5py
+# meta_array = np.array([], dtype=object)
+# # Open an HDF5 file in write mode
+# with h5py.File('train_hcp.hdf5', 'w') as h5f:
+#     flatmaps_dset = None
+    
+#     total_samples = 0
+
+#     for i, batch in tqdm(enumerate(train_dl), total = 120000):
+#         images = batch['image'][0]
+#         meta = batch['meta']
+#         batch_size = images.shape[0]
+#         meta_serializable = meta.copy()
+        
+        
+#         # Step 2: Serialize the dictionary to a JSON string
+#         meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)
+#         meta_array = np.append(meta_array, meta_str)
+#         if flatmaps_dset is None:
+#             # Initialize datasets with unlimited (None) maxshape along the first axis
+#             flatmaps_shape = (0,) + images.shape[1:]
+#             flatmaps_maxshape = (None,) + images.shape[1:]
+
+#             flatmaps_dset = h5f.create_dataset(
+#                 'flatmaps',
+#                 shape=flatmaps_shape,
+#                 maxshape=flatmaps_maxshape,
+#                 dtype=np.float16,
+#                 chunks=True  # Enable chunking for efficient resizing
+#             )
+
+#         # Resize datasets to accommodate new data
+#         flatmaps_dset.resize(total_samples + batch_size, axis=0)
+
+#         # Write data to the datasets
+#         flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)
+
+#         total_samples += batch_size
+        
+#     print(f"Processed {total_samples} samples")
+# np.save('metadata_test_HCP.npy', meta_array)
+
+
+# import h5py
+# meta_array = np.array([], dtype=object)
+# # Open an HDF5 file in write mode
+# with h5py.File('test_hcp.hdf5', 'w') as h5f:
+#     flatmaps_dset = None
+    
+#     total_samples = 0
+
+#     for i, batch in tqdm(enumerate(test_dl), total = 12000):
+#         images = batch['image'][0]
+#         meta = batch['meta']
+#         batch_size = images.shape[0]
+#         meta_serializable = meta.copy()
+        
+        
+#         # Step 2: Serialize the dictionary to a JSON string
+#         meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)
+#         meta_array = np.append(meta_array, meta_str)
+#         if flatmaps_dset is None:
+#             # Initialize datasets with unlimited (None) maxshape along the first axis
+#             flatmaps_shape = (0,) + images.shape[1:]
+#             flatmaps_maxshape = (None,) + images.shape[1:]
+
+#             flatmaps_dset = h5f.create_dataset(
+#                 'flatmaps',
+#                 shape=flatmaps_shape,
+#                 maxshape=flatmaps_maxshape,
+#                 dtype=np.float16,
+#                 chunks=True  # Enable chunking for efficient resizing
+#             )
+
+#         # Resize datasets to accommodate new data
+#         flatmaps_dset.resize(total_samples + batch_size, axis=0)
+
+#         # Write data to the datasets
+#         flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)
+
+#         total_samples += batch_size
+        
+#     print(f"Processed {total_samples} samples")
+# np.save('metadata_train_HCP.npy', meta_array)
+
+
+# ### Preparing data
+
+# In[4]:
+
+
+from sklearn.preprocessing import LabelEncoder
+
+INCLUDE_CONDS = {
+    "fear",
+    "neut",
+    "math",
+    "story",
+    "lf",
+    "lh",
+    "rf",
+    "rh",
+    "t",
+    "match",
+    "relation",
+    "mental",
+    "rnd",
+    "0bk_body",
+    "2bk_body",
+    "0bk_faces",
+    "2bk_faces",
+    "0bk_places",
+    "2bk_places",
+    "0bk_tools",
+    "2bk_tools",
+}
+
+# test_data = []
+
+# # Iterate over the DataLoader with a progress bar
+# for sample in tqdm(train_dl, desc="Processing samples"):
+#     x = sample['image']
+#     y = sample['meta']['trial_type']
+#     key = sample['meta']['key']
+#     print(x.shape, y, key)
+#     break
+# Initialize the label encoder
+label_encoder = LabelEncoder()
+label_encoder.fit(sorted(INCLUDE_CONDS))  # Ensure consistent ordering
+
+num_classes = len(label_encoder.classes_)
+print(f"Number of classes: {num_classes}")
+
+
+# In[5]:
+
+
+f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')
+flatmaps_train = f_train['flatmaps']
+
+f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.hdf5', 'r')
+flatmaps_test = f_test['flatmaps']
+
+metadata_train = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_train_HCP.npy', allow_pickle=True)
+metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)
+
+
+# In[6]:
+
+
+from torch.utils.data import Dataset, DataLoader
+
+class HCPFlatDataset(Dataset):
+    def __init__(self, flatmaps, metadata):
+        self.flatmaps = flatmaps
+        self.metadata = metadata
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, idx):
+        return self.flatmaps[idx], json.loads(self.metadata[idx])
+print("Moving datasets to ram")
+# Loading to cpu for faster training, this can take several minutes.  Remove this [:] if you want to move one at the time.
+train_dataset = HCPFlatDataset(flatmaps_train, metadata_train)
+train_dl = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+
+test_dataset = HCPFlatDataset(flatmaps_test, metadata_test)
+test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+print("Datasets ready")
+
+
+# ### Creating and loading Model
+
+# In[7]:
+
+
+from mae_utils.flat import load_hcp_flat_mask
+from mae_utils.flat import create_hcp_flat
+from mae_utils.flat import batch_unmask
+import mae_utils.visualize as vis
+
+flat_mask = load_hcp_flat_mask(hcp_flat_path)
+
+mae_model = flat_models.mae_vit_large_fmri(
+    patch_size=patch_size,
+    decoder_embed_dim=decoder_embed_dim,
+    t_patch_size=t_patch_size,
+    pred_t_dim=pred_t_dim,
+    decoder_depth=4,
+    cls_embed=cls_embed,
+    norm_pix_loss=norm_pix_loss,
+    no_qkv_bias=no_qkv_bias,
+    sep_pos_embed=sep_pos_embed,
+    trunc_init=trunc_init,
+    pct_masks_to_decode=pct_masks_to_decode,
+    img_mask=flat_mask,
+)
+
+
+# In[8]:
+
+
+checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]
+
+if utils.is_interactive():
+    latest_checkpoint = "epoch99.pth"
+else:
+    latest_checkpoint = sys.argv[2] 
+print(f"latest_checkpoint: {latest_checkpoint}")
+
+# Load the checkpoint
+checkpoint_path = os.path.join(outdir, latest_checkpoint)
+
+state = torch.load(checkpoint_path)
+mae_model.load_state_dict(state["model_state_dict"], strict=False)
+mae_model.to(device)
+
+print(f"\nLoaded checkpoint {latest_checkpoint} from {outdir}\n")
+
+
+# In[9]:
+
+
+class LinearClassifier(nn.Module):
+    def __init__(self, input_dim, num_classes):
+        super(LinearClassifier, self).__init__()
+        self.linear = nn.Linear(input_dim, num_classes)
+        
+    def forward(self, x):
+        # Flatten the input except for the batch dimension
+        x = x.view(x.size(0), -1)
+        out = self.linear(x)
+        return out  # Raw logits
+
+# Determine the input dimension from a single sample
+# Assuming images are of shape [1, 16, 144, 320]
+input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])
+print(f"Input dimension: {input_dim}")
+
+
+# In[10]:
+
+
+class FullModel(nn.Module):
+    def __init__(self, lc_model, mae_model):
+        super(FullModel, self).__init__()
+        self.lc_model = lc_model
+        self.mae_model = mae_model
+        
+        
+    def forward(self, x, gsr):
+        x = self.mae_model(x, global_pool=global_pool, forward_features = True)
+        x = self.lc_model(x)
+        return x
+
+
+# In[11]:
+
+
+# Initialize the model
+lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)
+
+model = FullModel(lc_model, mae_model)
+
+# Move the model to the GPU
+model.to(device)
+
+# Define loss function
+criterion = nn.CrossEntropyLoss()
+
+# Define optimizer with L2 regularization (weight_decay)
+learning_rate = 1e-4
+weight_decay = 1e-5  # Adjust based on your needs
+optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+num_epochs = 20  # Adjust as needed
+
+
+# ### Data
+
+# In[16]:
+
+
+import uuid
+
+myuuid = uuid.uuid4()
+str(myuuid)
+
+
+# In[17]:
+
+
+import wandb
+
+if utils.is_interactive():
+    print("Running in interactive notebook. Disabling W&B and ckpt saving.")
+    wandb_log = True
+    save_ckpt = True
+
+if wandb_log:
+    wandb_project = 'fMRI-foundation-model'
+    wandb_config = {
+        "model_name": model_name+'_HCP_FT',
+        "batch_size": batch_size,
+        "learning_rate": learning_rate,
+        "weight_decay": weight_decay,
+        "num_epochs": num_epochs,
+        "seed": seed,
+    }
+    print("wandb_config:\n", wandb_config)
+    random_id = str(uuid.uuid4())
+    print("wandb_id:", "HCPflat_raw" + f"_{random_id}")
+    wandb.init(
+        id=model_name+'_HCP_FT' + f"_{random_id}",
+        project=wandb_project,
+        name=model_name+'_HCP_FT',
+        config=wandb_config,
+        resume="allow",
+    )
+
+
+# In[13]:
+
+
+for epoch in range(num_epochs):
+    running_train_loss = 0.0
+    correct_train = 0
+    total_train = 0
+    step = 0
+
+    # with torch.amp.autocast(device_type='cuda'):
+    # Training Phase
+    model.train()
+    for  batch in tqdm(train_dl, desc=f"Epoch {epoch+1}/{num_epochs} - Training"):
+        optimizer.zero_grad()
+        images = batch[0].to(device).float().unsqueeze(1) #fix this  # Shape: [batch_size, 1, 16, 144, 320]
+        labels = batch[1]['trial_type']  # List of labels
+        
+        encoded_labels = label_encoder.transform(labels)
+        encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)  # Shape: [batch_size]
+        
+        # Forward pass
+        outputs = model(images, gsr=gsr)  # Shape: [num_train_samples, num_classes]
+        
+        # Compute loss
+        loss = criterion(outputs, encoded_labels)
+        
+        # Backward pass and optimization
+        loss.backward()
+        optimizer.step()
+        
+        # Accumulate loss
+        running_train_loss += loss.item() * images.size(0)
+
+        
+        # Calculate accuracy
+        _, predicted = torch.max(outputs, 1)
+        
+        correct_train += (predicted == encoded_labels).sum().item()
+        total_train += encoded_labels.size(0)
+        
+        step = step + 1
+        if step % 100 == 0:
+            print(f"Step [{step}/{len(train_dl)}] - Training Loss: {loss.item():.4f} - Training Accuracy: {100 * correct_train / total_train:.2f}%")
+        # thth
+
+    epoch_train_loss = running_train_loss / total_train if total_train > 0 else 0.0
+    train_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0
+    
+    # Validation Phase
+    model.eval()
+    running_val_loss = 0.0
+    correct_val = 0
+    total_val = 0
+    
+    with torch.no_grad():
+        for batch in tqdm(test_dl, desc=f"Epoch {epoch+1}/{num_epochs} - Validation"):
+            
+            images = batch[0].to(device).float().unsqueeze(1) #fix this
+            labels = batch[1]['trial_type']
+            
+            # Encode labels to integer indices
+            encoded_labels = label_encoder.transform(labels)
+            encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)
+            
+        
+            # Forward pass
+            outputs = model(images, gsr=gsr)
+            
+            # Compute loss
+            loss = criterion(outputs, encoded_labels)
+            
+            # Accumulate loss
+            running_val_loss += loss.item() * images.size(0)
+            
+            # Calculate accuracy
+            _, predicted = torch.max(outputs, 1)
+            correct_val += (predicted == encoded_labels).sum().item()
+            total_val += encoded_labels.size(0)
+
+            
+    
+    epoch_val_loss = running_val_loss / total_val if total_val > 0 else 0.0
+    val_accuracy = 100 * correct_val / total_val if total_val > 0 else 0.0
+    
+    print(f"Epoch [{epoch+1}/{num_epochs}] "
+        f"- Training Loss: {epoch_train_loss:.4f}, Training Accuracy: {train_accuracy:.2f}% "
+        f"- Validation Loss: {epoch_val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%")
+    
+    if wandb_log:
+        wandb.log({
+            "epoch_train_loss": epoch_train_loss,
+            "epoch_val_loss": epoch_val_loss,
+            "train_accuracy": train_accuracy,
+            "val_accuracy": val_accuracy,
+        })
+    if save_ckpt:
+        outdir = os.path.abspath(f'checkpoints/{model_name+"HCP_FT"}')
+        os.makedirs(outdir, exist_ok=True)
+        print("outdir", outdir)
+        # Save model and config
+        torch.save(model.state_dict(), f"{outdir}/model.pth")
+        with open(f"{outdir}/config.yaml", 'w') as f:
+            yaml.dump(wandb_config, f)
+        print(f"Saved model and config to {outdir}")
+    
+

fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/files/output.log

@@ -0,0 +1,23 @@
+Epoch 1/20 - Training:  16%|█▌        | 2209/13913 [13:24<1:07:17,  2.90it/s]
+Step [100/13913] - Training Loss: 1.9649 - Training Accuracy: 59.38%
+Step [200/13913] - Training Loss: 1.0904 - Training Accuracy: 71.19%
+Step [300/13913] - Training Loss: 0.0775 - Training Accuracy: 75.50%
+Step [400/13913] - Training Loss: 0.6052 - Training Accuracy: 78.84%
+Step [500/13913] - Training Loss: 0.0226 - Training Accuracy: 80.95%
+Step [600/13913] - Training Loss: 0.2728 - Training Accuracy: 82.54%
+Step [700/13913] - Training Loss: 0.1662 - Training Accuracy: 83.70%
+Step [800/13913] - Training Loss: 0.0385 - Training Accuracy: 84.89%
+Step [900/13913] - Training Loss: 0.2377 - Training Accuracy: 85.61%
+Step [1000/13913] - Training Loss: 0.8172 - Training Accuracy: 85.83%
+Step [1100/13913] - Training Loss: 0.2276 - Training Accuracy: 86.68%
+Step [1200/13913] - Training Loss: 0.0118 - Training Accuracy: 87.36%
+Step [1300/13913] - Training Loss: 1.0419 - Training Accuracy: 87.68%
+Step [1400/13913] - Training Loss: 0.8943 - Training Accuracy: 87.96%
+Step [1500/13913] - Training Loss: 0.2801 - Training Accuracy: 88.23%
+Step [1600/13913] - Training Loss: 0.6734 - Training Accuracy: 88.58%
+Step [1700/13913] - Training Loss: 0.6202 - Training Accuracy: 88.88%
+Step [1800/13913] - Training Loss: 0.0159 - Training Accuracy: 89.12%
+Step [1900/13913] - Training Loss: 0.0682 - Training Accuracy: 89.37%
+Step [2000/13913] - Training Loss: 0.3378 - Training Accuracy: 89.52%
+Step [2100/13913] - Training Loss: 0.0509 - Training Accuracy: 89.77%
+Step [2200/13913] - Training Loss: 0.1161 - Training Accuracy: 89.99%

fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/files/requirements.txt

@@ -0,0 +1,198 @@
+protobuf==5.28.2
+imageio==2.35.1
+MarkupSafe==3.0.0
+regex==2024.9.11
+matplotlib==3.9.2
+notebook==7.2.2
+debugpy==1.8.6
+aiosignal==1.3.1
+jupyter_core==5.7.2
+torchaudio==2.4.1+cu121
+python-json-logger==2.0.7
+six==1.16.0
+scikit-image==0.24.0
+types-python-dateutil==2.9.0.20241003
+PyYAML==6.0.2
+httpcore==1.0.6
+clip==1.0
+babel==2.16.0
+webcolors==24.8.0
+omegaconf==2.3.0
+webencodings==0.5.1
+kiwisolver==1.4.7
+uri-template==1.3.0
+diffusers==0.23.0
+idna==3.10
+fsspec==2024.9.0
+parso==0.8.4
+setuptools==65.5.0
+tornado==6.4.1
+webdataset==0.2.100
+decord==0.6.0
+nvidia-curand-cu12==10.3.2.106
+ipykernel==6.29.5
+jupyter==1.1.1
+pexpect==4.9.0
+kornia_rs==0.1.5
+iopath==0.1.10
+async-lru==2.0.4
+future==1.0.0
+torchvision==0.19.1+cu121
+botocore==1.34.162
+cycler==0.12.1
+tzdata==2024.2
+jupyter_server_terminals==0.5.3
+click==8.1.7
+einops==0.8.0
+pyzmq==26.2.0
+jupyter_client==8.6.3
+nbconvert==7.16.4
+scikit-learn==1.5.2
+executing==2.1.0
+asttokens==2.4.1
+docker-pycreds==0.4.0
+matplotlib-inline==0.1.7
+overrides==7.7.0
+websocket-client==1.8.0
+nbformat==5.10.4
+elbow==0.1.1
+contourpy==1.3.0
+nvidia-cudnn-cu12==9.1.0.70
+transformers==4.44.2
+gitdb==4.0.11
+jupyterlab_nvdashboard==0.11.0
+lazy_loader==0.4
+jsonpointer==3.0.0
+notebook_shim==0.2.4
+nvidia-nccl-cu12==2.20.5
+ffmpeg-python==0.2.0
+triton==3.0.0
+mistune==3.0.2
+python-dateutil==2.9.0.post0
+beautifulsoup4==4.12.3
+nbclient==0.10.0
+h5py==3.12.1
+ftfy==6.2.3
+zipp==3.20.2
+ptyprocess==0.7.0
+huggingface-hub==0.25.1
+pytz==2024.2
+jupyterlab_pygments==0.3.0
+nvidia-cublas-cu12==12.1.3.1
+pandocfilters==1.5.1
+Jinja2==3.1.4
+arrow==1.3.0
+rpds-py==0.20.0
+jupyter_server==2.14.2
+simplejson==3.19.3
+networkx==3.3
+packaging==24.1
+traitlets==5.14.3
+pandas==2.2.3
+xformers==0.0.22.post7
+lightning-utilities==0.11.7
+tifffile==2024.9.20
+nvidia-cuda-cupti-cu12==12.1.105
+mpmath==1.3.0
+GitPython==3.1.43
+scipy==1.14.1
+jsonschema==4.23.0
+prompt_toolkit==3.0.48
+s3transfer==0.10.2
+multidict==6.1.0
+bleach==6.1.0
+sentry-sdk==2.15.0
+nibabel==5.2.1
+accelerate==1.0.0
+pyarrow==17.0.0
+threadpoolctl==3.5.0
+attrs==24.2.0
+rfc3986-validator==0.1.1
+nvidia-cuda-runtime-cu12==12.1.105
+ipywidgets==8.1.5
+frozenlist==1.4.1
+pycparser==2.22
+jupyterlab_server==2.27.3
+nvidia-cuda-nvrtc-cu12==12.1.105
+yarl==1.13.1
+setproctitle==1.3.3
+isoduration==20.11.0
+Pygments==2.18.0
+jedi==0.19.1
+boto3==1.34.57
+tokenizers==0.19.1
+referencing==0.35.1
+rfc3339-validator==0.1.4
+pillow==10.4.0
+jupyterlab==4.2.5
+stack-data==0.6.3
+h11==0.14.0
+anyio==4.6.0
+nilearn==0.10.4
+nvidia-cusolver-cu12==11.4.5.107
+tinycss2==1.3.0
+defusedxml==0.7.1
+argon2-cffi-bindings==21.2.0
+soupsieve==2.6
+nest-asyncio==1.6.0
+torchmetrics==1.3.0.post0
+tqdm==4.66.5
+cffi==1.17.1
+charset-normalizer==3.3.2
+jsonschema-specifications==2023.12.1
+decorator==5.1.1
+open_clip_torch==2.26.1
+jupyter-events==0.10.0
+smart-open==7.0.5
+antlr4-python3-runtime==4.9.3
+prometheus_client==0.21.0
+kornia==0.7.3
+typing_extensions==4.12.2
+sniffio==1.3.1
+joblib==1.4.2
+comm==0.2.2
+aiohappyeyeballs==2.4.3
+numpy==2.1.2
+braceexpand==0.1.7
+certifi==2024.8.30
+psutil==6.0.0
+pyparsing==3.1.4
+pure_eval==0.2.3
+nvidia-cusparse-cu12==12.1.0.106
+wandb==0.18.3
+urllib3==2.2.3
+smmap==5.0.1
+platformdirs==4.3.6
+torch==2.4.1+cu121
+requests==2.32.3
+json5==0.9.25
+nvidia-nvjitlink-cu12==12.6.77
+jupyterlab_widgets==3.0.13
+lxml==5.3.0
+httpx==0.27.2
+opencv-python==4.6.0.66
+portalocker==2.10.1
+pytorch-lightning==2.0.1
+sympy==1.13.3
+wcwidth==0.2.13
+jmespath==1.0.1
+fqdn==1.5.1
+pynvml==11.5.3
+pip==24.0
+wrapt==1.16.0
+aiohttp==3.10.9
+filelock==3.16.1
+fonttools==4.54.1
+fastjsonschema==2.20.0
+jupyter-console==6.6.3
+widgetsnbextension==4.0.13
+timm==1.0.9
+nvidia-cufft-cu12==11.0.2.54
+ipython==8.28.0
+nvidia-nvtx-cu12==12.1.105
+jupyter-lsp==2.2.5
+safetensors==0.4.5
+terminado==0.18.1
+argon2-cffi==23.1.0
+Send2Trash==1.8.3
+importlib_metadata==8.5.0

fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/files/wandb-metadata.json

@@ -0,0 +1,131 @@
+{
+  "os":  "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+  "python":  "3.11.10",
+  "startedAt":  "2024-10-24T02:16:25.309620Z",
+  "args":  [
+    "HCPflat_large_gsrFalse_",
+    "epoch99.pth"
+  ],
+  "program":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py",
+  "codePath":  "src/HCP_downstream_finetune.py",
+  "git":  {
+    "remote":  "https://github.com/MedARC-AI/fMRI-foundation-model",
+    "commit":  "cf8214d4ebe437188b68b4ee5a34c5211a810db0"
+  },
+  "email":  "torrico.villanueva.cesar.kadir@gmail.com",
+  "root":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+  "host":  "ip-10-0-181-106",
+  "username":  "ckadirt",
+  "executable":  "/admin/home-ckadirt/foundation_env/bin/python",
+  "codePathLocal":  "HCP_downstream_finetune.py",
+  "cpu_count":  96,
+  "cpu_count_logical":  192,
+  "gpu":  "[NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3]",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "249555763200",
+      "used":  "184254554112"
+    }
+  },
+  "memory":  {
+    "total":  "2147443404800"
+  },
+  "cpu":  {
+    "count":  96,
+    "countLogical":  192
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    }
+  ],
+  "slurm":  {
+    "cluster_name":  "sagemaker2",
+    "conf":  "/opt/slurm/etc/slurm.conf",
+    "cpus_on_node":  "20",
+    "gpus_on_node":  "1",
+    "gpus_per_task":  "1",
+    "gtids":  "0",
+    "job_account":  "fmri",
+    "job_cpus_per_node":  "20",
+    "job_end_time":  "1729779337",
+    "job_gid":  "1879800513",
+    "job_gpus":  "0",
+    "job_id":  "528653",
+    "job_name":  "finetuneHCP",
+    "job_nodelist":  "ip-10-0-181-106",
+    "job_num_nodes":  "1",
+    "job_partition":  "p5",
+    "job_qos":  "normal",
+    "job_start_time":  "1729736136",
+    "job_uid":  "1879804696",
+    "job_user":  "ckadirt",
+    "jobid":  "528653",
+    "localid":  "0",
+    "mem_per_cpu":  "11500",
+    "nnodes":  "1",
+    "node_aliases":  "(null)",
+    "nodeid":  "0",
+    "nodelist":  "ip-10-0-181-106",
+    "nprocs":  "1",
+    "ntasks":  "1",
+    "ntasks_per_node":  "1",
+    "prio_process":  "0",
+    "procid":  "0",
+    "script_context":  "prolog_task",
+    "submit_dir":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+    "submit_host":  "ip-172-17-12-61",
+    "task_pid":  "2312183",
+    "tasks_per_node":  "1",
+    "topology_addr":  "ip-10-0-181-106",
+    "topology_addr_pattern":  "node",
+    "working_cluster":  "sagemaker2:ip-172-17-63-161:6817:9984:109"
+  },
+  "cudaVersion":  "12.2"
+}

fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/logs/debug-core.log

@@ -0,0 +1,7 @@
+{"time":"2024-10-24T02:16:24.2815064Z","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmpbnequ20j/port-2312212.txt","pid":2312212,"debug":false,"disable-analytics":false}
+{"time":"2024-10-24T02:16:24.281779825Z","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false}
+{"time":"2024-10-24T02:16:24.287252623Z","level":"INFO","msg":"Will exit if parent process dies.","ppid":2312212}
+{"time":"2024-10-24T02:16:24.287238992Z","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":43197,"Zone":""}}
+{"time":"2024-10-24T02:16:24.297514377Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:48772"}
+{"time":"2024-10-24T02:16:25.309558909Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427","id":"127.0.0.1:48772"}
+{"time":"2024-10-24T02:16:25.471977517Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427","id":"127.0.0.1:48772"}

fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/logs/debug-internal.log

@@ -0,0 +1,11 @@
+{"time":"2024-10-24T02:16:25.343111738Z","level":"INFO","msg":"using version","core version":"0.18.3"}
+{"time":"2024-10-24T02:16:25.343149488Z","level":"INFO","msg":"created symlink","path":"/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/logs/debug-core.log"}
+{"time":"2024-10-24T02:16:25.354611723Z","level":"ERROR","msg":"dialing: google: could not find default credentials. See https://cloud.google.com/docs/authentication/external/set-up-adc for more information"}
+{"time":"2024-10-24T02:16:25.471914526Z","level":"INFO","msg":"created new stream","id":"HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427"}
+{"time":"2024-10-24T02:16:25.471970027Z","level":"INFO","msg":"stream: started","id":"HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427"}
+{"time":"2024-10-24T02:16:25.472001617Z","level":"INFO","msg":"handler: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427"}}
+{"time":"2024-10-24T02:16:25.472008128Z","level":"INFO","msg":"writer: Do: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427"}}
+{"time":"2024-10-24T02:16:25.472045188Z","level":"INFO","msg":"sender: started","stream_id":{"value":"HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427"}}
+{"time":"2024-10-24T02:16:26.180235591Z","level":"INFO","msg":"wandb-core","!BADKEY":null}
+{"time":"2024-10-24T02:16:26.191656334Z","level":"INFO","msg":"Starting system monitor"}
+{"time":"2024-10-24T02:16:26.243500377Z","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}

fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/logs/debug.log

@@ -0,0 +1,25 @@
+2024-10-24 02:16:25,238 INFO    MainThread:2312212 [wandb_setup.py:_flush():79] Current SDK version is 0.18.3
+2024-10-24 02:16:25,239 INFO    MainThread:2312212 [wandb_setup.py:_flush():79] Configure stats pid to 2312212
+2024-10-24 02:16:25,239 INFO    MainThread:2312212 [wandb_setup.py:_flush():79] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-10-24 02:16:25,239 INFO    MainThread:2312212 [wandb_setup.py:_flush():79] Loading settings from /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/settings
+2024-10-24 02:16:25,239 INFO    MainThread:2312212 [wandb_setup.py:_flush():79] Loading settings from environment variables: {}
+2024-10-24 02:16:25,239 INFO    MainThread:2312212 [wandb_setup.py:_flush():79] Applying setup settings: {'mode': None, '_disable_service': None}
+2024-10-24 02:16:25,239 INFO    MainThread:2312212 [wandb_setup.py:_flush():79] Inferring run settings from compute environment: {'program_relpath': 'src/HCP_downstream_finetune.py', 'program_abspath': '/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py', 'program': '/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py'}
+2024-10-24 02:16:25,239 INFO    MainThread:2312212 [wandb_setup.py:_flush():79] Applying login settings: {}
+2024-10-24 02:16:25,242 INFO    MainThread:2312212 [wandb_init.py:_log_setup():532] Logging user logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/logs/debug.log
+2024-10-24 02:16:25,246 INFO    MainThread:2312212 [wandb_init.py:_log_setup():533] Logging internal logs to /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/logs/debug-internal.log
+2024-10-24 02:16:25,246 INFO    MainThread:2312212 [wandb_init.py:init():617] calling init triggers
+2024-10-24 02:16:25,246 INFO    MainThread:2312212 [wandb_init.py:init():624] wandb.init called with sweep_config: {}
+config: {'model_name': 'HCPflat_large_gsrFalse__HCP_FT', 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 1e-05, 'num_epochs': 20, 'seed': 42}
+2024-10-24 02:16:25,246 INFO    MainThread:2312212 [wandb_init.py:init():667] starting backend
+2024-10-24 02:16:25,246 INFO    MainThread:2312212 [wandb_init.py:init():671] sending inform_init request
+2024-10-24 02:16:25,305 INFO    MainThread:2312212 [backend.py:_multiprocessing_setup():104] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-10-24 02:16:25,305 INFO    MainThread:2312212 [wandb_init.py:init():684] backend started and connected
+2024-10-24 02:16:25,385 INFO    MainThread:2312212 [wandb_init.py:init():779] updated telemetry
+2024-10-24 02:16:25,618 INFO    MainThread:2312212 [wandb_init.py:init():812] communicating run to backend with 90.0 second timeout
+2024-10-24 02:16:26,156 INFO    MainThread:2312212 [wandb_init.py:init():863] starting run threads in backend
+2024-10-24 02:16:26,723 INFO    MainThread:2312212 [wandb_run.py:_console_start():2465] atexit reg
+2024-10-24 02:16:26,723 INFO    MainThread:2312212 [wandb_run.py:_redirect():2313] redirect: wrap_raw
+2024-10-24 02:16:26,723 INFO    MainThread:2312212 [wandb_run.py:_redirect():2378] Wrapping output streams.
+2024-10-24 02:16:26,723 INFO    MainThread:2312212 [wandb_run.py:_redirect():2403] Redirects installed.
+2024-10-24 02:16:26,731 INFO    MainThread:2312212 [wandb_init.py:init():907] run started, returning control to user process

fMRI-foundation-model/src/wandb/run-20241024_021624-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427/run-HCPflat_large_gsrFalse__HCP_FT_55946e57-be97-4417-b0d9-e535f9bdc427.wandb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fae4667e9bd6c4dd0115f39374a34a0cd6bde90440ba75d82671c57fa7cd83d1
+size 1343488

fMRI-foundation-model/src/wandb/run-20241024_213834-HCPflat_large_gsrFalse__HCP_FT_d07870e8-3c53-420c-b734-f32a1c5c3e5a/files/code/src/HCP_downstream_finetune.py

@@ -0,0 +1,597 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# Import packages and setup gpu configuration.
+# This code block shouldnt need to be adjusted!
+import os
+import sys
+import json
+import yaml
+import numpy as np
+import copy
+import math
+import time
+import random
+from tqdm.auto import tqdm
+import webdataset as wds
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import utils
+from mae_utils.flat_models import *
+import h5py
+from mae_utils import flat_models
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+# following fixes a Conv3D CUDNN_NOT_SUPPORTED error
+torch.backends.cudnn.benchmark = True
+
+# ## MODEL TO LOAD ##
+if utils.is_interactive():
+    model_name = "HCPflat_large_gsrFalse_"
+else:
+    model_name = sys.argv[1]
+    
+
+# outdir = os.path.abspath(f'checkpoints/{model_name}')
+outdir = os.path.abspath(f'checkpoints/{model_name}')
+
+print("outdir", outdir)
+# Load previous config.yaml if available
+if os.path.exists(f"{outdir}/config.yaml"):
+    config = yaml.load(open(f"{outdir}/config.yaml", 'r'), Loader=yaml.FullLoader)
+    print(f"Loaded config.yaml from ckpt folder {outdir}")
+    # create global variables from the config
+    print("\n__CONFIG__")
+    for attribute_name in config.keys():
+        print(f"{attribute_name} = {config[attribute_name]}")
+        globals()[attribute_name] = config[f'{attribute_name}']
+    print("\n")
+
+world_size = os.getenv('WORLD_SIZE')
+if world_size is None: 
+    world_size = 1
+else:
+    world_size = int(world_size)
+print(f"WORLD_SIZE={world_size}")
+
+if utils.is_interactive():
+    # Following allows you to change functions in models.py or utils.py and 
+    # have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    get_ipython().run_line_magic('autoreload', '2')
+
+batch_size = probe_batch_size
+num_epochs = probe_num_epochs
+
+data_type = torch.float32 # change depending on your mixed_precision
+global_batch_size = batch_size * world_size
+
+device = torch.device('cuda')
+
+hcp_flat_path = "/weka/proj-medarc/shared/HCP-Flat"
+# seed = 42
+# num_frames = 16
+# gsr = False
+# num_workers = 10
+# batch_size = 128
+save_ckpt = True
+wandb_log = True
+print("PID of this process =",os.getpid())
+utils.seed_everything(seed)
+
+
+# In[2]:
+
+
+if os.getenv('global_pool') == "False":
+    global_pool = False
+else:
+    global_pool = True
+print(f"global_pool = {global_pool}")
+
+try:
+    gsr
+except:
+    gsr = True
+    print("set gsr to True")
+print(f"gsr = {gsr}")
+
+
+# In[3]:
+
+
+#### UNCOMMENT THIS TO SAVE THE HCP-FLAT IN HDF5 FORMAT
+
+
+# from torch.utils.data import default_collate
+# from mae_utils.flat import load_hcp_flat_mask
+# from mae_utils.flat import create_hcp_flat
+# from mae_utils.flat import batch_unmask
+# import mae_utils.visualize as vis
+
+
+# batch_size = 26
+# print(f"changed batch_size to {batch_size}")
+
+# ## Test ##
+# datasets_to_include = "HCP"
+# assert "HCP" in datasets_to_include
+# test_dataset = create_hcp_flat(root=hcp_flat_path, 
+#                 clip_mode="event", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'test')
+# test_dl = wds.WebLoader(
+#     test_dataset.batched(batch_size, partial=False, collation_fn=default_collate),
+#     batch_size=None,
+#     shuffle=False,
+#     num_workers=num_workers,
+#     pin_memory=True,
+# )
+
+# ## Train ##
+# assert "HCP" in datasets_to_include
+# train_dataset = create_hcp_flat(root=hcp_flat_path, 
+#                 clip_mode="event", frames=num_frames, shuffle=False, gsr=gsr, sub_list = 'train')
+# train_dl = wds.WebLoader(
+#     train_dataset.batched(batch_size, partial=False, collation_fn=default_collate),
+#     batch_size=None,
+#     shuffle=False,
+#     num_workers=num_workers,
+#     pin_memory=True,
+# )
+
+# def flatten_meta(meta_dict):
+#     """
+#     Flatten the meta dictionary by:
+#     - Replacing single-item lists with the item itself.
+#     - Converting tensors to scalar numbers.
+#     """
+#     flattened = {}
+#     for key, value in meta_dict.items():
+#         if isinstance(value, list):
+#             if len(value) == 1:
+#                 flattened[key] = value[0]  # Replace list with its single item
+#             else:
+#                 flattened[key] = value  # Keep as is if multiple items
+#         elif isinstance(value, torch.Tensor):
+#             # Convert tensor to scalar
+#             if value.numel() == 1:
+#                 flattened[key] = value.item()
+#             else:
+#                 flattened[key] = value.tolist()  # Convert multi-element tensor to list
+#         else:
+#             flattened[key] = value  # Keep the value as is
+#     return flattened
+
+# import h5py
+# meta_array = np.array([], dtype=object)
+# # Open an HDF5 file in write mode
+# with h5py.File('train_hcp.hdf5', 'w') as h5f:
+#     flatmaps_dset = None
+    
+#     total_samples = 0
+
+#     for i, batch in tqdm(enumerate(train_dl), total = 120000):
+#         images = batch['image'][0]
+#         meta = batch['meta']
+#         batch_size = images.shape[0]
+#         meta_serializable = meta.copy()
+        
+        
+#         # Step 2: Serialize the dictionary to a JSON string
+#         meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)
+#         meta_array = np.append(meta_array, meta_str)
+#         if flatmaps_dset is None:
+#             # Initialize datasets with unlimited (None) maxshape along the first axis
+#             flatmaps_shape = (0,) + images.shape[1:]
+#             flatmaps_maxshape = (None,) + images.shape[1:]
+
+#             flatmaps_dset = h5f.create_dataset(
+#                 'flatmaps',
+#                 shape=flatmaps_shape,
+#                 maxshape=flatmaps_maxshape,
+#                 dtype=np.float16,
+#                 chunks=True  # Enable chunking for efficient resizing
+#             )
+
+#         # Resize datasets to accommodate new data
+#         flatmaps_dset.resize(total_samples + batch_size, axis=0)
+
+#         # Write data to the datasets
+#         flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)
+
+#         total_samples += batch_size
+        
+#     print(f"Processed {total_samples} samples")
+# np.save('metadata_test_HCP.npy', meta_array)
+
+
+# import h5py
+# meta_array = np.array([], dtype=object)
+# # Open an HDF5 file in write mode
+# with h5py.File('test_hcp.hdf5', 'w') as h5f:
+#     flatmaps_dset = None
+    
+#     total_samples = 0
+
+#     for i, batch in tqdm(enumerate(test_dl), total = 12000):
+#         images = batch['image'][0]
+#         meta = batch['meta']
+#         batch_size = images.shape[0]
+#         meta_serializable = meta.copy()
+        
+        
+#         # Step 2: Serialize the dictionary to a JSON string
+#         meta_str = json.dumps(flatten_meta(meta_serializable), indent=4)
+#         meta_array = np.append(meta_array, meta_str)
+#         if flatmaps_dset is None:
+#             # Initialize datasets with unlimited (None) maxshape along the first axis
+#             flatmaps_shape = (0,) + images.shape[1:]
+#             flatmaps_maxshape = (None,) + images.shape[1:]
+
+#             flatmaps_dset = h5f.create_dataset(
+#                 'flatmaps',
+#                 shape=flatmaps_shape,
+#                 maxshape=flatmaps_maxshape,
+#                 dtype=np.float16,
+#                 chunks=True  # Enable chunking for efficient resizing
+#             )
+
+#         # Resize datasets to accommodate new data
+#         flatmaps_dset.resize(total_samples + batch_size, axis=0)
+
+#         # Write data to the datasets
+#         flatmaps_dset[total_samples:total_samples + batch_size] = images.numpy().astype(np.float16)
+
+#         total_samples += batch_size
+        
+#     print(f"Processed {total_samples} samples")
+# np.save('metadata_train_HCP.npy', meta_array)
+
+
+# ### Preparing data
+
+# In[4]:
+
+
+from sklearn.preprocessing import LabelEncoder
+
+INCLUDE_CONDS = {
+    "fear",
+    "neut",
+    "math",
+    "story",
+    "lf",
+    "lh",
+    "rf",
+    "rh",
+    "t",
+    "match",
+    "relation",
+    "mental",
+    "rnd",
+    "0bk_body",
+    "2bk_body",
+    "0bk_faces",
+    "2bk_faces",
+    "0bk_places",
+    "2bk_places",
+    "0bk_tools",
+    "2bk_tools",
+}
+
+# test_data = []
+
+# # Iterate over the DataLoader with a progress bar
+# for sample in tqdm(train_dl, desc="Processing samples"):
+#     x = sample['image']
+#     y = sample['meta']['trial_type']
+#     key = sample['meta']['key']
+#     print(x.shape, y, key)
+#     break
+# Initialize the label encoder
+label_encoder = LabelEncoder()
+label_encoder.fit(sorted(INCLUDE_CONDS))  # Ensure consistent ordering
+
+num_classes = len(label_encoder.classes_)
+print(f"Number of classes: {num_classes}")
+
+
+# In[5]:
+
+
+f_train = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/train_hcp.hdf5', 'r')
+flatmaps_train = f_train['flatmaps']
+
+f_test = h5py.File('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/test_hcp.hdf5', 'r')
+flatmaps_test = f_test['flatmaps']
+
+metadata_train = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_train_HCP.npy', allow_pickle=True)
+metadata_test = np.load('/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/metadata_test_HCP.npy', allow_pickle=True)
+
+
+# In[6]:
+
+
+from torch.utils.data import Dataset, DataLoader
+
+class HCPFlatDataset(Dataset):
+    def __init__(self, flatmaps, metadata):
+        self.flatmaps = flatmaps
+        self.metadata = metadata
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, idx):
+        return self.flatmaps[idx], json.loads(self.metadata[idx])
+print("Moving datasets to ram")
+# Loading to cpu for faster training, this can take several minutes.  Remove this [:] if you want to move one at the time.
+train_dataset = HCPFlatDataset(flatmaps_train, metadata_train)
+train_dl = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+
+test_dataset = HCPFlatDataset(flatmaps_test, metadata_test)
+test_dl = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+print("Datasets ready")
+
+
+# ### Creating and loading Model
+
+# In[7]:
+
+
+from mae_utils.flat import load_hcp_flat_mask
+from mae_utils.flat import create_hcp_flat
+from mae_utils.flat import batch_unmask
+import mae_utils.visualize as vis
+
+flat_mask = load_hcp_flat_mask(hcp_flat_path)
+
+mae_model = flat_models.mae_vit_large_fmri(
+    patch_size=patch_size,
+    decoder_embed_dim=decoder_embed_dim,
+    t_patch_size=t_patch_size,
+    pred_t_dim=pred_t_dim,
+    decoder_depth=4,
+    cls_embed=cls_embed,
+    norm_pix_loss=norm_pix_loss,
+    no_qkv_bias=no_qkv_bias,
+    sep_pos_embed=sep_pos_embed,
+    trunc_init=trunc_init,
+    pct_masks_to_decode=pct_masks_to_decode,
+    img_mask=flat_mask,
+)
+
+
+# In[8]:
+
+
+checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]
+
+if utils.is_interactive():
+    latest_checkpoint = "epoch99.pth"
+else:
+    latest_checkpoint = sys.argv[2] 
+print(f"latest_checkpoint: {latest_checkpoint}")
+
+# Load the checkpoint
+checkpoint_path = os.path.join(outdir, latest_checkpoint)
+
+state = torch.load(checkpoint_path)
+mae_model.load_state_dict(state["model_state_dict"], strict=False)
+mae_model.to(device)
+
+print(f"\nLoaded checkpoint {latest_checkpoint} from {outdir}\n")
+
+
+# In[9]:
+
+
+class LinearClassifier(nn.Module):
+    def __init__(self, input_dim, num_classes):
+        super(LinearClassifier, self).__init__()
+        self.linear = nn.Linear(input_dim, num_classes)
+        
+    def forward(self, x):
+        # Flatten the input except for the batch dimension
+        x = x.view(x.size(0), -1)
+        out = self.linear(x)
+        return out  # Raw logits
+
+# Determine the input dimension from a single sample
+# Assuming images are of shape [1, 16, 144, 320]
+input_dim = np.prod(mae_model(torch.randn(1,1,16,144,320).to(device),global_pool=global_pool, forward_features = True).shape[1:])
+print(f"Input dimension: {input_dim}")
+
+
+# In[10]:
+
+
+class FullModel(nn.Module):
+    def __init__(self, lc_model, mae_model):
+        super(FullModel, self).__init__()
+        self.lc_model = lc_model
+        self.mae_model = mae_model
+        
+        
+    def forward(self, x, gsr):
+        x = self.mae_model(x, global_pool=global_pool, forward_features = True)
+        x = self.lc_model(x)
+        return x
+
+
+# In[11]:
+
+
+# Initialize the model
+lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)
+
+model = FullModel(lc_model, mae_model)
+
+# Move the model to the GPU
+model.to(device)
+
+# Define loss function
+criterion = nn.CrossEntropyLoss()
+
+# Define optimizer with L2 regularization (weight_decay)
+learning_rate = 1e-4
+weight_decay = 1e-5  # Adjust based on your needs
+optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+num_epochs = 20  # Adjust as needed
+
+
+# ### Data
+
+# In[16]:
+
+
+import uuid
+
+myuuid = uuid.uuid4()
+str(myuuid)
+
+
+# In[17]:
+
+
+import wandb
+
+if utils.is_interactive():
+    print("Running in interactive notebook. Disabling W&B and ckpt saving.")
+    wandb_log = True
+    save_ckpt = True
+
+if wandb_log:
+    wandb_project = 'fMRI-foundation-model'
+    wandb_config = {
+        "model_name": model_name+'_HCP_FT',
+        "batch_size": batch_size,
+        "learning_rate": learning_rate,
+        "weight_decay": weight_decay,
+        "num_epochs": num_epochs,
+        "seed": seed,
+    }
+    print("wandb_config:\n", wandb_config)
+    random_id = str(uuid.uuid4())
+    print("wandb_id:", "HCPflat_raw" + f"_{random_id}")
+    wandb.init(
+        id=model_name+'_HCP_FT' + f"_{random_id}",
+        project=wandb_project,
+        name=model_name+'_HCP_FT',
+        config=wandb_config,
+        resume="allow",
+    )
+
+
+# In[13]:
+
+
+for epoch in range(num_epochs):
+    running_train_loss = 0.0
+    correct_train = 0
+    total_train = 0
+    step = 0
+
+    # with torch.amp.autocast(device_type='cuda'):
+    # Training Phase
+    model.train()
+    for  batch in tqdm(train_dl, desc=f"Epoch {epoch+1}/{num_epochs} - Training"):
+        optimizer.zero_grad()
+        images = batch[0].to(device).float().unsqueeze(1) #fix this  # Shape: [batch_size, 1, 16, 144, 320]
+        labels = batch[1]['trial_type']  # List of labels
+        
+        encoded_labels = label_encoder.transform(labels)
+        encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)  # Shape: [batch_size]
+        
+        # Forward pass
+        outputs = model(images, gsr=gsr)  # Shape: [num_train_samples, num_classes]
+        
+        # Compute loss
+        loss = criterion(outputs, encoded_labels)
+        
+        # Backward pass and optimization
+        loss.backward()
+        optimizer.step()
+        
+        # Accumulate loss
+        running_train_loss += loss.item() * images.size(0)
+
+        
+        # Calculate accuracy
+        _, predicted = torch.max(outputs, 1)
+        
+        correct_train += (predicted == encoded_labels).sum().item()
+        total_train += encoded_labels.size(0)
+        
+        step = step + 1
+        if step % 100 == 0:
+            print(f"Step [{step}/{len(train_dl)}] - Training Loss: {loss.item():.4f} - Training Accuracy: {100 * correct_train / total_train:.2f}%")
+        # thth
+
+    epoch_train_loss = running_train_loss / total_train if total_train > 0 else 0.0
+    train_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0
+    
+    # Validation Phase
+    model.eval()
+    running_val_loss = 0.0
+    correct_val = 0
+    total_val = 0
+    
+    with torch.no_grad():
+        for batch in tqdm(test_dl, desc=f"Epoch {epoch+1}/{num_epochs} - Validation"):
+            
+            images = batch[0].to(device).float().unsqueeze(1) #fix this
+            labels = batch[1]['trial_type']
+            
+            # Encode labels to integer indices
+            encoded_labels = label_encoder.transform(labels)
+            encoded_labels = torch.tensor(encoded_labels, dtype=torch.long).to(device)
+            
+        
+            # Forward pass
+            outputs = model(images, gsr=gsr)
+            
+            # Compute loss
+            loss = criterion(outputs, encoded_labels)
+            
+            # Accumulate loss
+            running_val_loss += loss.item() * images.size(0)
+            
+            # Calculate accuracy
+            _, predicted = torch.max(outputs, 1)
+            correct_val += (predicted == encoded_labels).sum().item()
+            total_val += encoded_labels.size(0)
+
+            
+    
+    epoch_val_loss = running_val_loss / total_val if total_val > 0 else 0.0
+    val_accuracy = 100 * correct_val / total_val if total_val > 0 else 0.0
+    
+    print(f"Epoch [{epoch+1}/{num_epochs}] "
+        f"- Training Loss: {epoch_train_loss:.4f}, Training Accuracy: {train_accuracy:.2f}% "
+        f"- Validation Loss: {epoch_val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%")
+    
+    if wandb_log:
+        wandb.log({
+            "epoch_train_loss": epoch_train_loss,
+            "epoch_val_loss": epoch_val_loss,
+            "train_accuracy": train_accuracy,
+            "val_accuracy": val_accuracy,
+        })
+    if save_ckpt:
+        outdir = os.path.abspath(f'checkpoints/{model_name+"HCP_FT"}')
+        os.makedirs(outdir, exist_ok=True)
+        print("outdir", outdir)
+        # Save model and config
+        torch.save(model.state_dict(), f"{outdir}/model.pth")
+        with open(f"{outdir}/config.yaml", 'w') as f:
+            yaml.dump(wandb_config, f)
+        print(f"Saved model and config to {outdir}")
+    
+

fMRI-foundation-model/src/wandb/run-20241024_213834-HCPflat_large_gsrFalse__HCP_FT_d07870e8-3c53-420c-b734-f32a1c5c3e5a/files/requirements.txt

@@ -0,0 +1,198 @@
+protobuf==5.28.2
+imageio==2.35.1
+MarkupSafe==3.0.0
+regex==2024.9.11
+matplotlib==3.9.2
+notebook==7.2.2
+debugpy==1.8.6
+aiosignal==1.3.1
+jupyter_core==5.7.2
+torchaudio==2.4.1+cu121
+python-json-logger==2.0.7
+six==1.16.0
+scikit-image==0.24.0
+types-python-dateutil==2.9.0.20241003
+PyYAML==6.0.2
+httpcore==1.0.6
+clip==1.0
+babel==2.16.0
+webcolors==24.8.0
+omegaconf==2.3.0
+webencodings==0.5.1
+kiwisolver==1.4.7
+uri-template==1.3.0
+diffusers==0.23.0
+idna==3.10
+fsspec==2024.9.0
+parso==0.8.4
+setuptools==65.5.0
+tornado==6.4.1
+webdataset==0.2.100
+decord==0.6.0
+nvidia-curand-cu12==10.3.2.106
+ipykernel==6.29.5
+jupyter==1.1.1
+pexpect==4.9.0
+kornia_rs==0.1.5
+iopath==0.1.10
+async-lru==2.0.4
+future==1.0.0
+torchvision==0.19.1+cu121
+botocore==1.34.162
+cycler==0.12.1
+tzdata==2024.2
+jupyter_server_terminals==0.5.3
+click==8.1.7
+einops==0.8.0
+pyzmq==26.2.0
+jupyter_client==8.6.3
+nbconvert==7.16.4
+scikit-learn==1.5.2
+executing==2.1.0
+asttokens==2.4.1
+docker-pycreds==0.4.0
+matplotlib-inline==0.1.7
+overrides==7.7.0
+websocket-client==1.8.0
+nbformat==5.10.4
+elbow==0.1.1
+contourpy==1.3.0
+nvidia-cudnn-cu12==9.1.0.70
+transformers==4.44.2
+gitdb==4.0.11
+jupyterlab_nvdashboard==0.11.0
+lazy_loader==0.4
+jsonpointer==3.0.0
+notebook_shim==0.2.4
+nvidia-nccl-cu12==2.20.5
+ffmpeg-python==0.2.0
+triton==3.0.0
+mistune==3.0.2
+python-dateutil==2.9.0.post0
+beautifulsoup4==4.12.3
+nbclient==0.10.0
+h5py==3.12.1
+ftfy==6.2.3
+zipp==3.20.2
+ptyprocess==0.7.0
+huggingface-hub==0.25.1
+pytz==2024.2
+jupyterlab_pygments==0.3.0
+nvidia-cublas-cu12==12.1.3.1
+pandocfilters==1.5.1
+Jinja2==3.1.4
+arrow==1.3.0
+rpds-py==0.20.0
+jupyter_server==2.14.2
+simplejson==3.19.3
+networkx==3.3
+packaging==24.1
+traitlets==5.14.3
+pandas==2.2.3
+xformers==0.0.22.post7
+lightning-utilities==0.11.7
+tifffile==2024.9.20
+nvidia-cuda-cupti-cu12==12.1.105
+mpmath==1.3.0
+GitPython==3.1.43
+scipy==1.14.1
+jsonschema==4.23.0
+prompt_toolkit==3.0.48
+s3transfer==0.10.2
+multidict==6.1.0
+bleach==6.1.0
+sentry-sdk==2.15.0
+nibabel==5.2.1
+accelerate==1.0.0
+pyarrow==17.0.0
+threadpoolctl==3.5.0
+attrs==24.2.0
+rfc3986-validator==0.1.1
+nvidia-cuda-runtime-cu12==12.1.105
+ipywidgets==8.1.5
+frozenlist==1.4.1
+pycparser==2.22
+jupyterlab_server==2.27.3
+nvidia-cuda-nvrtc-cu12==12.1.105
+yarl==1.13.1
+setproctitle==1.3.3
+isoduration==20.11.0
+Pygments==2.18.0
+jedi==0.19.1
+boto3==1.34.57
+tokenizers==0.19.1
+referencing==0.35.1
+rfc3339-validator==0.1.4
+pillow==10.4.0
+jupyterlab==4.2.5
+stack-data==0.6.3
+h11==0.14.0
+anyio==4.6.0
+nilearn==0.10.4
+nvidia-cusolver-cu12==11.4.5.107
+tinycss2==1.3.0
+defusedxml==0.7.1
+argon2-cffi-bindings==21.2.0
+soupsieve==2.6
+nest-asyncio==1.6.0
+torchmetrics==1.3.0.post0
+tqdm==4.66.5
+cffi==1.17.1
+charset-normalizer==3.3.2
+jsonschema-specifications==2023.12.1
+decorator==5.1.1
+open_clip_torch==2.26.1
+jupyter-events==0.10.0
+smart-open==7.0.5
+antlr4-python3-runtime==4.9.3
+prometheus_client==0.21.0
+kornia==0.7.3
+typing_extensions==4.12.2
+sniffio==1.3.1
+joblib==1.4.2
+comm==0.2.2
+aiohappyeyeballs==2.4.3
+numpy==2.1.2
+braceexpand==0.1.7
+certifi==2024.8.30
+psutil==6.0.0
+pyparsing==3.1.4
+pure_eval==0.2.3
+nvidia-cusparse-cu12==12.1.0.106
+wandb==0.18.3
+urllib3==2.2.3
+smmap==5.0.1
+platformdirs==4.3.6
+torch==2.4.1+cu121
+requests==2.32.3
+json5==0.9.25
+nvidia-nvjitlink-cu12==12.6.77
+jupyterlab_widgets==3.0.13
+lxml==5.3.0
+httpx==0.27.2
+opencv-python==4.6.0.66
+portalocker==2.10.1
+pytorch-lightning==2.0.1
+sympy==1.13.3
+wcwidth==0.2.13
+jmespath==1.0.1
+fqdn==1.5.1
+pynvml==11.5.3
+pip==24.0
+wrapt==1.16.0
+aiohttp==3.10.9
+filelock==3.16.1
+fonttools==4.54.1
+fastjsonschema==2.20.0
+jupyter-console==6.6.3
+widgetsnbextension==4.0.13
+timm==1.0.9
+nvidia-cufft-cu12==11.0.2.54
+ipython==8.28.0
+nvidia-nvtx-cu12==12.1.105
+jupyter-lsp==2.2.5
+safetensors==0.4.5
+terminado==0.18.1
+argon2-cffi==23.1.0
+Send2Trash==1.8.3
+importlib_metadata==8.5.0

fMRI-foundation-model/src/wandb/run-20241024_213834-HCPflat_large_gsrFalse__HCP_FT_d07870e8-3c53-420c-b734-f32a1c5c3e5a/files/wandb-metadata.json

@@ -0,0 +1,131 @@
+{
+  "os":  "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+  "python":  "3.11.10",
+  "startedAt":  "2024-10-24T21:38:35.093209Z",
+  "args":  [
+    "HCPflat_large_gsrFalse_",
+    "epoch99.pth"
+  ],
+  "program":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src/HCP_downstream_finetune.py",
+  "codePath":  "src/HCP_downstream_finetune.py",
+  "git":  {
+    "remote":  "https://github.com/MedARC-AI/fMRI-foundation-model",
+    "commit":  "cf8214d4ebe437188b68b4ee5a34c5211a810db0"
+  },
+  "email":  "torrico.villanueva.cesar.kadir@gmail.com",
+  "root":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+  "host":  "ip-10-0-152-216",
+  "username":  "ckadirt",
+  "executable":  "/admin/home-ckadirt/foundation_env/bin/python",
+  "codePathLocal":  "HCP_downstream_finetune.py",
+  "cpu_count":  96,
+  "cpu_count_logical":  192,
+  "gpu":  "[NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3, NVIDIA H100 80GB HBM3]",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "249555763200",
+      "used":  "182052564992"
+    }
+  },
+  "memory":  {
+    "total":  "2147443372032"
+  },
+  "cpu":  {
+    "count":  96,
+    "countLogical":  192
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    },
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper"
+    }
+  ],
+  "slurm":  {
+    "cluster_name":  "sagemaker2",
+    "conf":  "/opt/slurm/etc/slurm.conf",
+    "cpus_on_node":  "20",
+    "gpus_on_node":  "1",
+    "gpus_per_task":  "1",
+    "gtids":  "0",
+    "job_account":  "fmri",
+    "job_cpus_per_node":  "20",
+    "job_end_time":  "1729921061",
+    "job_gid":  "1879800513",
+    "job_gpus":  "1",
+    "job_id":  "529188",
+    "job_name":  "finetuneHCP",
+    "job_nodelist":  "ip-10-0-152-216",
+    "job_num_nodes":  "1",
+    "job_partition":  "p5",
+    "job_qos":  "normal",
+    "job_start_time":  "1729805861",
+    "job_uid":  "1879804696",
+    "job_user":  "ckadirt",
+    "jobid":  "529188",
+    "localid":  "0",
+    "mem_per_cpu":  "11500",
+    "nnodes":  "1",
+    "node_aliases":  "(null)",
+    "nodeid":  "0",
+    "nodelist":  "ip-10-0-152-216",
+    "nprocs":  "1",
+    "ntasks":  "1",
+    "ntasks_per_node":  "1",
+    "prio_process":  "0",
+    "procid":  "0",
+    "script_context":  "prolog_task",
+    "submit_dir":  "/weka/proj-fmri/ckadirt/fMRI-foundation-model/src",
+    "submit_host":  "ip-172-17-12-61",
+    "task_pid":  "337193",
+    "tasks_per_node":  "1",
+    "topology_addr":  "ip-10-0-152-216",
+    "topology_addr_pattern":  "node",
+    "working_cluster":  "sagemaker2:ip-172-17-63-161:6817:9984:109"
+  },
+  "cudaVersion":  "12.2"
+}

fMRI-foundation-model/src/wandb/run-20241024_213834-HCPflat_large_gsrFalse__HCP_FT_d07870e8-3c53-420c-b734-f32a1c5c3e5a/logs/debug-core.log

@@ -0,0 +1,21 @@
+{"time":"2024-10-24T21:38:34.366295935Z","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmpwn5d9jvk/port-337263.txt","pid":337263,"debug":false,"disable-analytics":false}
+{"time":"2024-10-24T21:38:34.366306785Z","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmp_63dqr1q/port-337587.txt","pid":337587,"debug":false,"disable-analytics":false}
+{"time":"2024-10-24T21:38:34.366577281Z","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false}
+{"time":"2024-10-24T21:38:34.366603112Z","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false}
+{"time":"2024-10-24T21:38:34.372814525Z","level":"INFO","msg":"Will exit if parent process dies.","ppid":337587}
+{"time":"2024-10-24T21:38:34.372817685Z","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":43157,"Zone":""}}
+{"time":"2024-10-24T21:38:34.373811796Z","level":"INFO","msg":"Will exit if parent process dies.","ppid":337263}
+{"time":"2024-10-24T21:38:34.373817646Z","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":44789,"Zone":""}}
+{"time":"2024-10-24T21:38:34.528156376Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:38070"}
+{"time":"2024-10-24T21:38:34.528235848Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:51210"}
+{"time":"2024-10-24T21:38:35.06181568Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_large_gsrFalse__HCP_FT_cdc7b70b-b155-4483-907d-91148a52e1d7","id":"127.0.0.1:38070"}
+{"time":"2024-10-24T21:38:35.093628892Z","level":"INFO","msg":"handleInformInit: received","streamId":"HCPflat_large_gsrFalse__HCP_FT_d07870e8-3c53-420c-b734-f32a1c5c3e5a","id":"127.0.0.1:51210"}
+{"time":"2024-10-24T21:38:35.140878185Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_large_gsrFalse__HCP_FT_cdc7b70b-b155-4483-907d-91148a52e1d7","id":"127.0.0.1:38070"}
+{"time":"2024-10-24T21:38:35.145483394Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"HCPflat_large_gsrFalse__HCP_FT_d07870e8-3c53-420c-b734-f32a1c5c3e5a","id":"127.0.0.1:51210"}
+{"time":"2024-10-26T02:19:36.419933572Z","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"127.0.0.1:38070"}
+{"time":"2024-10-26T02:19:36.421381993Z","level":"INFO","msg":"server is shutting down"}
+{"time":"2024-10-26T02:19:36.421375983Z","level":"INFO","msg":"connection: Close: initiating connection closure","id":"127.0.0.1:38070"}
+{"time":"2024-10-26T02:19:36.421592698Z","level":"INFO","msg":"connection: Close: connection successfully closed","id":"127.0.0.1:38070"}
+{"time":"2024-10-26T02:19:38.264827227Z","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"127.0.0.1:38070"}
+{"time":"2024-10-26T02:19:38.264913688Z","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"127.0.0.1:38070"}
+{"time":"2024-10-26T02:19:38.264935459Z","level":"INFO","msg":"server is closed"}