fMRI-foundation-model/src/HCP_downstream_finetune.py

#!/usr/bin/env python
# coding: utf-8

# In[40]:


# 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
import pandas as pd
from sklearn.preprocessing import StandardScaler
from typing import List, Dict, Any, Tuple
import argparse

# 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




# In[48]:


# if running this interactively, can specify jupyter_args here for argparser to use
if utils.is_interactive():
    model_name_suffix = "testing"
    print("model_name_suffix:", model_name_suffix)
    
    # global_batch_size and batch_size should already be defined in the 2nd cell block
    jupyter_args = f"--found_model_name=HCPflat_large_gsrFalse_ --epoch_checkpoint epoch99.pth \
                    --hcp_flat_path=/weka/proj-medarc/shared/HCP-Flat \
                    --target=sex \
                    --model_suffix={model_name_suffix} \
                    --batch_size={batch_size} \
                    --max_lr=3e-4 --num_epochs=20 --no-save_ckpt --no-wandb_log --num_workers=10 \
                    --weight_decay=1e-5 \
                    --global_pool"
    # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow

    print(jupyter_args)
    jupyter_args = jupyter_args.split()
    
    from IPython.display import clear_output # function to clear print outputs in cell
    get_ipython().run_line_magic('load_ext', 'autoreload')
    # this 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('autoreload', '2')


# In[49]:


parser = argparse.ArgumentParser(description="Model Training Configuration")
parser.add_argument(
    "--found_model_name", type=str, default="Testing_flat",
    help="name of model, used for ckpt saving and wandb logging (if enabled)",
)
parser.add_argument(
    "--epoch_checkpoint", type=str, default="epoch99.pth",
    help="the epoch number of the found_model_name checkpoint",
)
parser.add_argument(
    "--model_suffix", type=str, default="Testing_flat",
    help="name of model, used for ckpt saving and wandb logging (if enabled)",
)
parser.add_argument(
    "--hcp_flat_path", type=str, default=os.getcwd(),
    help="Path to where NSD data is stored / where to download it to",
)
parser.add_argument(
    "--batch_size", type=int, default=128,
    help="Batch size can be increased by 10x if only training retreival submodule and not diffusion prior",
)
parser.add_argument(
    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
    help="whether to log to wandb",
)
parser.add_argument(
    "--num_epochs",type=int,default=150,
    help="number of epochs of training",
)
parser.add_argument(
    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
)
parser.add_argument(
    "--save_ckpt",action=argparse.BooleanOptionalAction,default=True,
)
parser.add_argument(
    "--seed",type=int,default=42,
)
parser.add_argument(
    "--max_lr",type=float,default=3e-4,
)
parser.add_argument(
    "--target",type=str,default='trial_type',choices=['trial_type','sex','age'],
)
parser.add_argument(
    "--num_workers",type=int,default=10,
)
parser.add_argument(
    "--weight_decay",type=float,default=1e-5,
)
parser.add_argument(
    "--global_pool",action=argparse.BooleanOptionalAction,default=False,
    help="not implemented yet",
)

if utils.is_interactive():
    args = parser.parse_args(jupyter_args)
else:
    args = parser.parse_args()

print(f"------ ARGS ------- \n {args}")

# create global variables without the args prefix
for attribute_name in vars(args).keys():
    globals()[attribute_name] = getattr(args, attribute_name)
    
# seed all random functions
utils.seed_everything(seed)


# In[ ]:


# ## MODEL TO LOAD ##
# if utils.is_interactive():
#     model_name = "HCPflat_large_gsrFalse_"
# else:
#     model_name = sys.argv[1]
    
# target = 'sex' # This can be 'trial_type' 'age' 'sex'


# In[50]:


# outdir = os.path.abspath(f'checkpoints/{model_name}')
outdir = os.path.abspath(f'checkpoints/{found_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[55]:


# 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}")

for attribute_name in vars(args).keys():
    globals()[attribute_name] = getattr(args, attribute_name)


# 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[56]:


from sklearn.preprocessing import LabelEncoder
if target == "trial_type":
    
    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[57]:


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[58]:


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("Creating datasets")
# 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")


# In[59]:


# open the file containing subject information
if target == "age" or target == "sex":
    subject_information_HCP_path = os.path.join(hcp_flat_path, "subjects_data_restricted.csv")
    try:
        subject_information_HCP = pd.read_csv(subject_information_HCP_path)
    except:
        try:
            subject_information_HCP = pd.read_csv('./unrestricted_clane9_4_23_2024_13_28_14.csv')   
        except:
            assert False, "Subject information file not found"

    ###### This is for unrestricted
    # age_related_columns = [
    #     'Age', 'PicSeq_AgeAdj', 'CardSort_AgeAdj', 'Flanker_AgeAdj',
    #     'ReadEng_AgeAdj', 'PicVocab_AgeAdj', 'ProcSpeed_AgeAdj',
    #     'CogFluidComp_AgeAdj', 'CogEarlyComp_AgeAdj', 'CogTotalComp_AgeAdj',
    #     'CogCrystalComp_AgeAdj', 'Endurance_AgeAdj', 'Dexterity_AgeAdj',
    #     'Strength_AgeAdj', 'Odor_AgeAdj', 'Taste_AgeAdj'
    # ]
    
    # sex_related_columns = [
    #     'Gender'
    # ]

    ###### This is for restricted
    gender_related_columns = [
        'Gender'
    ]

    age_related_columns = [
        'Age_in_Yrs',
        'Menstrual_AgeBegan',
        'Menstrual_AgeIrreg',
        'Menstrual_AgeStop',
        'SSAGA_Alc_Age_1st_Use',
        'SSAGA_TB_Age_1st_Cig',
        'SSAGA_Mj_Age_1st_Use',
        'Endurance_AgeAdj',
        'Dexterity_AgeAdj',
        'Strength_AgeAdj',
        'PicSeq_AgeAdj',
        'CardSort_AgeAdj',
        'Flanker_AgeAdj',
        'ReadEng_AgeAdj',
        'PicVocab_AgeAdj',
        'ProcSpeed_AgeAdj',
        'Odor_AgeAdj',
        'Taste_AgeAdj'
    ]

    # # show the first few rows of the subject information
    # subject_information_HCP[age_related_columns + sex_related_columns].head()

    # Handle missing values (e.g., impute with mean)
    mean_age = subject_information_HCP['Age_in_Yrs'].mean()
    
    # Initialize the scaler
    scaler = StandardScaler()
    
    # Perform z-score normalization
    subject_information_HCP['Age_in_Yrs_z'] = scaler.fit_transform(subject_information_HCP[['Age_in_Yrs']])


    
def get_label_unrestricted(subject_id: List[str], target: str, method_for_age: str = 'mean') -> List:
    """
    Get the label for the given subject id and target.

    For sex 0 is F and 1 is M
    """

    # convert to list of ints
    subject_id = [int(x) for x in subject_id]

    if target == "age":
        age_array = []
        for subject in subject_id:
            c_age = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Age'].values
            # if the subject is not in the subject information file trigger an error
            if len(c_age) == 0:
                assert False, f"Subject {subject} not found in subject information file"
            if len(c_age) > 1:
                print(f"Warning: Multiple entries for subject {subject}")

            c_age = c_age[0].split('-')
            if len(c_age) < 2:
                c_age = c_age[0].split('+')
                age_array.append(int(c_age[0]))
            else:
                if method_for_age == 'mean':
                    age_array.append(np.mean([int(x) for x in c_age]))
                elif method_for_age == 'min':
                    age_array.append(np.min([int(x) for x in c_age]))
                elif method_for_age == 'max':
                    age_array.append(np.max([int(x) for x in c_age]))
                else:
                    assert False, f"Method {method_for_age} not recognized"

        return np.array(age_array)  
    
    elif target == 'sex':
        sex_array = []
        for subject in subject_id:
            c_sex = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Gender'].values
            # if the subject is not in the subject information file trigger an error
            if len(c_sex) == 0:
                assert False, f"Subject {subject} not found in subject information file"
            if len(c_sex) > 1:
                print(f"Warning: Multiple entries for subject {subject}")
            sex_array.append(int(c_sex[0] == 'M'))
        return sex_array

def get_label_restricted(subject_id: List[str], target: str, normalized: bool = True) -> List:
    """
    Get the label for the given subject id and target.

    For sex 0 is F and 1 is M
    """

    # convert to list of ints
    subject_id = [int(x) for x in subject_id]

    if target == "age":
        age_array = []
        for subject in subject_id:
            c_age = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Age_in_Yrs' if not normalized else 'Age_in_Yrs_z'].values
            # if the subject is not in the subject information file trigger an error
            if len(c_age) == 0:
                assert False, f"Subject {subject} not found in subject information file"
            if len(c_age) > 1:
                print(f"Warning: Multiple entries for subject {subject}")

            age_array.append(np.int8(c_age[0]))

        return np.array(age_array)  
    
    elif target == 'sex':
        sex_array = []
        for subject in subject_id:
            c_sex = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Gender'].values
            # if the subject is not in the subject information file trigger an error
            if len(c_sex) == 0:
                assert False, f"Subject {subject} not found in subject information file"
            if len(c_sex) > 1:
                print(f"Warning: Multiple entries for subject {subject}")
            sex_array.append(int(c_sex[0] == 'M'))
        return sex_array


# ### Creating and loading Model

# In[60]:


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[61]:


checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]

if utils.is_interactive():
    latest_checkpoint = "epoch99.pth"
else:
    latest_checkpoint = epoch_checkpoint
    
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[62]:


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[63]:


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[64]:


# Initialize the model

if target == "trial_type":
    lc_model = LinearClassifier(input_dim=input_dim, num_classes=num_classes)
    criterion = nn.CrossEntropyLoss()

elif target == "age":
    lc_model = LinearClassifier(input_dim=input_dim, num_classes=1)
    criterion = nn.MSELoss()

elif target == "sex":
    lc_model = LinearClassifier(input_dim=input_dim, num_classes=1)
    criterion = nn.BCEWithLogitsLoss()


# 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 optimizer with L2 regularization (weight_decay)
# learning_rate = 1e-4
# weight_decay = 1e-5  # Adjust based on your needs

optimizer = torch.optim.AdamW(model.parameters(), lr=max_lr, weight_decay=weight_decay)

num_iterations_per_epoch = math.ceil(flatmaps_train.shape[0]/batch_size)

if lr_scheduler_type == 'linear':
    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
        optimizer,
        total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)),
        last_epoch=-1
    )
elif lr_scheduler_type == 'cycle':
    total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))
    print("total_steps", total_steps)
    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, 
        max_lr=max_lr,
        total_steps=total_steps,
        final_div_factor=1000,
        last_epoch=-1, pct_start=2/num_epochs
    )



# num_epochs = 20  # Adjust as needed


# In[29]:


# criterion


# ### Data

# In[65]:


import uuid

myuuid = uuid.uuid4()
str(myuuid)


# In[66]:


import wandb

if utils.is_interactive():
    print("Running in interactive notebook. Disabling W&B and ckpt saving.")
    wandb_log = False
    save_ckpt = False

if wandb_log:
    wandb_project = 'fMRI-foundation-model'
    wandb_config = {
        "model_name": f'{found_model_name}_HCP_FT_{target}',
        "batch_size": batch_size,
        "weight_decay": weight_decay,
        "num_epochs": num_epochs,
        "seed": seed,
        "lr_scheduler_type": lr_scheduler_type,
        "save_ckpt": save_ckpt,
        "seed": seed,
        "max_lr": max_lr,
        "target": target,
        "num_workers": num_workers,
        "weight_decay": weight_decay
    }
    print("wandb_config:\n", wandb_config)
    random_id = random.randint(0, 100000)
    wandb_id = f"{found_model_name}_{model_suffix}_{target}_HCPFT_{myuuid}"
    print("wandb_id:", wandb_id)
    wandb.init(
        id=wandb_id,
        project=wandb_project,
        name=f"{found_model_name}_{model_suffix}_{target}_HCPFT",
        config=wandb_config,
        resume="allow",
    )


# In[ ]:


for epoch in range(num_epochs):
    running_train_loss = 0.0
    correct_train = 0
    mse_age_train = 0.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()  # Shape: [batch_size, 1, 16, 144, 320]
        
        # Prepare labels based on target type
        if target == "trial_type":
            labels = batch[1]['trial_type']  # List of labels
            labels = label_encoder.transform(labels)
            labels = torch.tensor(labels, dtype=torch.long).to(device)  # Shape: [batch_size]
        elif target == "age":
            labels = get_label_restricted(batch[1]['sub'], 'age')
            labels = torch.tensor(labels, dtype=torch.float).to(device)  # Shape: [batch_size]
        elif target == "sex":
            labels = get_label_restricted(batch[1]['sub'], 'sex')
            labels = torch.tensor(labels, dtype=torch.float).to(device)  # Shape: [batch_size]
        # labels = labels.unsqueeze(1)
        
        # Forward pass
        outputs = model(images, gsr=gsr)  # Shape: [num_train_samples, num_classes]
        
        # Compute loss
        if target in ["trial_type", "sex"]:
            # For classification, ensure outputs are logits
            loss = criterion(outputs.squeeze(), labels.squeeze())
        elif target == "age":
            # For regression, ensure outputs are single values
            loss = criterion(outputs.squeeze(), labels.squeeze())
        
        # Backward pass and optimization
        loss.backward()
        optimizer.step()
        
        # Accumulate loss
        running_train_loss += loss.item() * images.size(0)

        
       # Calculate and accumulate metrics
        if target == "trial_type":
            _, predicted = torch.max(outputs, 1)
            correct_train += (predicted == labels).sum().item()
        elif target == "age":
            mse_age_train += (torch.sum((outputs.squeeze() - labels) ** 2).item()) / outputs.shape[0]
        elif target == "sex":
            threshold = 0.5
            predicted = (torch.sigmoid(outputs) > threshold).float()
            correct_train += (predicted == labels).sum().item()

        total_train += labels.size(0)
        step += 1

        # Print intermediate metrics every 100 steps
        if step % 100 == 0:
            if target in ["trial_type", "sex"]:
                current_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0
                print(f"Step [{step}/{len(train_dl)}] - Training Loss: {loss.item():.4f} - Training Accuracy: {current_accuracy:.2f}%")
            elif target == "age":
                current_mse = mse_age_train / total_train if total_train > 0 else 0.0
                print(f"Step [{step}/{len(train_dl)}] - Training Loss: {loss.item():.4f} - Training MSE: {current_mse:.4f}")

        if lr_scheduler_type is not None:
                lr_scheduler.step()

    # Calculate epoch-level metrics
    epoch_train_loss = running_train_loss / total_train if total_train > 0 else 0.0

    if target in ["trial_type", "sex"]:
        train_accuracy = 100 * correct_train / total_train if total_train > 0 else 0.0
    elif target == "age":
        train_mse = mse_age_train / total_train if total_train > 0 else 0.0
    
    # Validation Phase
    model.eval()
    running_val_loss = 0.0
    correct_val = 0
    mse_age_val = 0.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()
            labels = batch[1]['trial_type']
            
            # Prepare labels based on target type
            if target == "trial_type":
                labels = batch[1]['trial_type']  # List of labels
                labels = label_encoder.transform(labels)
                labels = torch.tensor(labels, dtype=torch.long).to(device)  # Shape: [batch_size]
            elif target == "age":
                labels = get_label_restricted(batch[1]['sub'], 'age')
                labels = torch.tensor(labels, dtype=torch.float).to(device)  # Shape: [batch_size]
            elif target == "sex":
                labels = get_label_restricted(batch[1]['sub'], 'sex')
                labels = torch.tensor(labels, dtype=torch.float).to(device)  # Shape: [batch_size]

            # labels = labels.unsqueeze(1)
            # Forward pass
            outputs = model(images, gsr=gsr)
            
            # Compute loss
            if target in ["trial_type", "sex"]:
                loss = criterion(outputs.squeeze(), labels.squeeze())
            elif target == "age":
                loss = criterion(outputs.squeeze(), labels.squeeze())

            # Accumulate loss
            running_val_loss += loss.item() * images.size(0)

            # Calculate and accumulate metrics
            if target == "trial_type":
                _, predicted = torch.max(outputs, 1)
                correct_val += (predicted == labels).sum().item()
            elif target == "age":
                mse_age_val += (torch.sum((outputs.squeeze() - labels) ** 2).item()) / outputs.shape[0]
            elif target == "sex":
                threshold = 0.5
                predicted = (torch.sigmoid(outputs) > threshold).float()
                correct_val += (predicted == labels).sum().item()

            total_val += labels.size(0)

    # Calculate epoch-level validation metrics
    epoch_val_loss = running_val_loss / total_val if total_val > 0 else 0.0

    if target in ["trial_type", "sex"]:
        val_accuracy = 100 * correct_val / total_val if total_val > 0 else 0.0
    elif target == "age":
        val_mse = mse_age_val / total_val if total_val > 0 else 0.0

    # Print epoch-level metrics
    if target in ["trial_type", "sex"]:
        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}%")
    elif target == "age":
        print(f"Epoch [{epoch+1}/{num_epochs}] "
              f"- Training Loss: {epoch_train_loss:.4f}, Training MSE: {train_mse:.4f} "
              f"- Validation Loss: {epoch_val_loss:.4f}, Validation MSE: {val_mse:.4f}")

    # Log metrics with wandb
    if wandb_log:
        log_dict = {
            "epoch_train_loss": epoch_train_loss,
            "epoch_val_loss": epoch_val_loss,
        }
        if target in ["trial_type", "sex"]:
            log_dict.update({
                f"train_accuracy_{target}": train_accuracy,
                f"val_accuracy_{target}": val_accuracy,
            })
        elif target == "age":
            log_dict.update({
                f"train_mse_{target}": train_mse,
                f"val_mse_{target}": val_mse,
            })
        wandb.log(log_dict)
        
    if save_ckpt:
        outdir = os.path.abspath(f'checkpoints/{f"{found_model_name}_{model_suffix}_{target}_HCPFT"}')
        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}")