Upload folder using huggingface_hub

ckpts/part1_voxel_indices.pt

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

ckpts/subj01_part1.pth

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

ckpts/subj01_part2.pth

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

ckpts/subj02_part1.pth

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

ckpts/subj02_part2.pth

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

ckpts/subj03_part1.pth

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

ckpts/subj03_part2.pth

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

ckpts/subj04_part1.pth

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

ckpts/subj04_part2.pth

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

ckpts/subj05_part1.pth

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

ckpts/subj05_part2.pth

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

ckpts/subj06_part1.pth

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

ckpts/subj06_part2.pth

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

ckpts/subj07_part1.pth

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

ckpts/subj07_part2.pth

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

ckpts/subj08_part1.pth

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

ckpts/subj08_part2.pth

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

config.yaml

@@ -9,10 +9,10 @@ DATAMODULE:
   PIN_MEMORY: true
 DATASET:
   CACHE_DIR: /data/cache
-  DARK_POSTFIX: 
+  DARK_POSTFIX: xvdb
   FILTER_BY_SESSION:
   - -1
-  FMRI_SPACE: full_fsaverage
+  FMRI_SPACE: fsaverage
   IMAGE_RESOLUTION:
   - 224
   - 224
@@ -22,13 +22,6 @@ DATASET:
   ROOT: /data/ALG23
   SUBJECT_LIST:
   - subj01
-  - subj02
-  - subj03
-  - subj04
-  - subj05
-  - subj06
-  - subj07
-  - subj08
 DESCRIPTION: for alex
 EXPERIMENTAL:
   ANOTHER_SPLIT: false
@@ -53,8 +46,8 @@ EXPERIMENTAL:
   USE_FTR_BEHV: false
 LOSS:
   DARK:
-    MAX_EPOCH: 150
-    USE: false
+    MAX_EPOCH: 50
+    USE: true
   NAME: SmoothL1Loss
   SMOOTH_L1_BETA: 0.01
   SYNC:
@@ -67,7 +60,7 @@ LOSS:
     SKIP_EPOCHS: 20
     STAGE: VAL
     UPDATE_RULE: raw
-    USE: false
+    USE: true
 MODEL:
   BACKBONE:
     ADAPTIVE_LN:
@@ -150,13 +143,13 @@ MODEL_SOUP:
   RECIPE: greedy
   USE: true
 OPTIMIZER:
-  LR: 0.0001
+  LR: 0.0003
   NAME: AdamW
   SCHEDULER:
     CYCLE_DECAY: 0.5
     CYCLE_LIMIT: 3
     K_DECAY: 1.5
-    LR_MIN: 0.0001
+    LR_MIN: 0.0003
     LR_MIN_WARMUP: 0.0001
     T_INITIAL: 1
     T_MULT: 1.0

drafts/cpfiles.py

@@ -0,0 +1,22 @@
+# %%
+import os
+# %%
+source_dir = "/nfscc/alg23/xalex_distill2/wb/"
+path_contains = "DARK.MAX_EPOCH=90"
+search_filename = "soup.pth"
+# %%
+files = []
+import glob
+for file in glob.glob(os.path.join(source_dir, f"*{path_contains}*", search_filename)):
+    files.append(file)
+files = sorted(files)
+print(files)
+# %%
+target_dir = "/workspace/model_packed2/ckpts/"
+
+for i_subj, file in enumerate(files):
+    print(f"copy {file} to {target_dir}")
+    target_filename = f"subj{i_subj+1:02d}_part2.pth"
+    target_path = os.path.join(target_dir, target_filename)
+    os.system(f"cp {file} {target_path}")
+# %%

drafts/cpvoxel.py

@@ -0,0 +1,15 @@
+# %%
+import numpy as np
+import torch
+import os
+# %%
+path = "/data/ALG23/subj01/data_mask/fsaverage/voxel_indices.npy"
+part1_voxel_indice_dict = {}
+for i in range(1, 9):
+    part1_voxel_indice_dict[f'subj0{i}'] = np.load(f"/data/ALG23/subj0{i}/data_mask/fsaverage/voxel_indices.npy")
+torch.save(part1_voxel_indice_dict, "/nfscc/alg23/xalex_distill2/high/voxel_indices_dict.pth")
+# %%
+part1_voxel_indice_dict = torch.load("/nfscc/alg23/xalex_distill2/high/voxel_indices_dict.pth")
+# %%
+part1_voxel_indice_dict['subj01']
+# %%

drafts/try_load.py

@@ -0,0 +1,10 @@
+# %%
+path = "/nfscc/alg23/xalex_distill2/high/t826c6_00016_DATASET.SUBJECT_LIST=subj01,LOSS.DARK.MAX_EPOCH=90,/soup.pth"
+import torch
+sd = torch.load(path, map_location='cpu')
+print(sd.keys())
+# %%
+sd['coord_dict.subj01'].shape
+# %%
+sd['model.voxel_outs_weight.subj01.weight'].shape
+# %%

example.py

@@ -1,13 +1,22 @@
-from model import BrainEncodingModel
+from model import _load_one_model, TowPartModel, BrainEncodingModel
 from config_utils import load_from_yaml
 import torch
 
-cfg = load_from_yaml("./config.yaml")
-model = BrainEncodingModel(cfg)
-sd_path = './ckpt.pth'
-sd = torch.load(sd_path)
-model.load_state_dict(sd)
-model.eval().cuda()
+subject = 'subj01'
+cfg_path = "/workspace/model_packed2/config.yaml"
+model_path1 = f"/workspace/model_packed2/ckpts/{subject}_part1.pth"
+model_path2 = f"/workspace/model_packed2/ckpts/{subject}_part2.pth"
+# model1 is for vertices with high noise ceiling (nsdgeneral)
+# model2 is for vertices from the rest of the brain
+model1: BrainEncodingModel = _load_one_model(model_path1, subject, cfg_path)
+model2: BrainEncodingModel = _load_one_model(model_path2, subject, cfg_path)
+# voxel_indices is a list of indices of vertices with high noise ceiling (for model1)
+voxel_indices_path = "/workspace/model_packed2/ckpts/part1_voxel_indices.pt"
+voxel_indices = torch.load(voxel_indices_path)[subject]
+model = TowPartModel(model1, model2, voxel_indices)
+
+model = model.cuda().eval()
+
 
 x = torch.randn(1, 3, 224, 224)
 def transform_image(x):
@@ -18,10 +27,8 @@ def transform_image(x):
 x = transform_image(x)
 x = x.cuda()
     
-    
-subject = 'subj01'  # could be 1 of 8 subjects
 
 with torch.no_grad():
-    out = model(x, subject)
+    out = model(x)
 print(out.shape)
 # torch.Size([1, 327684])

model.py

@@ -1,37 +1,26 @@
+#%%
 from functools import partial
 import logging
-from torch import nn, Tensor
 from einops import rearrange, repeat
 from typing import Dict, Optional, Union
 
 import torch
 import torch.nn.functional as F
+from torch import nn, Tensor
 from config import AutoConfig
 
 from backbone import (
-    SubjectTimeEmbed,
     build_backbone,
     AdaLNLoRADiNOv2ViT,
-    build_backbone_prev,
-    build_time_emd,
 )
 from blocks import (
-    PreviousFeatureMLPs,
-    SubjectPreviousFrameCompress,
-    build_class_token_mlp_prev,
     build_conv_blocks,
     build_class_token_mlp,
     DictConvBlocks,
     ClassTokenMLPs,
-    build_ftr_compress,
-    build_prev_compress,
-    build_prev_feat_mlp,
 )
-from behav_embed import build_behavior_embed, SubjectBehaviorEmbed
 from config_utils import load_from_yaml
 from topyneck import (
-    CoordsFreeWeights,
-    build_coords_free_weights,
     build_coords_mlp,
     CachedCoordsMLP,
     build_voxelouts_weight,
@@ -41,39 +30,16 @@ from topyneck import (
 
 import numpy as np
 
-def get_coords():
-    import nilearn
-    from nilearn import datasets, surface
-
-    fsaverage = nilearn.datasets.fetch_surf_fsaverage("fsaverage7")
-    lh_coords, lh_faces = nilearn.surface.load_surf_mesh(fsaverage["sphere_left"])
-    rh_coords, rh_faces = nilearn.surface.load_surf_mesh(fsaverage["sphere_right"])
-    lh_xmin, lh_xmax = np.min(lh_coords[:, 0]), np.max(lh_coords[:, 0])
-    lh_xmax = lh_xmin + (lh_xmax - lh_xmin) * 1.5
-    rh_xmin, rh_xmax = np.min(rh_coords[:, 0]), np.max(rh_coords[:, 0])
-    if rh_xmin < lh_xmax:
-        rh_coords[:, 0] += lh_xmax - rh_xmin
-    coords = np.concatenate((lh_coords, rh_coords), axis=0)
-    coords = torch.tensor(coords)
-    return coords
-
-# %%
 class BrainEncodingModel(nn.Module):
     def __init__(
         self,
         cfg: AutoConfig,
+        n_voxel_dict = {'subj01': 327684},
     ):
         
         super().__init__()
-        n_voxel_dict = {'subj01': 327684,
-                        'subj02': 327684,
-                        'subj03': 327684,
-                        'subj04': 327684,
-                        'subj05': 327684,
-                        'subj06': 327684,
-                        'subj07': 327684,
-                        'subj08': 327684}
         self.subject_list = list(n_voxel_dict.keys())
+        assert len(self.subject_list) == 1, "Only one subject is supported"
 
         self.layers = cfg.MODEL.BACKBONE.LAYERS
         self.layers_small = cfg.MODEL.BACKBONE_SMALL.LAYERS
@@ -82,19 +48,13 @@ class BrainEncodingModel(nn.Module):
         cfg.MODEL.CONV_HEAD.WIDTH = int(cfg.MODEL.CONV_HEAD.WIDTH * r)
         self.cfg = cfg
 
-        self.behav_embed: SubjectBehaviorEmbed = build_behavior_embed(cfg) 
-        # behavior is not used, just a placeholder
-        
         self.backbone: AdaLNLoRADiNOv2ViT = build_backbone(cfg)
         self.conv_blocks: DictConvBlocks = build_conv_blocks(cfg)
         self.cls_blocks: ClassTokenMLPs = build_class_token_mlp(cfg)
         
         def build_each_subject(fn, subject_list):
             return nn.ModuleDict({subject: fn() for subject in subject_list})
-        
-        self.coords = get_coords()  # [327684, 3], for layer selector and retina mapper
-        self.coords = nn.Parameter(self.coords, requires_grad=False)
-        
+                
         self.layer_selector: Dict[str, CachedCoordsMLP] = build_each_subject(
             partial(
                 build_coords_mlp,
@@ -131,25 +91,26 @@ class BrainEncodingModel(nn.Module):
                 for subject in self.subject_list
             }
         )
+        
+        self.coords : nn.Parameter = None
 
         
     def forward(
         self,
         x: Tensor,  # [B, C, H, W]
-        subject: str,
         voxel_indices: Optional[Tensor] = None,
         chunk_size=4096,
+        **kwargs,
     ):
+        coords = self.coords
+        subject = self.subject_list[0]
+        
         bsz = x.shape[0]
         device = x.device
         dtype = x.dtype
         
-        # bhv is not used, just a placeholder
-        bhv = torch.zeros((bsz, self.cfg.MODEL.COND.IN_DIM), device=device, dtype=dtype)  # [B, D_B=35]
-        c = self.behav_embed(bhv, subject=subject)  # [B, D_C]
-        
         x_retina_grid, x_cls_dict = self.backbone.get_intermediate_layers(
-            x, n=self.layers, c=c
+            x, n=self.layers, c=None
         )
         x_retina_grid = self.conv_blocks(x_retina_grid)
         x_cls_dict = self.cls_blocks(x_cls_dict)
@@ -161,7 +122,6 @@ class BrainEncodingModel(nn.Module):
         #############################
 
         # divide voxels into chunks to avoid OOM
-        coords = self.coords
         n_voxels = coords.shape[0]
         if voxel_indices is None or voxel_indices == ...:
             voxel_indices = torch.arange(n_voxels, device=coords.device)
@@ -186,7 +146,7 @@ class BrainEncodingModel(nn.Module):
         reg_layer = torch.cat(reg_layers, dim=0).mean()  # [1]
 
         # if self.training:
-        #     return out_y, reg_layer
+            # return out_y, reg_layer
         # else:
         return out_y
 
@@ -260,3 +220,64 @@ class BrainEncodingModel(nn.Module):
 
         return out_y, reg_layer  # [B, N], [N]
 
+
+        
+def _load_one_model(model_path: str, subject: str='subj01', cfg_path: str=None):
+    cfg = load_from_yaml(cfg_path)
+    
+    # load model weights
+    sd = torch.load(model_path, map_location='cpu')
+    n_voxels = sd[f'model.voxel_outs_weight.{subject}.weight'].shape[0]
+    # create model
+    model = BrainEncodingModel(cfg, {subject: n_voxels})
+    
+    # save voxel's coordinates to model
+    coords = sd[f'coord_dict.{subject}']
+    model.coords = nn.Parameter(coords)
+    
+    # load weights
+    filtered_sd = {k: v for k, v in sd.items() if k.startswith('model')}
+    filtered_sd = {k[6:]: v for k, v in filtered_sd.items() if k.startswith('model')}
+    filtered_sd['coords'] = model.coords  # add coordinates of voxels
+    model.load_state_dict(filtered_sd)
+    
+    model = model.eval()
+    return model
+
+
+class TowPartModel(nn.Module):
+    def __init__(self, model_part1, model_part2, part1_voxel_indices):
+        super().__init__()
+        self.model_part1 = model_part1
+        self.model_part2 = model_part2
+        self.part1_voxel_indices = part1_voxel_indices
+        
+    def forward(self, x):
+        # x: [B, 3, 224, 224]   # image after normalization
+        out1 = self.model_part1(x)
+        out2 = self.model_part2(x)
+        out = out2
+        out[:, self.part1_voxel_indices] = out1
+        return out
+    
+
+
+
+# %%
+if __name__ == '__main__':
+    # model_path = "/nfscc/alg23/xalex_distill2/high/t826c6_00016_DATASET.SUBJECT_LIST=subj01,LOSS.DARK.MAX_EPOCH=90,/soup.pth"
+    subject = 'subj01'
+    cfg_path = "/workspace/model_packed2/config.yaml"
+    model_path1 = f"/workspace/model_packed2/ckpts/{subject}_part1.pth"
+    model_path2 = f"/workspace/model_packed2/ckpts/{subject}_part2.pth"
+    model1 = _load_one_model(model_path1, subject, cfg_path)
+    model2 = _load_one_model(model_path2, subject, cfg_path)
+    voxel_indices_path = "/workspace/model_packed2/ckpts/part1_voxel_indices.pt"
+    voxel_indices = torch.load(voxel_indices_path)[subject]
+    model = TowPartModel(model1, model2, voxel_indices)
+
+    x = torch.randn(1, 3, 224, 224)
+    x = x.cuda()
+    model = model.cuda()
+    out = model(x)
+    print(out.shape)