Basic files for implementing trained model

Checkpoint/checkpoint49_2024-03-28_Zdim_2_Vae-beta_20.0_Lr_0.0001_Batch-size_128_washu120_subsample10_train100_val10.pth.tar

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

Checkpoint/checkpoint49_2024-06-21_Zdim_4_Vae-beta_20.0_Lr_0.0001_Batch-size_128_washu120_subsample10_train100_val10.pth.tar

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

Checkpoint/checkpoint49_2024-11-28_Zdim_3_Vae-beta_20.0_Lr_0.0001_Batch-size_128_washu120_subsample10_train100_val10.pth.tar

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

VAE_inference_example.py

@@ -0,0 +1,83 @@
+import torch # tested on version 2.1.2+cu118
+import scipy.io as io
+import argparse
+import logging
+from utils import load_dataset_test, save_image_mat
+from fMRIVAE_Model import BetaVAE
+import os
+
+def main():
+
+    parser = argparse.ArgumentParser(description='VAE for fMRI generation')
+    parser.add_argument('--batch-size', type=int, metavar='N', help='how many samples per saved file?')
+    parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)')
+    parser.add_argument('--zdim', type=int, default=256, metavar='N', help='dimension of latent variables')
+    parser.add_argument('--data-path', type=str, metavar='DIR', help='path to dataset')
+    parser.add_argument('--z-path', type=str, default='./result/latent/', help='path to saved z files')
+    parser.add_argument('--resume', type=str, default='./checkpoint/checkpoint.pth.tar', help='the VAE checkpoint') 
+    parser.add_argument('--img-path', type=str, default='./result/recon', help='path to save reconstructed images')
+    parser.add_argument('--mode', type=str, default='both', help='choose from \'encode\',\'decode\' or \'both\'')
+    parser.add_argument('--debug', action='store_true', help='Enable debug mode for detailed logging')
+
+    args = parser.parse_args()
+
+    if not os.path.isdir(args.z_path):
+        os.system('mkdir '+ args.z_path + ' -p')
+    if (args.mode != 'encode') and not os.path.isdir(args.img_path):
+        os.system('mkdir '+ args.img_path + ' -p')
+
+    # Set logging level based on debug flag
+    logging_level = logging.DEBUG if args.debug else logging.INFO
+    logging.basicConfig(level=logging_level, format='%(asctime)s - %(levelname)s - %(message)s')
+
+    logging.debug("Starting the VAE inference script.")
+    args = parser.parse_args()
+    logging.debug(f"Parsed arguments: {args}")
+
+    try:
+        torch.manual_seed(args.seed)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logging.debug(f"Using device: {device}")
+
+        logging.debug(f"Loading VAE model from {args.resume}.")
+        model = BetaVAE(z_dim=args.zdim, nc=1).to(device)
+        if os.path.isfile(args.resume):
+            checkpoint = torch.load(args.resume, map_location=device)
+            model.load_state_dict(checkpoint['state_dict'])
+            logging.debug("Checkpoint loaded.")
+        else:
+            logging.error(f"Checkpoint not found at {args.resume}")
+            raise RuntimeError("Checkpoint not found.")
+
+        if (args.mode == 'encode') or (args.mode == 'both'):
+            logging.debug("Starting encoding process...")
+            test_loader = load_dataset_test(args.data_path, args.batch_size)
+            logging.debug(f"Loaded test dataset from {args.data_path}")
+            for batch_idx, (xL, xR) in enumerate(test_loader):
+                xL = xL.to(device)
+                xR = xR.to(device)
+                z_distribution = model._encode(xL, xR)
+                save_data = {'z_distribution': z_distribution.detach().cpu().numpy()}
+                io.savemat(os.path.join(args.z_path, f'save_z{batch_idx}.mat'), save_data)
+                logging.debug(f"Encoded batch {batch_idx}")
+
+        if (args.mode == 'decode') or (args.mode == 'both'):
+            logging.debug("Starting decoding process...")
+            filelist = [f for f in os.listdir(args.z_path) if f.split('_')[0] == 'save']
+            logging.debug(f"Filelist: {filelist}")
+            for batch_idx, filename in enumerate(filelist):
+                logging.debug(f"Decoding file {filename}")
+                z_dist = io.loadmat(os.path.join(args.z_path, f'save_z{batch_idx}.mat'))
+                z_dist = z_dist['z_distribution']
+                mu = z_dist[:, :args.zdim]
+                z = torch.tensor(mu).to(device)
+                x_recon_L, x_recon_R = model._decode(z)
+                save_image_mat(x_recon_R, x_recon_L, args.img_path, batch_idx)
+                logging.debug(f"Decoded and saved batch {batch_idx}")
+
+    except Exception as e:
+        logging.error(f"An error occurred: {e}")
+        raise
+
+if __name__ == "__main__":
+    main()

call_VAE_inference_example.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+SECONDS=0
+
+# MSC
+subjs=(01 02 03 04 05 06 07 08 09 10)
+parcelcount=(602 567 620 616 633 580 628 710 613 649)
+zdim=2
+checkpoint="./VAE_Model/Checkpoint/checkpoint49_2024-03-28_Zdim_2_Vae-beta_20.0_Lr_0.0001_Batch-size_128_washu120_subsample10_train100_val10.pth.tar"
+for i in "${!subjs[@]}"; do 
+    subj="${subjs[$i]}" 
+    curr_parcel_count="${parcelcount[$i]}"
+    echo $curr_parcel_count
+    namestr="sub-MSC${subj}_sub-MSC${subj}Parcel"
+
+    python3 VAE_inference_example.py --data-path ./data/$namestr --zdim $zdim \
+        --resume  "${checkpoint}" \
+        --z-path './result/latent/'$namestr'_Zdim'$zdim --mode 'encode' --batch-size $curr_parcel_count
+
+    echo "The command took $SECONDS seconds." 
+done

fMRIVAE_Model.py

@@ -0,0 +1,140 @@
+"""model.py"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+from collections import Iterable
+from torch.autograd import Variable
+
+class BetaVAE(nn.Module):
+    """Model proposed in original beta-VAE paper(Higgins et al, ICLR, 2017)."""
+
+    def __init__(self, z_dim=64, nc=1, cirpad_dire=(False, True)):
+        super(BetaVAE, self).__init__()
+        self.z_dim = z_dim
+        self.nc = nc
+        self.cirpad_dire = cirpad_dire
+   
+        self.ocs = [64, 128, 128, 256, 256]
+        self.nLays = len(self.ocs)
+        self.topW = int(192/2**self.nLays)
+
+        # encoder
+        self.ConvL = nn.Conv2d(1,int(self.ocs[0]/2),8,2,0)  # pad=3, only in forward
+        self.ConvR = nn.Conv2d(1,int(self.ocs[0]/2),8,2,0)  # pad=3, only in forward # B, 128, 96, 96
+        self.EncConvs = nn.ModuleList([nn.Conv2d(self.ocs[i-1], self.ocs[i], 4, 2, 0) for i in range(1, self.nLays)]) # pad=1 only in forward     
+        self.fc1 = nn.Linear(self.ocs[-1]*self.topW**2, z_dim*2)
+
+        # decoder
+        self.fc2 = nn.Linear(z_dim, self.ocs[-1]*self.topW**2)
+        self.DecConvs = nn.ModuleList([nn.ConvTranspose2d(self.ocs[i], self.ocs[i-1], 4, 2, 3) for i in range(4,0,-1)]) # pad=1; dilation * (kernel_size - 1) - padding = 6  (later in forward)
+        self.tConvL = nn.ConvTranspose2d(int(self.ocs[0]/2), nc, 8, 2, 9) # pad=3 later; dilation * (kernel_size - 1) - padding = 4  (later in forward)
+        self.tConvR = nn.ConvTranspose2d(int(self.ocs[0]/2), nc, 8, 2, 9) # pad=3 later
+
+        self.relu = nn.ReLU(inplace=True)
+
+        self.weight_init()
+
+    def cirpad(self, x, padding, cirpad_dire):
+        # x            is    input
+        # padding      is    the size of pading
+        # cirpad_dire  is    (last_dim_pad, second_to_last_dim_pad)
+        
+        # >>> t4d = torch.empty(3, 3, 4, 2)
+        # >>> p2d = (1, 1, 2, 2) # pad last dim by (1, 1) and 2nd to last by (2, 2)
+        # >>> out = F.pad(t4d, p2d, "constant", 0)
+        # >>> print(out.size())
+        # torch.Size([3, 3, 8, 4])
+
+        # last dim
+        if cirpad_dire[0] is True:
+            x = F.pad(x, (padding, padding, 0, 0), 'circular')
+        else:
+            x = F.pad(x, (padding, padding, 0, 0), "constant", 0)
+        
+        # second last dim
+        if cirpad_dire[1] is True:
+            x = F.pad(x, (0, 0, padding, padding), 'circular')
+        else:
+            x = F.pad(x, (0, 0, padding, padding), "constant", 0)
+            
+        return x
+
+
+    def weight_init(self):
+        for block in self._modules:
+            if isinstance(self._modules[block], Iterable):
+                for m in self._modules[block]:
+                    m.apply(kaiming_init)
+            else:
+                self._modules[block].apply(kaiming_init)
+
+    def _encode(self, xL, xR):
+        xL = self.cirpad(xL, 3, self.cirpad_dire) 
+        xR = self.cirpad(xR, 3, self.cirpad_dire) 
+        x = torch.cat((self.ConvL(xL), self.ConvR(xR)), 1)
+        x = self.relu(x)
+        for lay in range(self.nLays-1):
+            x = self.cirpad(x, 1, self.cirpad_dire) 
+            x = self.relu(self.EncConvs[lay](x))
+        x = x.view(-1, self.ocs[-1]*self.topW*self.topW)
+        x = self.fc1(x)
+        return x
+
+    def _decode(self, z):
+        x = self.relu(self.fc2(z).view(-1 , self.ocs[-1], self.topW, self.topW))
+        for lay in range(self.nLays-1):
+            x = self.cirpad(x, 1, self.cirpad_dire) 
+            x = self.relu(self.DecConvs[lay](x))
+
+        xL, xR = torch.chunk(x, 2, dim=1)
+
+        xrL = self.tConvL(self.cirpad(xL, 3, self.cirpad_dire))
+
+        xrR = self.tConvR(self.cirpad(xR, 3, self.cirpad_dire))
+        return xrL, xrR
+
+    def reparametrize(self, mu, logvar):
+        std = logvar.div(2).exp()
+        eps = Variable(std.data.new(std.size()).normal_())
+        return mu + std*eps
+
+
+    def forward(self, xL, xR):
+        distributions = self._encode(xL, xR)
+        mu = distributions[:, :self.z_dim]
+        logvar = distributions[:, self.z_dim:]
+        z = self.reparametrize(mu, logvar)
+        x_recon_L, x_recon_R = self._decode(z)
+        return x_recon_L, x_recon_R, mu, logvar
+
+def kaiming_init(m):
+    if isinstance(m, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)): # Shall we apply init to ConvTranspose2d?
+        init.kaiming_normal_(m.weight)
+        if m.bias is not None:
+            m.bias.data.fill_(0)
+    elif isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d)):
+        m.weight.data.fill_(1)
+        if m.bias is not None:
+            m.bias.data.fill_(0)
+
+def normal_init(m, mean, std):
+    if isinstance(m, (nn.Linear, nn.Conv2d)):
+        m.weight.data.normal_(mean, std)
+        if m.bias.data is not None:
+            m.bias.data.zero_()
+    elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
+        m.weight.data.fill_(1)
+        if m.bias.data is not None:
+            m.bias.data.zero_()
+
+#if __name__ == "__main__":
+#    m = BetaVAE_H()
+#    a=torch.ones(1,1,192,192)
+#    out1, out2, _, _ = m(a,a)
+#    print(out1.size())
+
+
+
+

requirements.txt

@@ -0,0 +1,73 @@
+asttokens==2.4.1
+backcall==0.2.0
+certifi==2023.11.17
+charset-normalizer==3.3.2
+comm==0.2.0
+contourpy==1.1.1
+cycler==0.12.1
+debugpy==1.8.0
+decorator==5.1.1
+executing==2.0.1
+filelock==3.13.1
+fonttools==4.47.0
+fsspec==2023.12.2
+h5py==3.10.0
+idna==3.6
+importlib-metadata==7.0.0
+importlib-resources==6.1.1
+ipykernel==6.27.1
+ipython==8.12.3
+jedi==0.19.1
+Jinja2==3.1.2
+jupyter_client==8.6.0
+jupyter_core==5.5.1
+kiwisolver==1.4.5
+MarkupSafe==2.1.3
+matplotlib==3.7.4
+matplotlib-inline==0.1.6
+mpmath==1.3.0
+nest-asyncio==1.5.8
+networkx==3.1
+nibabel==5.2.0
+numpy==1.24.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.18.1
+nvidia-nvjitlink-cu12==12.3.101
+nvidia-nvtx-cu12==12.1.105
+packaging==23.2
+parso==0.8.3
+pexpect==4.9.0
+pickleshare==0.7.5
+Pillow==10.1.0
+platformdirs==4.1.0
+prompt-toolkit==3.0.43
+psutil==5.9.7
+ptyprocess==0.7.0
+pure-eval==0.2.2
+Pygments==2.17.2
+pyparsing==3.1.1
+python-dateutil==2.8.2
+pyzmq==25.1.2
+requests==2.31.0
+scipy==1.10.1
+six==1.16.0
+stack-data==0.6.3
+sympy==1.12
+torch==2.1.2+cu118
+torchaudio==2.1.2+cu118
+torchvision==0.16.2+cu118
+tornado==6.4
+traitlets==5.14.0
+triton==2.1.0
+typing_extensions==4.9.0
+urllib3==2.1.0
+wcwidth==0.2.12
+zipp==3.17.0

utils.py

@@ -0,0 +1,80 @@
+import h5py
+import torch
+import torch.utils.data as data
+import torch.multiprocessing
+import scipy.io as sio
+from torch.nn import functional as F
+# torch.multiprocessing.set_start_method('spawn')
+
+class H5Dataset(data.Dataset):
+    def __init__(self, H5Path):
+        super(H5Dataset, self).__init__()
+        self.H5File = h5py.File(H5Path,'r')       
+        self.LeftData = self.H5File['LeftData']
+        self.RightData = self.H5File['RightData']
+        #self.LeftMask = self.H5File['LeftMask'][:] # update 2024.01.11 Masks loaded separately
+        #self.RightMask = self.H5File['RightMask'][:]
+
+    def __getitem__(self, index):
+        return (torch.from_numpy(self.LeftData[index,:,:,:]).float(),
+                torch.from_numpy(self.RightData[index,:,:,:]).float())
+ 
+    def __len__(self):
+        return self.LeftData.shape[0]
+
+def save_image_mat(img_r, img_l, result_path, idx):
+    save_data = {}
+    save_data['recon_L'] = img_l.detach().cpu().numpy()
+    save_data['recon_R'] = img_r.detach().cpu().numpy()
+    sio.savemat(result_path+'img{}.mat'.format(idx), save_data)
+
+def load_dataset(data_path, batch_size):
+    kwargs = {'num_workers': 1, 'pin_memory': True} if torch.cuda.is_available() else {}
+    train_dir = data_path + '_train.h5'
+    val_dir = data_path + '_val.h5'
+    train_set = H5Dataset(train_dir)
+    val_set = H5Dataset(val_dir)
+    train_loader = torch.utils.data.DataLoader(train_set,batch_size=batch_size, shuffle=False, **kwargs)
+    val_loader = torch.utils.data.DataLoader(val_set,batch_size=batch_size, shuffle=False, **kwargs)
+    return train_loader, val_loader
+
+def load_dataset_test(data_path, batch_size):
+    kwargs = {'num_workers': 1, 'pin_memory': True} if torch.cuda.is_available() else {}
+    test_dir = data_path + '.h5'
+    test_set = H5Dataset(test_dir)
+    test_loader = torch.utils.data.DataLoader(test_set,batch_size=batch_size, shuffle=False, **kwargs)
+    return test_loader
+
+# loss function # update 20240109 mask out zeros
+def loss_function(xL, xR, x_recon_L, x_recon_R, mu, logvar, beta, left_mask, right_mask):
+
+    Image_Size=xL.size(3)
+
+    beta/=Image_Size**2
+    
+    # print('====> Image_Size: {} Beta: {:.8f}'.format(Image_Size, beta))
+
+    # R_batch_size=xR.size(0)
+    # Tutorial on VAE Page-14
+    # log[P(X|z)] = C - \frac{1}{2} ||X-f(z)||^2 // \sigma^2 
+    #             = C - \frac{1}{2} \sum_{i=1}^{N} ||X^{(i)}-f(z^{(i)}||^2 // \sigma^2
+    #             = C - \farc{1}{2} N * F.mse_loss(Xhat-Xtrue) // \sigma^2
+    # log[P(X|z)]-C = - \frac{1}{2}*2*192*192//\sigma^2 * F.mse_loss
+    # Therefore, vae_beta = \frac{1}{36864//\sigma^2}
+        
+    # mask out zeros
+    valid_mask_L = xL!=0
+    valid_mask_R = xR!=0 
+
+    if left_mask is not None:
+        valid_mask_L = valid_mask_L & (left_mask.detach().to(torch.int32)==1)
+        valid_mask_R = valid_mask_R & (right_mask.detach().to(torch.int32)==1)
+    
+    MSE_L = F.mse_loss(x_recon_L*valid_mask_L, xL*valid_mask_L, size_average=True)
+    MSE_R = F.mse_loss(x_recon_R*valid_mask_R, xR *valid_mask_R, size_average=True)  
+
+    # KLD is averaged across batch-samples
+    KLD = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(1).mean()
+
+    return KLD * beta + MSE_L + MSE_R
+