commit

adrd/__init__.py

@@ -1,22 +0,0 @@
-__version__ = '0.0.1'
-
-from . import nn
-from . import model
-
-# # load pretrained transformer
-# pretrained_transformer = model.Transformer.from_ckpt('{}/ckpt/ckpt.pt'.format(__path__[0]))
-# from . import shap_adrd
-# from .model import DynamicCalibratedClassifier
-# from .model import StaticCalibratedClassifier
-
-# load fitted transformer and calibrated wrapper
-# try:
-#     fitted_resnet3d = model.CNNResNet3DWithLinearClassifier.from_ckpt('{}/ckpt/ckpt_img_072523.pt'.format(__path__[0]))
-#     fitted_calibrated_classifier_nonimg = StaticCalibratedClassifier.from_ckpt(
-#         filepath_state_dict = '{}/ckpt/static_calibrated_classifier_073023.pkl'.format(__path__[0]),
-#         filepath_wrapped_model = '{}/ckpt/ckpt_080823.pt'.format(__path__[0]),
-#     )
-#     fitted_transformer_nonimg = fitted_calibrated_classifier_nonimg.model
-#     shap_explainer = shap_adrd.SamplingExplainer(fitted_transformer_nonimg)
-# except:
-#     print('Fail to load checkpoints.')

adrd/_ds/lddl.py

@@ -1,71 +0,0 @@
-from typing import Any, Self, overload
-
-
-class lddl:
-    ''' ... '''
-    def __init__(self) -> None:
-        ''' ... '''
-        self.dat_ld: list[dict[str, Any]] = None
-        self.dat_dl: dict[str, list[Any]] = None
-
-    @overload
-    def __getitem__(self, idx: int) -> dict[str, Any]: ...
-
-    @overload
-    def __getitem__(self, idx: str) -> list[Any]: ...
-
-    def __getitem__(self, idx: str | int) -> list[Any] | dict[str, Any]:
-        ''' ... '''
-        if isinstance(idx, str):
-            return self.dat_dl[idx]
-        elif isinstance(idx, int):
-            return self.dat_ld[idx]
-        else:
-            raise TypeError('Unexpected key type: {}'.format(type(idx)))
-
-    @classmethod
-    def from_ld(cls, dat: list[dict[str, Any]]) -> Self:
-        ''' Construct from list of dicts. '''
-        obj = cls()
-        obj.dat_ld = dat
-        obj.dat_dl = {k: [dat[i][k] for i in range(len(dat))] for k in dat[0]}
-        return obj
-
-    @classmethod
-    def from_dl(cls, dat: dict[str, list[Any]]) -> Self:
-        ''' Construct from dict of lists. '''
-        obj = cls()
-        obj.dat_ld = [dict(zip(dat, v)) for v in zip(*dat.values())]
-        obj.dat_dl = dat
-        return obj
-    
-
-if __name__ == '__main__':
-    ''' for testing purpose only '''
-    dl = {
-        'a': [0, 1, 2],
-        'b': [3, 4, 5],
-    }
-
-    ld = [
-        {'a': 0, 'b': 1, 'c': 2},
-        {'a': 3, 'b': 4, 'c': 5},
-    ]
-
-    # test constructing from ld
-    dat_ld = lddl.from_ld(ld)
-    print(dat_ld.dat_ld)
-    print(dat_ld.dat_dl)
-
-    # test constructing from dl
-    dat_dl = lddl.from_dl(dl)
-    print(dat_dl.dat_ld)
-    print(dat_dl.dat_dl)
-
-    # test __getitem__
-    print(dat_dl['a'])
-    print(dat_dl[0])
-
-    # mouse hover to check if type hints are correct
-    v = dat_dl['a']
-    v = dat_dl[0]

adrd/model/__init__.py

@@ -1,6 +0,0 @@
-from .adrd_model import ADRDModel
-from .imaging_model import ImagingModel
-from .train_resnet import TrainResNet
-# from .transformer import Transformer
-from .calibration import DynamicCalibratedClassifier
-from .calibration import StaticCalibratedClassifier

adrd/model/adrd_model.py

@@ -1,976 +0,0 @@
-__all__ = ['ADRDModel']
-
-import wandb
-import torch
-from torch.utils.data import DataLoader
-import numpy as np
-import tqdm
-import multiprocessing
-from sklearn.base import BaseEstimator
-from sklearn.utils.validation import check_is_fitted
-from sklearn.model_selection import train_test_split
-from scipy.special import expit
-from copy import deepcopy
-from contextlib import suppress
-from typing import Any, Self, Type
-from functools import wraps
-from tqdm import tqdm
-Tensor = Type[torch.Tensor]
-Module = Type[torch.nn.Module]
-
-# for DistributedDataParallel
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from torch.nn.parallel import DistributedDataParallel as DDP
-
-from .. import nn
-from ..nn import Transformer
-from ..utils import TransformerTrainingDataset, TransformerBalancedTrainingDataset, TransformerValidationDataset, TransformerTestingDataset, Transformer2ndOrderBalancedTrainingDataset
-from ..utils.misc import ProgressBar
-from ..utils.misc import get_metrics_multitask, print_metrics_multitask
-from ..utils.misc import convert_args_kwargs_to_kwargs
-
-
-def _manage_ctx_fit(func):
-    ''' ... '''
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        # format arguments
-        kwargs = convert_args_kwargs_to_kwargs(func, args, kwargs)
-
-        if kwargs['self']._device_ids is None:
-            return func(**kwargs)
-        else:
-            # change primary device
-            default_device = kwargs['self'].device
-            kwargs['self'].device = kwargs['self']._device_ids[0]
-            rtn = func(**kwargs)
-            kwargs['self'].to(default_device)
-            return rtn
-    return wrapper
-
-
-class ADRDModel(BaseEstimator):
-    """Primary model class for ADRD framework.
-
-    The ADRDModel encapsulates the core pipeline of the ADRD framework, 
-    permitting users to train and validate with the provided data. Designed for 
-    user-friendly operation, the ADRDModel is derived from 
-    ``sklearn.base.BaseEstimator``, ensuring compliance with the well-established 
-    API design conventions of scikit-learn.
-    """    
-    def __init__(self,
-        src_modalities: dict[str, dict[str, Any]],
-        tgt_modalities: dict[str, dict[str, Any]],
-        label_fractions: dict[str, float],
-        d_model: int = 32,
-        nhead: int = 1,
-        num_encoder_layers: int = 1,
-        num_decoder_layers: int = 1,
-        num_epochs: int = 32,
-        batch_size: int = 8,
-        batch_size_multiplier: int = 1,
-        lr: float = 1e-2,
-        weight_decay: float = 0.0,
-        beta: float = 0.9999,
-        gamma: float = 2.0,
-        criterion: str | None = None,
-        device: str = 'cpu',
-        cuda_devices: list = [1],
-        img_net: str | None = None,
-        imgnet_layers: int | None = 2,
-        img_size: int | None = 128,
-        fusion_stage: str = 'middle',
-        patch_size: int | None = 16,
-        imgnet_ckpt: str | None = None,
-        train_imgnet: bool = False,
-        ckpt_path: str = './adrd_tool/adrd/dev/ckpt/ckpt.pt',
-        load_from_ckpt: bool = False,
-        save_intermediate_ckpts: bool = False,
-        data_parallel: bool = False,
-        verbose: int = 0,
-        wandb_ = 0,
-        balanced_sampling: bool = False,
-        label_distribution: dict = {},
-        ranking_loss: bool = False,
-        _device_ids: list | None = None,
-
-        _dataloader_num_workers: int = 4,
-        _amp_enabled: bool = False,
-    ) -> None:  
-        """Create a new ADRD model.
-
-        :param src_modalities: _description_
-        :type src_modalities: dict[str, dict[str, Any]]
-        :param tgt_modalities: _description_
-        :type tgt_modalities: dict[str, dict[str, Any]]
-        :param label_fractions: _description_
-        :type label_fractions: dict[str, float]
-        :param d_model: _description_, defaults to 32
-        :type d_model: int, optional
-        :param nhead: number of transformer heads, defaults to 1
-        :type nhead: int, optional
-        :param num_encoder_layers: number of encoder layers, defaults to 1
-        :type num_encoder_layers: int, optional
-        :param num_decoder_layers: number of decoder layers, defaults to 1
-        :type num_decoder_layers: int, optional
-        :param num_epochs: number of training epochs, defaults to 32
-        :type num_epochs: int, optional
-        :param batch_size: batch size, defaults to 8
-        :type batch_size: int, optional
-        :param batch_size_multiplier: _description_, defaults to 1
-        :type batch_size_multiplier: int, optional
-        :param lr: learning rate, defaults to 1e-2
-        :type lr: float, optional
-        :param weight_decay: _description_, defaults to 0.0
-        :type weight_decay: float, optional
-        :param beta: _description_, defaults to 0.9999
-        :type beta: float, optional
-        :param gamma: The focusing parameter for the focal loss. Higher values of gamma make easy-to-classify examples contribute less to the loss relative to hard-to-classify examples. Must be non-negative., defaults to 2.0
-        :type gamma: float, optional
-        :param criterion: The criterion to select the best model, defaults to None
-        :type criterion: str | None, optional
-        :param device: 'cuda' or 'cpu', defaults to 'cpu'
-        :type device: str, optional
-        :param cuda_devices: A list of gpu numbers to data parallel training. The device must be set to 'cuda' and data_parallel must be set to True, defaults to [1]
-        :type cuda_devices: list, optional
-        :param img_net: _description_, defaults to None
-        :type img_net: str | None, optional
-        :param imgnet_layers: _description_, defaults to 2
-        :type imgnet_layers: int | None, optional
-        :param img_size: _description_, defaults to 128
-        :type img_size: int | None, optional
-        :param fusion_stage: _description_, defaults to 'middle'
-        :type fusion_stage: str, optional
-        :param patch_size: _description_, defaults to 16
-        :type patch_size: int | None, optional
-        :param imgnet_ckpt: _description_, defaults to None
-        :type imgnet_ckpt: str | None, optional
-        :param train_imgnet: Set to True to finetune the img_net backbone, defaults to False
-        :type train_imgnet: bool, optional
-        :param ckpt_path: The model checkpoint point path, defaults to './adrd_tool/adrd/dev/ckpt/ckpt.pt'
-        :type ckpt_path: str, optional
-        :param load_from_ckpt: Set to True to load the model weights from checkpoint ckpt_path, defaults to False
-        :type load_from_ckpt: bool, optional
-        :param save_intermediate_ckpts: Set to True to save intermediate model checkpoints, defaults to False
-        :type save_intermediate_ckpts: bool, optional
-        :param data_parallel: Set to True to enable data parallel trsining, defaults to False
-        :type data_parallel: bool, optional
-        :param verbose: _description_, defaults to 0
-        :type verbose: int, optional
-        :param wandb_: Set to 1 to track the loss and accuracy curves on wandb, defaults to 0
-        :type wandb_: int, optional
-        :param balanced_sampling: _description_, defaults to False
-        :type balanced_sampling: bool, optional
-        :param label_distribution: _description_, defaults to {}
-        :type label_distribution: dict, optional
-        :param ranking_loss: _description_, defaults to False
-        :type ranking_loss: bool, optional
-        :param _device_ids: _description_, defaults to None
-        :type _device_ids: list | None, optional
-        :param _dataloader_num_workers: _description_, defaults to 4
-        :type _dataloader_num_workers: int, optional
-        :param _amp_enabled: _description_, defaults to False
-        :type _amp_enabled: bool, optional
-        """
-        # for multiprocessing
-        self._rank = 0
-        self._lock = None
-
-        # positional parameters
-        self.src_modalities = src_modalities
-        self.tgt_modalities = tgt_modalities
-
-        # training parameters
-        self.label_fractions = label_fractions
-        self.d_model = d_model
-        self.nhead = nhead
-        self.num_encoder_layers = num_encoder_layers
-        self.num_decoder_layers = num_decoder_layers
-        self.num_epochs = num_epochs
-        self.batch_size = batch_size
-        self.batch_size_multiplier = batch_size_multiplier
-        self.lr = lr
-        self.weight_decay = weight_decay
-        self.beta = beta
-        self.gamma = gamma
-        self.criterion = criterion
-        self.device = device
-        self.cuda_devices = cuda_devices
-        self.img_net = img_net
-        self.patch_size = patch_size
-        self.img_size = img_size
-        self.fusion_stage = fusion_stage
-        self.imgnet_ckpt = imgnet_ckpt
-        self.imgnet_layers = imgnet_layers
-        self.train_imgnet = train_imgnet
-        self.ckpt_path = ckpt_path
-        self.load_from_ckpt = load_from_ckpt
-        self.save_intermediate_ckpts = save_intermediate_ckpts
-        self.data_parallel = data_parallel
-        self.verbose = verbose
-        self.label_distribution = label_distribution
-        self.wandb_ = wandb_
-        self.balanced_sampling = balanced_sampling
-        self.ranking_loss = ranking_loss
-        self._device_ids = _device_ids
-        self._dataloader_num_workers = _dataloader_num_workers
-        self._amp_enabled = _amp_enabled
-        self.scaler = torch.cuda.amp.GradScaler()
-        # self._init_net()
-
-    @_manage_ctx_fit
-    def fit(self, x_trn, x_vld, y_trn, y_vld, img_train_trans=None, img_vld_trans=None, img_mode=0) -> Self:
-    # def fit(self, x, y) -> Self:
-        ''' ... '''
-        
-        # start a new wandb run to track this script
-        if self.wandb_ == 1:
-            wandb.init(
-                # set the wandb project where this run will be logged
-                project="ADRD_main",
-                
-                # track hyperparameters and run metadata
-                config={
-                "Loss": 'Focalloss',
-                "ranking_loss": self.ranking_loss,
-                "img architecture": self.img_net,
-                "EMB": "ALL_SEQ",
-                "epochs": self.num_epochs,
-                "d_model": self.d_model,
-                # 'positional encoding': 'Diff PE',
-                'Balanced Sampling': self.balanced_sampling,
-                'Shared CNN': 'Yes',
-                }
-            )
-            wandb.run.log_code(".")
-        else:
-            wandb.init(mode="disabled") 
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-        # print(img_train_trans)
-        # initialize neural network
-        print(self.criterion)
-        print(f"Ranking loss: {self.ranking_loss}")
-        print(f"Batch size: {self.batch_size}")
-        print(f"Batch size multiplier: {self.batch_size_multiplier}")
-
-        if img_mode in [0,1,2]:
-            for k, info in self.src_modalities.items():
-                if info['type'] == 'imaging':
-                    if 'densenet' in self.img_net.lower() and 'emb' not in self.img_net.lower():
-                        info['shape'] = (1,) + self.img_size
-                        info['img_shape'] = (1,) + self.img_size
-                    elif 'emb' not in self.img_net.lower():
-                        info['shape'] = (1,) + (self.img_size,) * 3
-                        info['img_shape'] = (1,) + (self.img_size,) * 3
-                    elif 'swinunetr' in self.img_net.lower():
-                        info['shape'] = (1, 768, 4, 4, 4)
-                        info['img_shape'] = (1, 768, 4, 4, 4)
-        
-        
-        
-        self._init_net()
-        ldr_trn, ldr_vld = self._init_dataloader(x_trn, x_vld, y_trn, y_vld, img_train_trans=img_train_trans, img_vld_trans=img_vld_trans)
-
-        # initialize optimizer and scheduler
-        if not self.load_from_ckpt:
-            self.optimizer = self._init_optimizer()
-        self.scheduler = self._init_scheduler(self.optimizer)
-
-        # gradient scaler for AMP
-        if self._amp_enabled: 
-            self.scaler = torch.cuda.amp.GradScaler()
-
-        # initialize the focal losses 
-        self.loss_fn = {}
-
-        for k in self.tgt_modalities:
-            if self.label_fractions[k] >= 0.3:
-                alpha = -1
-            else:
-                alpha = pow((1 - self.label_fractions[k]), 2)
-            # alpha = -1
-            self.loss_fn[k] = nn.SigmoidFocalLoss(
-                alpha = alpha,
-                gamma = self.gamma,
-                reduction = 'none'
-            )
-
-        # to record the best validation performance criterion
-        if self.criterion is not None:
-            best_crit = None
-            best_crit_AUPR = None
-
-        # progress bar for epoch loops
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch = tqdm.tqdm(
-                    desc = 'Rank {:02d}'.format(self._rank),
-                    total = self.num_epochs,
-                    position = self._rank,
-                    ascii = True,
-                    leave = False,
-                    bar_format='{l_bar}{r_bar}'
-                )
-
-        self.skip_embedding = {}
-        for k, info in self.src_modalities.items():
-            # if info['type'] == 'imaging':
-            #     if not self.img_net:
-            #         self.skip_embedding[k] = True
-            # else:
-            self.skip_embedding[k] = False
-
-        self.grad_list = []
-        # Define a hook function to print and store the gradient of a layer
-        def print_and_store_grad(grad):
-            self.grad_list.append(grad)
-            # print(grad)
-            
-       
-        # initialize the ranking loss
-        self.lambda_coeff = 0.005
-        self.margin = 0.25
-        self.margin_loss = torch.nn.MarginRankingLoss(reduction='sum', margin=self.margin)
-        
-        # training loop
-        for epoch in range(self.start_epoch, self.num_epochs):
-            met_trn = self.train_one_epoch(ldr_trn, epoch)
-            met_vld = self.validate_one_epoch(ldr_vld, epoch)
-            
-            print(self.ckpt_path.split('/')[-1])
-
-            # save the model if it has the best validation performance criterion by far
-            if self.criterion is None: continue
-            
-            # is current criterion better than previous best?
-            curr_crit = np.mean([met_vld[i][self.criterion] for i in range(len(self.tgt_modalities))])
-            curr_crit_AUPR = np.mean([met_vld[i]["AUC (PR)"] for i in range(len(self.tgt_modalities))])
-            # AUROC
-            if best_crit is None or np.isnan(best_crit):
-                is_better = True
-            elif self.criterion == 'Loss' and best_crit >= curr_crit:
-                is_better = True
-            elif self.criterion != 'Loss' and best_crit <= curr_crit :
-                is_better = True
-            else:
-                is_better = False
-
-            # AUPR
-            if best_crit_AUPR is None or np.isnan(best_crit_AUPR):
-                is_better_AUPR = True
-            elif best_crit_AUPR <= curr_crit_AUPR :
-                is_better_AUPR = True
-            else:
-                is_better_AUPR = False
-            # update best criterion
-            if is_better_AUPR:
-                best_crit_AUPR = curr_crit_AUPR
-                if self.save_intermediate_ckpts:
-                    print(f"Saving the model to {self.ckpt_path[:-3]}_AUPR.pt...")
-                    self.save(self.ckpt_path[:-3]+"_AUPR.pt", epoch)
-            if is_better:
-                best_crit = curr_crit
-                best_state_dict = deepcopy(self.net_.state_dict())
-                if self.save_intermediate_ckpts:
-                    print(f"Saving the model to {self.ckpt_path}...")
-                    self.save(self.ckpt_path, epoch)
-
-            if self.verbose > 2:
-                print('Best {}: {}'.format(self.criterion, best_crit))
-                print('Best {}: {}'.format('AUC (PR)', best_crit_AUPR))
-
-            if self.verbose == 1:
-                with self._lock if self._lock is not None else suppress():
-                    pbr_epoch.update(1)
-                    pbr_epoch.refresh()
-
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch.close()
-
-        return self
-    
-    def train_one_epoch(self, ldr_trn, epoch):
-        # progress bar for batch loops
-        if self.verbose > 1: 
-            pbr_batch = ProgressBar(len(ldr_trn.dataset), 'Epoch {:03d} (TRN)'.format(epoch))
-
-        # set model to train mode
-        torch.set_grad_enabled(True)
-        self.net_.train()
-
-        scores_trn, y_true_trn, y_mask_trn = [], [], []
-        losses_trn = [[] for _ in self.tgt_modalities]
-        iters = len(ldr_trn)
-        for n_iter, (x_batch, y_batch, mask, y_mask) in enumerate(ldr_trn):
-           
-            # mount data to the proper device
-            x_batch = {k: x_batch[k].to(self.device) for k in x_batch}
-            y_batch = {k: y_batch[k].to(torch.float).to(self.device) for k in y_batch}
-            mask = {k: mask[k].to(self.device) for k in mask}
-            y_mask = {k: y_mask[k].to(self.device) for k in y_mask}
-            
-            with torch.autocast(
-                device_type = 'cpu' if self.device == 'cpu' else 'cuda',
-                dtype = torch.bfloat16 if self.device == 'cpu' else torch.float16,
-                enabled = self._amp_enabled,
-            ):
-
-                outputs = self.net_(x_batch, mask, skip_embedding=self.skip_embedding)
-                
-                # calculate multitask loss
-                loss = 0
-
-                # for initial 10 epochs, only the focal loss is used for stable training
-                if self.ranking_loss:
-                    if epoch < 10:
-                        loss = 0
-                    else:
-                        for i, k in enumerate(self.tgt_modalities):
-                           for ii, kk in enumerate(self.tgt_modalities):
-                               if ii>i:
-                                   pairs = (y_mask[k] == 1) & (y_mask[kk] == 1)
-                                   total_elements = (torch.abs(y_batch[k][pairs]-y_batch[kk][pairs])).sum()
-                                   if  total_elements != 0:
-                                       loss += self.lambda_coeff * (self.margin_loss(torch.sigmoid(outputs[k])[pairs],torch.sigmoid(outputs[kk][pairs]),y_batch[k][pairs]-y_batch[kk][pairs]))/total_elements
-
-                for i, k in enumerate(self.tgt_modalities):
-                    loss_task = self.loss_fn[k](outputs[k], y_batch[k])
-                    msk_loss_task = loss_task * y_mask[k]
-                    msk_loss_mean = msk_loss_task.sum() / y_mask[k].sum()
-                    # msk_loss_mean = msk_loss_task.sum()
-                    loss += msk_loss_mean
-                    losses_trn[i] += msk_loss_task.detach().cpu().numpy().tolist()
-
-            # backward
-            loss = loss / self.batch_size_multiplier
-            if self._amp_enabled:
-                self.scaler.scale(loss).backward()
-            else:
-                loss.backward()
-
-            if len(self.grad_list) > 0:
-                print(len(self.grad_list), len(self.grad_list[-1]))
-                print(f"Gradient at {n_iter}: {self.grad_list[-1][0]}")
-           
-            # print("img_MRI_T1_1 ", self.net_.modules_emb_src.img_MRI_T1_1.img_model.features[0].weight)
-            # print("img_MRI_T1_1 ", self.net_.modules_emb_src.img_MRI_T1_1.downsample[0].weight)
-            
-            # update parameters
-            if n_iter != 0 and n_iter % self.batch_size_multiplier == 0:
-                if self._amp_enabled:
-                    self.scaler.step(self.optimizer)
-                    self.scaler.update()
-                    self.optimizer.zero_grad()
-                else: 
-                    self.optimizer.step()
-                    self.optimizer.zero_grad()
-                
-                # set self.scheduler
-                self.scheduler.step(epoch + n_iter / iters)
-
-            ''' TODO: change array to dictionary later '''
-            outputs = torch.stack(list(outputs.values()), dim=1)
-            y_batch = torch.stack(list(y_batch.values()), dim=1)
-            y_mask = torch.stack(list(y_mask.values()), dim=1)
-
-            # save outputs to evaluate performance later
-            scores_trn.append(outputs.detach().to(torch.float).cpu())
-            y_true_trn.append(y_batch.cpu())
-            y_mask_trn.append(y_mask.cpu())
-            
-            # update progress bar
-            if self.verbose > 1:
-                batch_size = len(next(iter(x_batch.values())))
-                pbr_batch.update(batch_size, {})
-                pbr_batch.refresh()
-
-            # clear cuda cache
-            if "cuda" in self.device:
-                torch.cuda.empty_cache()
-
-        # for better tqdm progress bar display
-        if self.verbose > 1:
-            pbr_batch.close()
-
-        # calculate and print training performance metrics
-        scores_trn = torch.cat(scores_trn)
-        y_true_trn = torch.cat(y_true_trn)
-        y_mask_trn = torch.cat(y_mask_trn)
-        y_pred_trn = (scores_trn > 0).to(torch.int)
-        y_prob_trn = torch.sigmoid(scores_trn)
-        met_trn = get_metrics_multitask(
-            y_true_trn.numpy(),
-            y_pred_trn.numpy(),
-            y_prob_trn.numpy(),
-            y_mask_trn.numpy()
-        )
-
-        # add loss to metrics
-        for i in range(len(self.tgt_modalities)):
-            met_trn[i]['Loss'] = np.mean(losses_trn[i])
-        
-        # log metrics to wandb
-        wandb.log({f"Train loss {list(self.tgt_modalities)[i]}": met_trn[i]['Loss']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        wandb.log({f"Train Balanced Accuracy {list(self.tgt_modalities)[i]}": met_trn[i]['Balanced Accuracy']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        
-        wandb.log({f"Train AUC (ROC) {list(self.tgt_modalities)[i]}": met_trn[i]['AUC (ROC)']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        wandb.log({f"Train AUPR {list(self.tgt_modalities)[i]}": met_trn[i]['AUC (PR)']  for i in range(len(self.tgt_modalities))}, step=epoch)
-
-        if self.verbose > 2:
-            print_metrics_multitask(met_trn)
-            
-        return met_trn
-    
-    def validate_one_epoch(self, ldr_vld, epoch):
-        # # progress bar for validation
-        if self.verbose > 1:
-            pbr_batch = ProgressBar(len(ldr_vld.dataset), 'Epoch {:03d} (VLD)'.format(epoch))
-
-        # set model to validation mode
-        torch.set_grad_enabled(False)
-        self.net_.eval()
-
-        scores_vld, y_true_vld, y_mask_vld = [], [], []
-        losses_vld = [[] for _ in self.tgt_modalities]
-        for x_batch, y_batch, mask, y_mask in ldr_vld:
-            # if len(next(iter(x_batch.values()))) < self.batch_size:
-            #     break
-            # mount data to the proper device
-            x_batch = {k: x_batch[k].to(self.device) for k in x_batch} # if 'img' not in k}
-            # x_img_batch = {k: x_img_batch[k].to(self.device) for k in x_img_batch}
-            y_batch = {k: y_batch[k].to(torch.float).to(self.device) for k in y_batch}
-            mask = {k: mask[k].to(self.device) for k in mask}
-            y_mask = {k: y_mask[k].to(self.device) for k in y_mask}
-
-            # forward
-            with torch.autocast(
-                device_type = 'cpu' if self.device == 'cpu' else 'cuda',
-                dtype = torch.bfloat16 if self.device == 'cpu' else torch.float16,
-                enabled = self._amp_enabled
-            ):
-                
-                outputs = self.net_(x_batch, mask, skip_embedding=self.skip_embedding)
-
-                # calculate multitask loss
-                for i, k in enumerate(self.tgt_modalities):
-                    loss_task = self.loss_fn[k](outputs[k], y_batch[k])
-                    msk_loss_task = loss_task * y_mask[k]
-                    losses_vld[i] += msk_loss_task.detach().cpu().numpy().tolist()
-
-            ''' TODO: change array to dictionary later '''
-            outputs = torch.stack(list(outputs.values()), dim=1)
-            y_batch = torch.stack(list(y_batch.values()), dim=1)
-            y_mask = torch.stack(list(y_mask.values()), dim=1)
-
-            # save outputs to evaluate performance later
-            scores_vld.append(outputs.detach().to(torch.float).cpu())
-            y_true_vld.append(y_batch.cpu())
-            y_mask_vld.append(y_mask.cpu())
-
-            # update progress bar
-            if self.verbose > 1:
-                batch_size = len(next(iter(x_batch.values())))
-                pbr_batch.update(batch_size, {})
-                pbr_batch.refresh()
-
-            # clear cuda cache
-            if "cuda" in self.device:
-                torch.cuda.empty_cache()
-
-        # for better tqdm progress bar display
-        if self.verbose > 1:
-            pbr_batch.close()
-
-        # calculate and print validation performance metrics
-        scores_vld = torch.cat(scores_vld)
-        y_true_vld = torch.cat(y_true_vld)
-        y_mask_vld = torch.cat(y_mask_vld)
-        y_pred_vld = (scores_vld > 0).to(torch.int)
-        y_prob_vld = torch.sigmoid(scores_vld)
-        met_vld = get_metrics_multitask(
-            y_true_vld.numpy(),
-            y_pred_vld.numpy(),
-            y_prob_vld.numpy(),
-            y_mask_vld.numpy()
-        )
-
-        # add loss to metrics
-        for i in range(len(self.tgt_modalities)):
-            met_vld[i]['Loss'] = np.mean(losses_vld[i])
-            
-        wandb.log({f"Validation loss {list(self.tgt_modalities)[i]}": met_vld[i]['Loss'] for i in range(len(self.tgt_modalities))}, step=epoch)
-        wandb.log({f"Validation Balanced Accuracy {list(self.tgt_modalities)[i]}": met_vld[i]['Balanced Accuracy']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        
-        wandb.log({f"Validation AUC (ROC) {list(self.tgt_modalities)[i]}": met_vld[i]['AUC (ROC)']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        wandb.log({f"Validation AUPR {list(self.tgt_modalities)[i]}": met_vld[i]['AUC (PR)']  for i in range(len(self.tgt_modalities))}, step=epoch)
-
-        if self.verbose > 2:
-            print_metrics_multitask(met_vld)
-        
-        return met_vld
-        
-
-    def predict_logits(self,
-        x: list[dict[str, Any]],
-        _batch_size: int | None = None,
-        skip_embedding: dict | None = None,
-        img_transform: Any | None = None,
-    ) -> list[dict[str, float]]:
-        '''
-        The input x can be a single sample or a list of samples.
-        '''
-        # input validation
-        check_is_fitted(self)
-        print(self.device)
-        
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # set model to eval mode
-        torch.set_grad_enabled(False)
-        self.net_.eval()
-
-        # intialize dataset and dataloader object
-        dat = TransformerTestingDataset(x, self.src_modalities, img_transform=img_transform)
-        ldr = DataLoader(
-            dataset = dat,
-            batch_size = _batch_size if _batch_size is not None else len(x),
-            shuffle = False,
-            drop_last = False,
-            num_workers = 0,
-            collate_fn = TransformerTestingDataset.collate_fn,
-        )
-        # print("dataloader done")
-
-        # run model and collect results
-        logits: list[dict[str, float]] = []
-        for x_batch, mask in tqdm(ldr):
-            # mount data to the proper device
-            # print(x_batch['his_SEX'])
-            x_batch = {k: x_batch[k].to(self.device) for k in x_batch} 
-            mask = {k: mask[k].to(self.device) for k in mask}
-
-            # forward
-            output: dict[str, Tensor] = self.net_(x_batch, mask, skip_embedding)
-            
-            # convert output from dict-of-list to list of dict, then append
-            tmp = {k: output[k].tolist() for k in self.tgt_modalities}
-            tmp = [{k: tmp[k][i] for k in self.tgt_modalities} for i in range(len(next(iter(tmp.values()))))]
-            logits += tmp
-
-        return logits
-        
-    def predict_proba(self,
-        x: list[dict[str, Any]],
-        skip_embedding: dict | None = None,
-        temperature: float = 1.0,
-        _batch_size: int | None = None,
-        img_transform: Any | None = None,
-    ) -> list[dict[str, float]]:
-        ''' ... '''
-        logits = self.predict_logits(x=x, _batch_size=_batch_size, img_transform=img_transform, skip_embedding=skip_embedding)
-        print("got logits")
-        return logits, [{k: expit(smp[k] / temperature) for k in self.tgt_modalities} for smp in logits]
-
-    def predict(self,
-        x: list[dict[str, Any]],
-        skip_embedding: dict | None = None,
-        fpr: dict[str, Any] | None = None,
-        tpr: dict[str, Any] | None = None,
-        thresholds: dict[str, Any] | None = None,
-        _batch_size: int | None = None,
-        img_transform: Any | None = None,
-    ) -> list[dict[str, int]]:
-        ''' ... '''
-        if fpr is None or tpr is None or thresholds is None:
-            logits, proba = self.predict_proba(x, _batch_size=_batch_size, img_transform=img_transform, skip_embedding=skip_embedding)
-            print("got proba")
-            return logits, proba, [{k: int(smp[k] > 0.5) for k in self.tgt_modalities} for smp in proba]
-        else:
-            logits, proba = self.predict_proba(x, _batch_size=_batch_size, img_transform=img_transform, skip_embedding=skip_embedding)
-            print("got proba")
-            youden_index = {}
-            thr = {}
-            for i, k in enumerate(self.tgt_modalities):
-                youden_index[k] = tpr[i] - fpr[i]
-                thr[k] = thresholds[i][np.argmax(youden_index[k])]
-            #     print(thr[k])
-            # print(thr)
-            return logits, proba, [{k: int(smp[k] > thr[k]) for k in self.tgt_modalities} for smp in proba]
-
-    def save(self, filepath: str, epoch: int) -> None:
-        """Save the model to the given file stream.
-
-        :param filepath: _description_
-        :type filepath: str
-        :param epoch: _description_
-        :type epoch: int
-        """        
-        check_is_fitted(self)
-        if self.data_parallel:
-            state_dict = self.net_.module.state_dict()
-        else:
-            state_dict = self.net_.state_dict()
-
-        # attach model hyper parameters
-        state_dict['src_modalities'] = self.src_modalities
-        state_dict['tgt_modalities'] = self.tgt_modalities
-        state_dict['d_model'] = self.d_model
-        state_dict['nhead'] = self.nhead
-        state_dict['num_encoder_layers'] = self.num_encoder_layers
-        state_dict['num_decoder_layers'] = self.num_decoder_layers
-        state_dict['optimizer'] = self.optimizer
-        state_dict['img_net'] = self.img_net
-        state_dict['imgnet_layers'] = self.imgnet_layers
-        state_dict['img_size'] = self.img_size
-        state_dict['patch_size'] = self.patch_size
-        state_dict['imgnet_ckpt'] = self.imgnet_ckpt
-        state_dict['train_imgnet'] = self.train_imgnet
-        state_dict['epoch'] = epoch
-
-        if self.scaler is not None:
-            state_dict['scaler'] = self.scaler.state_dict()
-        if self.label_distribution:
-            state_dict['label_distribution'] = self.label_distribution
-
-        torch.save(state_dict, filepath)
-
-    def load(self, filepath: str, map_location: str = 'cpu', img_dict=None) -> None:
-        """Load a model from the given file stream.
-
-        :param filepath: _description_
-        :type filepath: str
-        :param map_location: _description_, defaults to 'cpu'
-        :type map_location: str, optional
-        :param img_dict: _description_, defaults to None
-        :type img_dict: _type_, optional
-        """        
-        # load state_dict
-        state_dict = torch.load(filepath, map_location=map_location)
-
-        # load data modalities
-        self.src_modalities: dict[str, dict[str, Any]] = state_dict.pop('src_modalities')
-        self.tgt_modalities: dict[str, dict[str, Any]] = state_dict.pop('tgt_modalities')
-        if 'label_distribution' in state_dict:
-            self.label_distribution: dict[str, dict[int, int]] = state_dict.pop('label_distribution')
-        if 'optimizer' in state_dict:
-            self.optimizer = state_dict.pop('optimizer')
-
-        # initialize model
-        self.d_model = state_dict.pop('d_model')
-        self.nhead = state_dict.pop('nhead')
-        self.num_encoder_layers = state_dict.pop('num_encoder_layers')
-        self.num_decoder_layers = state_dict.pop('num_decoder_layers')
-        if 'epoch' in state_dict.keys():
-            self.start_epoch = state_dict.pop('epoch')
-        if img_dict is None:
-            self.img_net = state_dict.pop('img_net')
-            self.imgnet_layers = state_dict.pop('imgnet_layers')
-            self.img_size = state_dict.pop('img_size')
-            self.patch_size = state_dict.pop('patch_size')
-            self.imgnet_ckpt = state_dict.pop('imgnet_ckpt')
-            self.train_imgnet = state_dict.pop('train_imgnet')
-        else:
-            self.img_net  = img_dict['img_net']
-            self.imgnet_layers  = img_dict['imgnet_layers']
-            self.img_size  = img_dict['img_size']
-            self.patch_size  = img_dict['patch_size']
-            self.imgnet_ckpt  = img_dict['imgnet_ckpt']
-            self.train_imgnet  = img_dict['train_imgnet']
-            state_dict.pop('img_net')
-            state_dict.pop('imgnet_layers')
-            state_dict.pop('img_size')
-            state_dict.pop('patch_size')
-            state_dict.pop('imgnet_ckpt')
-            state_dict.pop('train_imgnet')
-            
-        for k, info in self.src_modalities.items():
-            if info['type'] == 'imaging':
-                if 'emb' not in self.img_net.lower():
-                    info['shape'] = (1,) + (self.img_size,) * 3
-                    info['img_shape'] = (1,) + (self.img_size,) * 3
-                elif 'swinunetr' in self.img_net.lower():
-                    info['shape'] = (1, 768, 4, 4, 4)
-                    info['img_shape'] = (1, 768, 4, 4, 4)
-                # print(info['shape'])
-
-        self.net_ = Transformer(self.src_modalities, self.tgt_modalities, self.d_model, self.nhead, self.num_encoder_layers, self.num_decoder_layers, self.device, self.cuda_devices, self.img_net, self.imgnet_layers, self.img_size, self.patch_size, self.imgnet_ckpt, self.train_imgnet, self.fusion_stage)
-
-       
-        if 'scaler' in state_dict and state_dict['scaler']:
-            self.scaler.load_state_dict(state_dict.pop('scaler'))
-        self.net_.load_state_dict(state_dict)
-        check_is_fitted(self)
-        self.net_.to(self.device)
-
-    def to(self, device: str) -> Self:
-        """Mount the model to the given device. 
-
-        :param device: _description_
-        :type device: str
-        :return: _description_
-        :rtype: Self
-        """        
-        self.device = device
-        if hasattr(self, 'model'): self.net_ = self.net_.to(device)
-        if hasattr(self, 'img_model'): self.img_model = self.img_model.to(device)
-        return self
-    
-    @classmethod
-    def from_ckpt(cls, filepath: str, device='cpu', img_dict=None) -> Self:
-        """Create a new ADRD model and load parameters from the checkpoint. 
-
-        This is an alternative constructor.
-
-        :param filepath: _description_
-        :type filepath: str
-        :param device: _description_, defaults to 'cpu'
-        :type device: str, optional
-        :param img_dict: _description_, defaults to None
-        :type img_dict: _type_, optional
-        :return: _description_
-        :rtype: Self
-        """ 
-        obj = cls(None, None, None,device=device)
-        if device == 'cuda':
-            obj.device = "{}:{}".format(obj.device, str(obj.cuda_devices[0]))
-        print(obj.device)
-        obj.load(filepath, map_location=obj.device, img_dict=img_dict)
-        return obj
-    
-    def _init_net(self):
-        """ ... """
-        # set the device for use
-        if self.device == 'cuda':
-            self.device = "{}:{}".format(self.device, str(self.cuda_devices[0]))
-        print("Device: " + self.device)
-        
-        self.start_epoch = 0
-        if self.load_from_ckpt:
-            try:
-                print("Loading model from checkpoint...")
-                self.load(self.ckpt_path, map_location=self.device)
-            except:
-                print("Cannot load from checkpoint. Initializing new model...")
-                self.load_from_ckpt = False
-
-        if not self.load_from_ckpt:
-            self.net_ = nn.Transformer(
-                src_modalities = self.src_modalities, 
-                tgt_modalities = self.tgt_modalities, 
-                d_model = self.d_model, 
-                nhead = self.nhead, 
-                num_encoder_layers = self.num_encoder_layers, 
-                num_decoder_layers = self.num_decoder_layers, 
-                device = self.device, 
-                cuda_devices = self.cuda_devices, 
-                img_net = self.img_net, 
-                layers = self.imgnet_layers, 
-                img_size = self.img_size, 
-                patch_size = self.patch_size, 
-                imgnet_ckpt = self.imgnet_ckpt, 
-                train_imgnet = self.train_imgnet,
-                fusion_stage = self.fusion_stage,
-            )  
-            
-            # intialize model parameters using xavier_uniform
-            for name, p in self.net_.named_parameters():
-                if p.dim() > 1:
-                    torch.nn.init.xavier_uniform_(p)
-        
-        self.net_.to(self.device)
-
-        # Initialize the number of GPUs
-        if self.data_parallel and torch.cuda.device_count() > 1:
-            print("Available", torch.cuda.device_count(), "GPUs!")
-            self.net_ = torch.nn.DataParallel(self.net_, device_ids=self.cuda_devices)
-
-        # return net
-
-    def _init_dataloader(self, x_trn, x_vld, y_trn, y_vld, img_train_trans=None, img_vld_trans=None):    
-        # initialize dataset and dataloader
-        if self.balanced_sampling:
-            dat_trn = Transformer2ndOrderBalancedTrainingDataset(
-                x_trn, y_trn,
-                self.src_modalities,
-                self.tgt_modalities,
-                dropout_rate = .5,
-                dropout_strategy = 'permutation',
-                img_transform=img_train_trans,
-            )
-        else:
-            dat_trn = TransformerTrainingDataset(
-                x_trn, y_trn,
-                self.src_modalities,
-                self.tgt_modalities,
-                dropout_rate = .5,
-                dropout_strategy = 'permutation',
-                img_transform=img_train_trans,
-            )
-
-        dat_vld = TransformerValidationDataset(
-            x_vld, y_vld,
-            self.src_modalities,
-            self.tgt_modalities,
-            img_transform=img_vld_trans,
-        )
-
-        ldr_trn = DataLoader(
-            dataset = dat_trn,
-            batch_size = self.batch_size,
-            shuffle = True,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = TransformerTrainingDataset.collate_fn,
-            # pin_memory = True
-        )
-
-        ldr_vld = DataLoader(
-            dataset = dat_vld,
-            batch_size = self.batch_size,
-            shuffle = False,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = TransformerValidationDataset.collate_fn,
-            # pin_memory = True
-        )
-
-        return ldr_trn, ldr_vld
-    
-    def _init_optimizer(self):
-        """ ... """
-        params = list(self.net_.parameters())
-        return torch.optim.AdamW(
-            params,
-            lr = self.lr,
-            betas = (0.9, 0.98),
-            weight_decay = self.weight_decay
-        )
-    
-    def _init_scheduler(self, optimizer):
-        """ ... """       
-
-        return torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
-            optimizer=optimizer,
-            T_0=64,
-            T_mult=2,
-            eta_min = 0,
-            verbose=(self.verbose > 2)
-        )
-    
-    def _init_loss_func(self, 
-        num_per_cls: dict[str, tuple[int, int]],
-    ) -> dict[str, Module]:
-        """ ... """
-        return {k: nn.SigmoidFocalLossBeta(
-            beta = self.beta,
-            gamma = self.gamma,
-            num_per_cls = num_per_cls[k],
-            reduction = 'none',
-        ) for k in self.tgt_modalities}
-    
-    def _proc_fit(self):
-        """ ... """

adrd/model/calibration.py

@@ -1,450 +0,0 @@
-import numpy as np
-from sklearn.base import BaseEstimator
-from sklearn.utils.validation import check_is_fitted
-from sklearn.linear_model import LogisticRegression
-from sklearn.isotonic import IsotonicRegression
-from functools import lru_cache
-from functools import cached_property
-from typing import Self, Any
-from pickle import dump
-from pickle import load
-from abc import ABC, abstractmethod
-
-from . import ADRDModel
-from ..utils import Formatter
-from ..utils import MissingMasker
-
-
-def calibration_curve(
-    y_true: list[int],
-    y_pred: list[float],
-    n_bins: int = 10,
-    ratio: float = 1.0,
-) -> tuple[list[float], list[float]]:
-    """
-    Compute true and predicted probabilities for a calibration curve. The method
-    assumes the inputs come from a binary classifier, and discretize the [0, 1] 
-    interval into bins.
-
-    Note that this function is an alternative to
-    sklearn.calibration.calibration_curve() which can only estimate the absolute
-    proportion of positive cases in each bin.
-
-    Parameters
-    ----------
-    y_true : list[int]
-        True targets.
-    y_pred : list[float]
-        Probabilities of the positive class.
-    n_bins : int, default=10
-        Number of bins to discretize the [0, 1] interval. A bigger number
-        requires more data. Bins with no samples (i.e. without corresponding
-        values in y_prob) will not be returned, thus the returned arrays may
-        have less than n_bins values.
-    ratio : float, default=1.0
-        Used to adjust the class balance.
-
-    Returns
-    -------
-    prob_true : list[float]
-        The proportion of positive samples in each bin.
-    prob_pred : list[float]
-        The mean predicted probability in each bin.
-    """
-    # generate "n_bin" intervals
-    tmp = np.around(np.linspace(0, 1, n_bins + 1), decimals=6)
-    intvs = [(tmp[i - 1], tmp[i]) for i in range(1, len(tmp))]
-    
-    # pair up (pred, true) and group them by intervals
-    tmp = list(zip(y_pred, y_true))
-    intv_pairs = {(l, r): [p for p in tmp if l <= p[0] < r] for l, r in intvs}
-
-    # calculate balanced proportion of POSITIVE cases for each intervel
-    # along with the balanced averaged predictions
-    intv_prob_true: dict[tuple, float] = dict()
-    intv_prob_pred: dict[tuple, float] = dict()
-    for intv, pairs in intv_pairs.items():
-        # number of cases that fall into the interval
-        n_pairs = len(pairs)
-
-        # it's likely that no predictions fall into the interval
-        if n_pairs == 0: continue
-
-        # count number of positives and negatives in the interval
-        n_pos = sum([p[1] for p in pairs])
-        n_neg = n_pairs - n_pos
-
-        # calculate adjusted proportion of positives
-        intv_prob_true[intv] = n_pos / (n_pos + n_neg * ratio)
-
-        # calculate adjusted avg. predictions
-        sum_pred_pos = sum([p[0] for p in pairs if p[1] == 1])
-        sum_pred_neg = sum([p[0] for p in pairs if p[1] == 0])
-        intv_prob_pred[intv] = (sum_pred_pos + sum_pred_neg * ratio)
-        intv_prob_pred[intv] /= (n_pos + n_neg * ratio)
-
-    prob_true = list(intv_prob_true.values())
-    prob_pred = list(intv_prob_pred.values())
-    return prob_true, prob_pred
-
-
-class CalibrationCore(BaseEstimator):
-    """
-    A wrapper class of multiple regressors to predict the proportions of
-    positive samples from the predicted probabilities. The method for
-    calibration can be 'sigmoid' which corresponds to Platt's method (i.e. a 
-    logistic regression model) or 'isotonic' which is a non-parametric approach.
-    It is not advised to use isotonic calibration with too few calibration
-    samples (<<1000) since it tends to overfit.
-
-    TODO
-    ----
-    - 'sigmoid' method is not trivial to implement.
-    """
-    def __init__(self, 
-        method: str = 'isotonic',
-    ) -> None:
-        """
-        Initialization function of CalibrationCore class.
-
-        Parameters
-        ----------
-        method : {'sigmoid', 'isotonic'}, default='isotonic'
-            The method to use for calibration. can be 'sigmoid' which
-            corresponds to Platt's method (i.e. a logistic regression model) or
-            'isotonic' which is a non-parametric approach. It is not advised to
-            use isotonic calibration with too few calibration samples (<<1000)
-            since it tends to overfit.
-
-        Raises
-        ------
-        ValueError
-            Sigmoid approach has not been implemented.
-        """        
-        assert method in ('sigmoid', 'isotonic')
-        if method == 'sigmoid':
-            raise ValueError('Sigmoid approach has not been implemented.')
-        self.method = method
-
-    def fit(self, 
-        prob_pred: list[float], 
-        prob_true: list[float],
-    ) -> Self:
-        """
-        Fit the underlying regressor using prob_pred, prob_true as training
-        data.
-
-        Parameters
-        ----------
-        prob_pred : list[float]
-            Probabilities predicted directly by a model.
-        prob_true : list[float]
-            Target probabilities to calibrate to.
-
-        Returns
-        -------
-        Self
-            CalibrationCore object.
-        """              
-        # using Platt's method for calibration
-        if self.method == 'sigmoid':
-            self.model_ = LogisticRegression()
-            self.model_.fit(prob_pred, prob_true)
-
-        # using isotonic calibration
-        elif self.method == 'isotonic':
-            self.model_ = IsotonicRegression(y_min=0, y_max=1, out_of_bounds='clip')
-            self.model_.fit(prob_pred, prob_true)
-
-        return self
-
-    def predict(self,
-        prob_pred: list[float],
-    ) -> list[float]:
-        """
-        Calibrate the input probabilities using the fitted regressor.
-
-        Parameters
-        ----------
-        prob_pred : list[float]
-            Probabilities predicted directly by a model.
-
-        Returns
-        -------
-        prob_cali : list[float]
-            Calibrated probabilities.
-        """        
-        # as usual, the core needs to be fitted
-        check_is_fitted(self)
-
-        # note that logistic regression is classification model, we need to call
-        # 'predict_proba' instead of 'predict' to get the calibrated results
-        if self.method == 'sigmoid':
-            prob_cali = self.model_.predict_proba(prob_pred)
-        elif self.method == 'isotonic':
-            prob_cali = self.model_.predict(prob_pred)
-
-        return prob_cali
-    
-
-class CalibratedClassifier(ABC):
-    """
-    Abstract class of calibrated classifier.
-    """
-    def __init__(self, 
-        model: ADRDModel,
-        background_src: list[dict[str, Any]],
-        background_tgt: list[dict[str, Any]],
-        background_is_embedding: dict[str, bool] | None = None,
-        method: str = 'isotonic',
-    ) -> None:
-        """
-        Constructor of Calibrator class.
-
-        Parameters
-        ----------
-        model : ADRDModel
-            Fitted model to calibrate.
-        background_src : list[dict[str, Any]]
-            Features of the background dataset.
-        background_tgt : list[dict[str, Any]]
-            Labels of the background dataset.
-        method : {'sigmoid', 'isotonic'}, default='isotonic'
-            Method used by the underlying regressor. 
-        """
-        self.method = method
-        self.model = model
-        self.src_modalities = model.src_modalities
-        self.tgt_modalities = model.tgt_modalities
-        self.background_is_embedding = background_is_embedding
-
-        # format background data
-        fmt_src = Formatter(self.src_modalities)
-        fmt_tgt = Formatter(self.tgt_modalities)
-        self.background_src = [fmt_src(smp) for smp in background_src]
-        self.background_tgt = [fmt_tgt(smp) for smp in background_tgt]
-    
-    @abstractmethod
-    def predict_proba(self, 
-        src: list[dict[str, Any]],
-        is_embedding: dict[str, bool] | None = None,
-    ) -> list[dict[str, float]]:
-        """
-        This method returns calibrated probabilities of classification.
-
-        Parameters
-        ----------
-        src : list[dict[str, Any]]
-            Features of the input samples.
-
-        Returns
-        -------
-        list[dict[str, float]]
-            Calibrated probabilities.
-        """ 
-        pass
-
-    def predict(self,
-        src: list[dict[str, Any]],
-        is_embedding: dict[str, bool] | None = None,
-    ) -> list[dict[str, int]]:
-        """
-        Make predictions based on the results of predict_proba().
-
-        Parameters
-        ----------
-        x : list[dict[str, Any]]
-            Input features.
-
-        Returns
-        -------
-        list[dict[str, int]]
-            Calibrated predictions.
-        """
-        proba = self.predict_proba(src, is_embedding)
-        return [{k: int(smp[k] > 0.5) for k in self.tgt_modalities} for smp in proba]
-
-    def save(self,
-        filepath_state_dict: str,
-    ) -> None:
-        """
-        Save the state dict and the underlying model to the given paths.
-
-        Parameters
-        ----------
-        filepath_state_dict : str
-            File path to save the state_dict which includes the background
-            dataset and the regressor information.
-        filepath_wrapped_model : str | None, default=None
-            File path to save the wrapped model. If None, the model won't be
-            saved. 
-        """
-        # save state dict
-        state_dict = {
-            'background_src': self.background_src,
-            'background_tgt': self.background_tgt,
-            'background_is_embedding': self.background_is_embedding,
-            'method': self.method,
-        }
-        with open(filepath_state_dict, 'wb') as f:
-            dump(state_dict, f)
-
-    @classmethod
-    def from_ckpt(cls,
-        filepath_state_dict: str,
-        filepath_wrapped_model: str,
-    ) -> Self:
-        """
-        Alternative constructor which loads from checkpoint.
-
-        Parameters
-        ----------
-        filepath_state_dict : str
-            File path to load the state_dict which includes the background
-            dataset and the regressor information.
-        filepath_wrapped_model : str
-            File path of the wrapped model.
-
-        Returns
-        -------
-        Self
-            CalibratedClassifier class object.
-        """
-        with open(filepath_state_dict, 'rb') as f:
-            kwargs = load(f)
-        kwargs['model'] = ADRDModel.from_ckpt(filepath_wrapped_model)
-        return cls(**kwargs)
-
-
-class DynamicCalibratedClassifier(CalibratedClassifier):
-    """
-    The dynamic approach generates background predictions based on the
-    missingness pattern of each input. With an astronomical number of
-    missingness patterns, calibrating each sample requires a comprehensive
-    process that involves running the ADRDModel on the majority of the
-    background data and training a corresponding regressor. This results in a
-    computationally intensive calculation.
-    """
-    def predict_proba(self,
-        src: list[dict[str, Any]],
-        is_embedding: dict[str, bool] | None = None,
-    ) -> list[dict[str, float]]:
-        
-        # initialize mask generator and format inputs
-        msk_gen = MissingMasker(self.src_modalities)
-        fmt_src = Formatter(self.src_modalities)
-        src = [fmt_src(smp) for smp in src]
-
-        # calculate calibrated probabilities
-        calibrated_prob: list[dict[str, float]] = []
-        for smp in src:
-            # model output and missingness pattern
-            prob = self.model.predict_proba([smp], is_embedding)[0]
-            mask = tuple(msk_gen(smp).values())
-
-            # get/fit core and calculate calibrated probabilities
-            core = self._fit_core(mask)
-            calibrated_prob.append({k: core[k].predict([prob[k]])[0] for k in self.tgt_modalities})
-
-        return calibrated_prob
-    
-    # @lru_cache(maxsize = None)
-    def _fit_core(self,
-        missingness_pattern: tuple[bool],
-    ) -> dict[str, CalibrationCore]:
-        ''' ... ''' 
-        # remove features from all background samples accordingly
-        background_src, background_tgt = [], []
-        for src, tgt in zip(self.background_src, self.background_tgt):
-            src = {k: v for j, (k, v) in enumerate(src.items()) if missingness_pattern[j] == False}
-
-            # make sure there is at least one feature available
-            if len([v is not None for v in src.values()]) == 0: continue
-            background_src.append(src)
-            background_tgt.append(tgt)
-
-        # run model on background samples and collection predictions
-        background_prob = self.model.predict_proba(background_src, self.background_is_embedding, _batch_size=1024)
-
-        # list[dict] -> dict[list]
-        N = len(background_src)
-        background_prob = {k: [background_prob[i][k] for i in range(N)] for k in self.tgt_modalities}
-        background_true = {k: [background_tgt[i][k] for i in range(N)] for k in self.tgt_modalities}
-
-        # now, fit cores
-        core: dict[str, CalibrationCore] = dict()
-        for k in self.tgt_modalities:
-            prob_true, prob_pred = calibration_curve(
-                background_true[k], background_prob[k],
-                ratio = self.background_ratio[k],
-            )
-            core[k] = CalibrationCore(self.method).fit(prob_pred, prob_true)
-        
-        return core
-    
-    @cached_property
-    def background_ratio(self) -> dict[str, float]:
-        ''' The ratio of positives over negatives in the background dataset. '''
-        return {k: self.background_n_pos[k] / self.background_n_neg[k] for k in self.tgt_modalities}
-
-    @cached_property
-    def background_n_pos(self) -> dict[str, int]:
-        ''' Number of positives w.r.t each target in the background dataset. '''
-        return {k: sum([d[k] for d in self.background_tgt]) for k in self.tgt_modalities}
-
-    @cached_property
-    def background_n_neg(self) -> dict[str, int]:
-        ''' Number of negatives w.r.t each target in the background dataset. '''
-        return {k: len(self.background_tgt) - self.background_n_pos[k] for k in self.tgt_modalities}
-
-
-class StaticCalibratedClassifier(CalibratedClassifier):
-    """
-    The static approach generates background predictions without considering the
-    missingness patterns.
-    """
-    def predict_proba(self,
-        src: list[dict[str, Any]],
-        is_embedding: dict[str, bool] | None = None,
-    ) -> list[dict[str, float]]:
-
-        # number of input samples
-        N = len(src)
-
-        # format inputs, and run ADRDModel, and convert to dict[list]
-        fmt_src = Formatter(self.src_modalities)
-        src = [fmt_src(smp) for smp in src]
-        prob = self.model.predict_proba(src, is_embedding)
-        prob = {k: [prob[i][k] for i in range(N)] for k in self.tgt_modalities}
-
-        # calibrate probabilities
-        core = self._fit_core()
-        calibrated_prob = {k: core[k].predict(prob[k]) for k in self.tgt_modalities}
-
-        # convert back to list[dict]
-        calibrated_prob: list[dict[str, float]] = [
-            {k: calibrated_prob[k][i] for k in self.tgt_modalities} for i in range(N)
-        ]
-        return calibrated_prob
-    
-    @lru_cache(maxsize = None)
-    def _fit_core(self) -> dict[str, CalibrationCore]:
-        ''' ... '''
-        # run model on background samples and collection predictions
-        background_prob = self.model.predict_proba(self.background_src, self.background_is_embedding, _batch_size=1024)
-
-        # list[dict] -> dict[list]
-        N = len(self.background_src)
-        background_prob = {k: [background_prob[i][k] for i in range(N)] for k in self.tgt_modalities}
-        background_true = {k: [self.background_tgt[i][k] for i in range(N)] for k in self.tgt_modalities}
-
-        # now, fit cores
-        core: dict[str, CalibrationCore] = dict()
-        for k in self.tgt_modalities:
-            prob_true, prob_pred = calibration_curve(
-                background_true[k], background_prob[k],
-                ratio = 1.0,
-            )
-            core[k] = CalibrationCore(self.method).fit(prob_pred, prob_true)
-        
-        return core

adrd/model/cnn_resnet3d_with_linear_classifier.py

@@ -1,533 +0,0 @@
-__all__ = ['CNNResNet3DWithLinearClassifier']
-
-import torch
-from torch.utils.data import DataLoader
-import numpy as np
-import tqdm
-from sklearn.base import BaseEstimator
-from sklearn.utils.validation import check_is_fitted
-from sklearn.model_selection import train_test_split
-from scipy.special import expit
-from copy import deepcopy
-from contextlib import suppress
-from typing import Any, Self, Type
-from functools import wraps
-Tensor = Type[torch.Tensor]
-Module = Type[torch.nn.Module]
-
-from ..utils.misc import ProgressBar
-from ..utils.misc import get_metrics_multitask, print_metrics_multitask
-
-from .. import nn
-from ..utils import TransformerTrainingDataset
-from ..utils import Transformer2ndOrderBalancedTrainingDataset
-from ..utils import TransformerValidationDataset
-from ..utils import TransformerTestingDataset
-from ..utils.misc import ProgressBar
-from ..utils.misc import get_metrics_multitask, print_metrics_multitask
-from ..utils.misc import convert_args_kwargs_to_kwargs
-
-
-def _manage_ctx_fit(func):
-    ''' ... '''
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        # format arguments
-        kwargs = convert_args_kwargs_to_kwargs(func, args, kwargs)
-
-        if kwargs['self']._device_ids is None:
-            return func(**kwargs)
-        else:
-            # change primary device
-            default_device = kwargs['self'].device
-            kwargs['self'].device = kwargs['self']._device_ids[0]
-            rtn = func(**kwargs)
-
-            # the actual module is wrapped
-            kwargs['self'].net_ = kwargs['self'].net_.module
-            kwargs['self'].to(default_device)
-            return rtn
-    return wrapper
-
-
-class CNNResNet3DWithLinearClassifier(BaseEstimator):
-
-    def __init__(self,
-        src_modalities: dict[str, dict[str, Any]],
-        tgt_modalities: dict[str, dict[str, Any]],
-        num_epochs: int = 32,
-        batch_size: int = 8,
-        batch_size_multiplier: int = 1,
-        lr: float = 1e-2,
-        weight_decay: float = 0.0,
-        beta: float = 0.9999,
-        gamma: float = 2.0,
-        scale: float = 1.0,
-        criterion: str | None = None,
-        device: str = 'cpu',
-        verbose: int = 0,
-        _device_ids: list | None = None,
-        _dataloader_num_workers: int = 0,
-        _amp_enabled: bool = False,
-        _tmp_ckpt_filepath: str | None = None,
-    ) -> None:  
-        ''' ... '''
-        # for multiprocessing
-        self._rank = 0
-        self._lock = None
-
-        # positional parameters
-        self.src_modalities = src_modalities
-        self.tgt_modalities = tgt_modalities
-
-        # training parameters
-        self.num_epochs = num_epochs
-        self.batch_size = batch_size
-        self.batch_size_multiplier = batch_size_multiplier
-        self.lr = lr
-        self.weight_decay = weight_decay
-        self.beta = beta
-        self.gamma = gamma
-        self.scale = scale
-        self.criterion = criterion
-        self.device = device
-        self.verbose = verbose
-        self._device_ids = _device_ids
-        self._dataloader_num_workers = _dataloader_num_workers
-        self._amp_enabled = _amp_enabled
-        self._tmp_ckpt_filepath = _tmp_ckpt_filepath
-
-    
-    @_manage_ctx_fit
-    def fit(self, x, y) -> Self:
-        ''' ... '''
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # initialize neural network
-        self.net_ = self._init_net()
-
-        # initialize dataloaders
-        ldr_trn, ldr_vld = self._init_dataloader(x, y)
-
-        # initialize optimizer and scheduler
-        optimizer = self._init_optimizer()
-        scheduler = self._init_scheduler(optimizer)
-
-        # gradient scaler for AMP
-        if self._amp_enabled: scaler = torch.cuda.amp.GradScaler()
-        
-        # initialize loss function (binary cross entropy)
-        loss_func = self._init_loss_func({
-            k: (
-                sum([_[k] == 0 for _ in ldr_trn.dataset.tgt]),
-                sum([_[k] == 1 for _ in ldr_trn.dataset.tgt]),
-            ) for k in self.tgt_modalities
-        })
-
-        # to record the best validation performance criterion
-        if self.criterion is not None: best_crit = None
-
-        # progress bar for epoch loops
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch = tqdm.tqdm(
-                    desc = 'Rank {:02d}'.format(self._rank),
-                    total = self.num_epochs,
-                    position = self._rank,
-                    ascii = True,
-                    leave = False,
-                    bar_format='{l_bar}{r_bar}'
-                )
-
-        # training loop
-        for epoch in range(self.num_epochs):
-            # progress bar for batch loops
-            if self.verbose > 1: 
-                pbr_batch = ProgressBar(len(ldr_trn.dataset), 'Epoch {:03d} (TRN)'.format(epoch))
-
-            # set model to train mode
-            torch.set_grad_enabled(True)
-            self.net_.train()
-
-            scores_trn: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            y_true_trn: dict[str, list[int]]   = {k: [] for k in self.tgt_modalities}
-            losses_trn: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            for n_iter, (x_batch, y_batch, _, mask_y) in enumerate(ldr_trn):
-                # mount data to the proper device
-                x_batch = {k: x_batch[k].to(self.device) for k in self.src_modalities}
-                y_batch = {k: y_batch[k].to(torch.float).to(self.device) for k in self.tgt_modalities}
-                # mask_x = {k: mask_x[k].to(self.device) for k in self.src_modalities}
-                mask_y = {k: mask_y[k].to(self.device) for k in self.tgt_modalities}
-
-                # forward
-                with torch.autocast(
-                    device_type = 'cpu' if self.device == 'cpu' else 'cuda',
-                    dtype = torch.bfloat16 if self.device == 'cpu' else torch.float16,
-                    enabled = self._amp_enabled,
-                ):
-                    outputs = self.net_(x_batch)
-
-                    # calculate multitask loss
-                    loss = 0
-                    for i, tgt_k in enumerate(self.tgt_modalities):
-                        loss_k = loss_func[tgt_k](outputs[tgt_k], y_batch[tgt_k])
-                        loss_k = torch.masked_select(loss_k, torch.logical_not(mask_y[tgt_k].squeeze()))
-                        loss += loss_k.mean()
-                        losses_trn[tgt_k] += loss_k.detach().cpu().numpy().tolist()
-                
-                # backward
-                if self._amp_enabled:
-                    scaler.scale(loss).backward()
-                else:
-                    loss.backward()
-                
-                # update parameters
-                if n_iter != 0 and n_iter % self.batch_size_multiplier == 0:
-                    if self._amp_enabled:
-                        scaler.step(optimizer)
-                        scaler.update()
-                        optimizer.zero_grad()
-                    else: 
-                        optimizer.step()
-                        optimizer.zero_grad()
-
-                # save outputs to evaluate performance later
-                for tgt_k in self.tgt_modalities:
-                    tmp = torch.masked_select(outputs[tgt_k], torch.logical_not(mask_y[tgt_k].squeeze()))
-                    scores_trn[tgt_k] += tmp.detach().cpu().numpy().tolist()
-                    tmp = torch.masked_select(y_batch[tgt_k], torch.logical_not(mask_y[tgt_k].squeeze()))
-                    y_true_trn[tgt_k] += tmp.cpu().numpy().tolist()
-
-                # update progress bar
-                if self.verbose > 1:
-                    batch_size = len(next(iter(x_batch.values())))
-                    pbr_batch.update(batch_size, {})
-                    pbr_batch.refresh()
-
-            # for better tqdm progress bar display
-            if self.verbose > 1:
-                pbr_batch.close()
-
-            # set scheduler
-            scheduler.step()
-
-            # calculate and print training performance metrics
-            y_pred_trn: dict[str, list[int]]   = {k: [] for k in self.tgt_modalities}
-            y_prob_trn: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            for tgt_k in self.tgt_modalities:
-                for i in range(len(scores_trn[tgt_k])):
-                    y_pred_trn[tgt_k].append(1 if scores_trn[tgt_k][i] > 0 else 0)
-                    y_prob_trn[tgt_k].append(expit(scores_trn[tgt_k][i]))
-            met_trn = get_metrics_multitask(y_true_trn, y_pred_trn, y_prob_trn)
-
-            # add loss to metrics
-            for tgt_k in self.tgt_modalities:
-                met_trn[tgt_k]['Loss'] = np.mean(losses_trn[tgt_k])
-
-            if self.verbose > 2:
-                print_metrics_multitask(met_trn)
-
-            # progress bar for validation
-            if self.verbose > 1:
-                pbr_batch = ProgressBar(len(ldr_vld.dataset), 'Epoch {:03d} (VLD)'.format(epoch))
-
-            # set model to validation mode
-            torch.set_grad_enabled(False)
-            self.net_.eval()
-
-            scores_vld: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            y_true_vld: dict[str, list[int]]   = {k: [] for k in self.tgt_modalities}
-            losses_vld: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            for x_batch, y_batch, _, mask_y in ldr_vld:
-                # mount data to the proper device
-                x_batch = {k: x_batch[k].to(self.device) for k in self.src_modalities}
-                y_batch = {k: y_batch[k].to(torch.float).to(self.device) for k in self.tgt_modalities}
-                # mask_x = {k: mask_x[k].to(self.device) for k in self.src_modalities}
-                mask_y = {k: mask_y[k].to(self.device) for k in self.tgt_modalities}
-
-                # forward
-                with torch.autocast(
-                    device_type = 'cpu' if self.device == 'cpu' else 'cuda',
-                    dtype = torch.bfloat16 if self.device == 'cpu' else torch.float16,
-                    enabled = self._amp_enabled
-                ):
-                    outputs = self.net_(x_batch)
-
-                    # calculate multitask loss
-                    for i, tgt_k in enumerate(self.tgt_modalities):
-                        loss_k = loss_func[tgt_k](outputs[tgt_k], y_batch[tgt_k])
-                        loss_k = torch.masked_select(loss_k, torch.logical_not(mask_y[tgt_k].squeeze()))
-                        losses_vld[tgt_k] += loss_k.detach().cpu().numpy().tolist()
-
-                # save outputs to evaluate performance later
-                for tgt_k in self.tgt_modalities:
-                    tmp = torch.masked_select(outputs[tgt_k], torch.logical_not(mask_y[tgt_k].squeeze()))
-                    scores_vld[tgt_k] += tmp.detach().cpu().numpy().tolist()
-                    tmp = torch.masked_select(y_batch[tgt_k], torch.logical_not(mask_y[tgt_k].squeeze()))
-                    y_true_vld[tgt_k] += tmp.cpu().numpy().tolist()
-
-                # update progress bar
-                if self.verbose > 1:
-                    batch_size = len(next(iter(x_batch.values())))
-                    pbr_batch.update(batch_size, {})
-                    pbr_batch.refresh()
-
-            # for better tqdm progress bar display
-            if self.verbose > 1:
-                pbr_batch.close()
-
-            # calculate and print validation performance metrics
-            y_pred_vld: dict[str, list[int]]   = {k: [] for k in self.tgt_modalities}
-            y_prob_vld: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            for tgt_k in self.tgt_modalities:
-                for i in range(len(scores_vld[tgt_k])):
-                    y_pred_vld[tgt_k].append(1 if scores_vld[tgt_k][i] > 0 else 0)
-                    y_prob_vld[tgt_k].append(expit(scores_vld[tgt_k][i]))
-            met_vld = get_metrics_multitask(y_true_vld, y_pred_vld, y_prob_vld)
-
-            # add loss to metrics
-            for tgt_k in self.tgt_modalities:
-                met_vld[tgt_k]['Loss'] = np.mean(losses_vld[tgt_k])
-
-            if self.verbose > 2:
-                print_metrics_multitask(met_vld)
-
-            # save the model if it has the best validation performance criterion by far
-            if self.criterion is None: continue
-            
-            # is current criterion better than previous best?
-            curr_crit = np.mean([met_vld[k][self.criterion] for k in self.tgt_modalities])
-            if best_crit is None or np.isnan(best_crit):
-                is_better = True
-            elif self.criterion == 'Loss' and best_crit >= curr_crit:
-                is_better = True
-            elif self.criterion != 'Loss' and best_crit <= curr_crit:
-                is_better = True
-            else:
-                is_better = False
-
-            # update best criterion
-            if is_better:
-                best_crit = curr_crit
-                best_state_dict = deepcopy(self.net_.state_dict())
-
-                if self._tmp_ckpt_filepath is not None:
-                    self.save(self._tmp_ckpt_filepath)
-
-            if self.verbose > 2:
-                print('Best {}: {}'.format(self.criterion, best_crit))
-
-            if self.verbose == 1:
-                with self._lock if self._lock is not None else suppress():
-                    pbr_epoch.update(1)
-                    pbr_epoch.refresh()
-
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch.close()
-
-        # restore the model of the best validation performance across all epoches
-        if ldr_vld is not None and self.criterion is not None:
-            self.net_.load_state_dict(best_state_dict)
-
-        return self
-
-    def predict_logits(self,
-        x: list[dict[str, Any]],
-        _batch_size: int | None = None,
-    ) -> list[dict[str, float]]:
-        """
-        The input x can be a single sample or a list of samples.
-        """
-        # input validation
-        check_is_fitted(self)
-        
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # set model to eval mode
-        torch.set_grad_enabled(False)
-        self.net_.eval()
-
-        # intialize dataset and dataloader object
-        dat = TransformerTestingDataset(x, self.src_modalities)
-        ldr = DataLoader(
-            dataset = dat,
-            batch_size = _batch_size if _batch_size is not None else len(x),
-            shuffle = False,
-            drop_last = False,
-            num_workers = 0,
-            collate_fn = TransformerTestingDataset.collate_fn,
-        )
-
-        # run model and collect results
-        logits: list[dict[str, float]] = []
-        for x_batch, _ in ldr:
-            # mount data to the proper device
-            x_batch = {k: x_batch[k].to(self.device) for k in self.src_modalities}
-
-            # forward
-            output: dict[str, Tensor] = self.net_(x_batch)
-            
-            # convert output from dict-of-list to list of dict, then append
-            tmp = {k: output[k].tolist() for k in self.tgt_modalities}
-            tmp = [{k: tmp[k][i] for k in self.tgt_modalities} for i in range(len(next(iter(tmp.values()))))]
-            logits += tmp
-
-        return logits
-        
-    def predict_proba(self,
-        x: list[dict[str, Any]],
-        temperature: float = 1.0,
-        _batch_size: int | None = None,
-    ) -> list[dict[str, float]]:
-        ''' ... '''
-        logits = self.predict_logits(x, _batch_size)
-        return [{k: expit(smp[k] / temperature) for k in self.tgt_modalities} for smp in logits]
-
-    def predict(self,
-        x: list[dict[str, Any]],
-        _batch_size: int | None = None,
-    ) -> list[dict[str, int]]:
-        ''' ... '''
-        logits = self.predict_logits(x, _batch_size)
-        return [{k: int(smp[k] > 0.0) for k in self.tgt_modalities} for smp in logits]
-
-    def save(self, filepath: str) -> None:
-        ''' ... '''
-        check_is_fitted(self)
-        state_dict = self.net_.state_dict()
-
-        # attach model hyper parameters
-        state_dict['src_modalities'] = self.src_modalities
-        state_dict['tgt_modalities'] = self.tgt_modalities
-        print('Saving model checkpoint to {} ... '.format(filepath), end='')
-        torch.save(state_dict, filepath)
-        print('Done.')
-
-    def load(self, filepath: str) -> None:
-        ''' ... '''
-        # load state_dict
-        state_dict = torch.load(filepath, map_location='cpu')
-
-        # load essential parameters
-        self.src_modalities: dict[str, dict[str, Any]] = state_dict.pop('src_modalities')
-        self.tgt_modalities: dict[str, dict[str, Any]] = state_dict.pop('tgt_modalities')
-
-        # initialize model
-        self.net_ = nn.CNNResNet3DWithLinearClassifier(
-            self.src_modalities,
-            self.tgt_modalities,
-        )
-
-        # load model parameters
-        self.net_.load_state_dict(state_dict)
-        self.to(self.device)
-
-    def to(self, device: str) -> Self:
-        ''' Mount model to the given device. '''
-        self.device = device
-        if hasattr(self, 'net_'): self.net_ = self.net_.to(device)
-        return self
-    
-    @classmethod
-    def from_ckpt(cls, filepath: str) -> Self:
-        ''' ... '''
-        obj = cls(None, None)
-        obj.load(filepath)
-        return obj
-
-    def _init_net(self):
-        """ ... """
-        net = nn.CNNResNet3DWithLinearClassifier(
-            self.src_modalities,
-            self.tgt_modalities,
-        ).to(self.device)
-
-        # train on multiple GPUs using torch.nn.DataParallel
-        if self._device_ids is not None:
-            net = torch.nn.DataParallel(net, device_ids=self._device_ids)
-
-        # intialize model parameters using xavier_uniform
-        for p in net.parameters():
-            if p.dim() > 1:
-                torch.nn.init.xavier_uniform_(p)
-        
-        return net
-
-    def _init_dataloader(self, x, y):
-        """ ... """
-        # split dataset
-        x_trn, x_vld, y_trn, y_vld = train_test_split(
-            x, y, test_size = 0.2, random_state = 0,
-        )
-
-        # initialize dataset and dataloader
-        # dat_trn = TransformerTrainingDataset(
-        dat_trn = Transformer2ndOrderBalancedTrainingDataset(
-            x_trn, y_trn,
-            self.src_modalities,
-            self.tgt_modalities,
-            dropout_rate = .5,
-            # dropout_strategy = 'compensated',
-            dropout_strategy = 'permutation',
-        )
-
-        dat_vld = TransformerValidationDataset(
-            x_vld, y_vld,
-            self.src_modalities,
-            self.tgt_modalities,
-        )
-
-        ldr_trn = DataLoader(
-            dataset = dat_trn,
-            batch_size = self.batch_size,
-            shuffle = True,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = TransformerTrainingDataset.collate_fn,
-            # pin_memory = True
-        )
-
-        ldr_vld = DataLoader(
-            dataset = dat_vld,
-            batch_size = self.batch_size,
-            shuffle = False,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = TransformerValidationDataset.collate_fn,
-            # pin_memory = True
-        )
-
-        return ldr_trn, ldr_vld
-    
-    def _init_optimizer(self):
-        """ ... """
-        return torch.optim.AdamW(
-            self.net_.parameters(),
-            lr = self.lr,
-            betas = (0.9, 0.98),
-            weight_decay = self.weight_decay
-        )
-    
-    def _init_scheduler(self, optimizer):
-        """ ... """
-        return torch.optim.lr_scheduler.OneCycleLR(
-            optimizer = optimizer, 
-            max_lr = self.lr,
-            total_steps = self.num_epochs,
-            verbose = (self.verbose > 2)
-        )
-    
-    def _init_loss_func(self, 
-        num_per_cls: dict[str, tuple[int, int]],
-    ) -> dict[str, Module]:
-        """ ... """
-        return {k: nn.SigmoidFocalLoss(
-            beta = self.beta,
-            gamma = self.gamma,
-            scale = self.scale,
-            num_per_cls = num_per_cls[k],
-            reduction = 'none',
-        ) for k in self.tgt_modalities}

adrd/model/imaging_model.py

@@ -1,843 +0,0 @@
-__all__ = ['Transformer']
-
-import wandb
-import torch
-import numpy as np
-import functools
-import inspect
-import monai
-import random
-
-from tqdm import tqdm
-from functools import wraps
-from sklearn.base import BaseEstimator
-from sklearn.utils.validation import check_is_fitted
-from sklearn.model_selection import train_test_split
-from scipy.special import expit
-from copy import deepcopy
-from contextlib import suppress
-from typing import Any, Self, Type
-Tensor = Type[torch.Tensor]
-Module = Type[torch.nn.Module]
-from torch.utils.data import DataLoader
-from monai.utils.type_conversion import convert_to_tensor
-from monai.transforms import (
-    LoadImaged,
-    Compose,
-    CropForegroundd,
-    CopyItemsd,
-    SpatialPadd,
-    EnsureChannelFirstd,
-    Spacingd,
-    OneOf,
-    ScaleIntensityRanged,
-    HistogramNormalized,
-    RandSpatialCropSamplesd,
-    RandSpatialCropd,
-    CenterSpatialCropd,
-    RandCoarseDropoutd,
-    RandCoarseShuffled,
-    Resized,
-)
-
-# for DistributedDataParallel
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from torch.nn.parallel import DistributedDataParallel as DDP
-
-from .. import nn
-from ..utils.misc import ProgressBar
-from ..utils.misc import get_metrics_multitask, print_metrics_multitask
-from ..utils.misc import convert_args_kwargs_to_kwargs
-
-import warnings
-warnings.filterwarnings("ignore")
-
-
-def _manage_ctx_fit(func):
-    ''' ... '''
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        # format arguments
-        kwargs = convert_args_kwargs_to_kwargs(func, args, kwargs)
-
-        if kwargs['self']._device_ids is None:
-            return func(**kwargs)
-        else:
-            # change primary device
-            default_device = kwargs['self'].device
-            kwargs['self'].device = kwargs['self']._device_ids[0]
-            rtn = func(**kwargs)
-            kwargs['self'].to(default_device)
-            return rtn
-    return wrapper
-
-def collate_handle_corrupted(samples_list, dataset, labels, dtype=torch.half):
-    # print(len(samples_list))
-    orig_len = len(samples_list)
-    # for the loss to be consistent, we drop samples with NaN values in any of their corresponding crops
-    for i, s in enumerate(samples_list):
-        ic(s is None)
-        if s is None:
-            continue
-    samples_list = list(filter(lambda x: x is not None, samples_list))
-
-    if len(samples_list) == 0:
-        ic('recursive call')
-        return collate_handle_corrupted([dataset[random.randint(0, len(dataset)-1)] for _ in range(orig_len)], dataset, labels)
-    
-    # collated_images = torch.stack([convert_to_tensor(s["image"]) for s in samples_list])
-    try:
-        if "image" in samples_list[0]:
-            samples_list = [s for s in samples_list if not torch.isnan(s["image"]).any()]
-            # print('samples list: ', len(samples_list))
-            collated_images = torch.stack([convert_to_tensor(s["image"]) for s in samples_list])
-            # print("here1")
-            collated_labels = {k: torch.Tensor([s["label"][k] if s["label"][k] is not None else 0 for s in samples_list]) for k in labels}
-            # print("here2")
-            collated_mask = {k: torch.Tensor([1 if s["label"][k] is not None else 0 for s in samples_list]) for k in labels}
-            # print("here3")
-            return {"image": collated_images,
-                    "label": collated_labels,
-                    "mask": collated_mask}
-    except:
-        return collate_handle_corrupted([dataset[random.randint(0, len(dataset)-1)] for _ in range(orig_len)], dataset, labels)
-    
-    
-     
-def get_backend(img_backend):
-    if img_backend == 'C3D':
-        return nn.C3D
-    elif img_backend == 'DenseNet':
-        return nn.DenseNet
-
-
-class ImagingModel(BaseEstimator):
-    ''' ... '''
-    def __init__(self,
-        tgt_modalities: list[str],
-        label_fractions: dict[str, float],
-        num_epochs: int = 32,
-        batch_size: int = 8,
-        batch_size_multiplier: int = 1,
-        lr: float = 1e-2,
-        weight_decay: float = 0.0,
-        beta: float = 0.9999,
-        gamma: float = 2.0,
-        bn_size: int = 4,
-        growth_rate: int = 12, 
-        block_config: tuple = (3, 3, 3), 
-        compression: float = 0.5,
-        num_init_features: int = 16, 
-        drop_rate: float = 0.2,
-        criterion: str | None = None,
-        device: str = 'cpu',
-        cuda_devices: list = [1],
-        ckpt_path: str = '/home/skowshik/ADRD_repo/adrd_tool/dev/ckpt/ckpt.pt',
-        load_from_ckpt: bool = True,
-        save_intermediate_ckpts: bool = False,
-        data_parallel: bool = False,
-        verbose: int = 0,
-        img_backend: str | None = None,
-        label_distribution: dict = {},
-        wandb_ = 1,
-        _device_ids: list | None = None,
-        _dataloader_num_workers: int = 4,
-        _amp_enabled: bool = False,
-    ) -> None:  
-        ''' ... '''
-        # for multiprocessing
-        self._rank = 0
-        self._lock = None
-
-        # positional parameters
-        self.tgt_modalities = tgt_modalities
-
-        # training parameters
-        self.label_fractions = label_fractions
-        self.num_epochs = num_epochs
-        self.batch_size = batch_size
-        self.batch_size_multiplier = batch_size_multiplier
-        self.lr = lr
-        self.weight_decay = weight_decay
-        self.beta = beta
-        self.gamma = gamma
-        self.bn_size = bn_size
-        self.growth_rate = growth_rate
-        self.block_config = block_config
-        self.compression = compression
-        self.num_init_features = num_init_features
-        self.drop_rate = drop_rate
-        self.criterion = criterion
-        self.device = device
-        self.cuda_devices = cuda_devices
-        self.ckpt_path = ckpt_path
-        self.load_from_ckpt = load_from_ckpt
-        self.save_intermediate_ckpts = save_intermediate_ckpts
-        self.data_parallel = data_parallel
-        self.verbose = verbose
-        self.img_backend = img_backend
-        self.label_distribution = label_distribution
-        self.wandb_ = wandb_
-        self._device_ids = _device_ids
-        self._dataloader_num_workers = _dataloader_num_workers
-        self._amp_enabled = _amp_enabled
-        self.scaler = torch.cuda.amp.GradScaler()
-
-    @_manage_ctx_fit
-    def fit(self, trn_list, vld_list, img_train_trans=None, img_vld_trans=None) -> Self:
-    # def fit(self, x, y) -> Self:
-        ''' ... '''
-        
-        # start a new wandb run to track this script
-        if self.wandb_ == 1:
-            wandb.init(
-                # set the wandb project where this run will be logged
-                project="ADRD_main",
-                
-                # track hyperparameters and run metadata
-                config={
-                "Model": "DenseNet",
-                "Loss": 'Focalloss',
-                "EMB": "ALL_EMB",
-                "epochs": 256,
-                }
-            )
-            wandb.run.log_code("/home/skowshik/ADRD_repo/pipeline_v1_main/adrd_tool")
-        else:
-            wandb.init(mode="disabled") 
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-        print(self.criterion)
-
-        # initialize neural network
-        self._init_net()
-
-        # for k, info in self.src_modalities.items():
-        #     if info['type'] == 'imaging' and self.img_net != 'EMB':
-        #         info['shape'] = (1,) + (self.img_size,) * 3
-        #         info['img_shape'] = (1,) + (self.img_size,) * 3
-                # print(info['shape'])
-
-        # initialize dataloaders
-        # ldr_trn, ldr_vld = self._init_dataloader(x, y)
-        # ldr_trn, ldr_vld = self._init_dataloader(x_trn, x_vld, y_trn, y_vld)
-        ldr_trn, ldr_vld = self._init_dataloader(trn_list, vld_list, img_train_trans=img_train_trans, img_vld_trans=img_vld_trans)
-
-        # initialize optimizer and scheduler
-        if not self.load_from_ckpt:
-            self.optimizer = self._init_optimizer()
-        self.scheduler = self._init_scheduler(self.optimizer)
-
-        # gradient scaler for AMP
-        if self._amp_enabled: 
-            self.scaler = torch.cuda.amp.GradScaler()
-
-        # initialize focal loss function 
-        self.loss_fn = {}
-            
-        for k in self.tgt_modalities:
-            if self.label_fractions[k] >= 0.3:
-                alpha = -1
-            else:
-                alpha = pow((1 - self.label_fractions[k]), 2)
-            # alpha = -1
-            self.loss_fn[k] = nn.SigmoidFocalLoss(
-                alpha = alpha,
-                gamma = self.gamma,
-                reduction = 'none'
-            )
-
-        # to record the best validation performance criterion
-        if self.criterion is not None: 
-            best_crit = None
-            best_crit_AUPR = None
-
-        # progress bar for epoch loops
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch = tqdm(
-                    desc = 'Rank {:02d}'.format(self._rank),
-                    total = self.num_epochs,
-                    position = self._rank,
-                    ascii = True,
-                    leave = False,
-                    bar_format='{l_bar}{r_bar}'
-                )
-
-        # Define a hook function to print and store the gradient of a layer
-        def print_and_store_grad(grad, grad_list):
-            grad_list.append(grad)
-            # print(grad)
-
-        # grad_list = []
-        # self.net_.modules_emb_src['img_MRI_T1'].downsample[0].weight.register_hook(lambda grad: print_and_store_grad(grad, grad_list))
-
-        # lambda_coeff = 0.0001
-        # margin_loss = torch.nn.MarginRankingLoss(reduction='sum', margin=0.05)
-        
-        # training loop
-        for epoch in range(self.start_epoch, self.num_epochs):
-            met_trn = self.train_one_epoch(ldr_trn, epoch)
-            met_vld = self.validate_one_epoch(ldr_vld, epoch)
-            
-            print(self.ckpt_path.split('/')[-1])
-
-            # save the model if it has the best validation performance criterion by far
-            if self.criterion is None: continue
-            
-                
-            # is current criterion better than previous best?
-            curr_crit = np.mean([met_vld[i][self.criterion] for i in range(len(self.tgt_modalities))])
-            curr_crit_AUPR = np.mean([met_vld[i]["AUC (PR)"] for i in range(len(self.tgt_modalities))])
-            # AUROC
-            if best_crit is None or np.isnan(best_crit):
-                is_better = True
-            elif self.criterion == 'Loss' and best_crit >= curr_crit:
-                is_better = True
-            elif self.criterion != 'Loss' and best_crit <= curr_crit :
-                is_better = True
-            else:
-                is_better = False
-
-            # AUPR
-            if best_crit_AUPR is None or np.isnan(best_crit_AUPR):
-                is_better_AUPR = True
-            elif best_crit_AUPR <= curr_crit_AUPR :
-                is_better_AUPR = True
-            else:
-                is_better_AUPR = False
-                
-            # update best criterion
-            if is_better_AUPR:
-                best_crit_AUPR = curr_crit_AUPR
-                if self.save_intermediate_ckpts:
-                    print(f"Saving the model to {self.ckpt_path[:-3]}_AUPR.pt...")
-                    self.save(self.ckpt_path[:-3]+"_AUPR.pt", epoch)
-                    
-            if is_better:
-                best_crit = curr_crit
-                best_state_dict = deepcopy(self.net_.state_dict())
-                if self.save_intermediate_ckpts:
-                    print(f"Saving the model to {self.ckpt_path}...")
-                    self.save(self.ckpt_path, epoch)
-
-            if self.verbose > 2:
-                print('Best {}: {}'.format(self.criterion, best_crit))
-                print('Best {}: {}'.format('AUC (PR)', best_crit_AUPR))
-
-            if self.verbose == 1:
-                with self._lock if self._lock is not None else suppress():
-                    pbr_epoch.update(1)
-                    pbr_epoch.refresh()
-
-        return self
-    
-    def train_one_epoch(self, ldr_trn, epoch):
-        
-        # progress bar for batch loops
-        if self.verbose > 1: 
-            pbr_batch = ProgressBar(len(ldr_trn.dataset), 'Epoch {:03d} (TRN)'.format(epoch))
-        
-        torch.set_grad_enabled(True)
-        self.net_.train()
-        
-        scores_trn, y_true_trn, y_mask_trn = [], [], []
-        losses_trn = [[] for _ in self.tgt_modalities]
-        iters = len(ldr_trn)
-        print(iters)
-        for n_iter, batch_data in enumerate(ldr_trn):
-            # if len(batch_data["image"]) < self.batch_size:
-            #     continue
-            
-            x_batch = batch_data["image"].to(self.device, non_blocking=True)
-            y_batch = {k: v.to(self.device, non_blocking=True) for k,v in batch_data["label"].items()}
-            y_mask = {k: v.to(self.device, non_blocking=True) for k,v in batch_data["mask"].items()}
-            
-            with torch.autocast(
-                device_type = 'cpu' if self.device == 'cpu' else 'cuda',
-                dtype = torch.bfloat16 if self.device == 'cpu' else torch.float16,
-                enabled = self._amp_enabled,
-            ):
-
-                outputs = self.net_(x_batch, shap=False)
-                # print(outputs.shape)
-                # calculate multitask loss
-                loss = 0
-                for i, k in enumerate(self.tgt_modalities):
-                    loss_task = self.loss_fn[k](outputs[k], y_batch[k])
-                    msk_loss_task = loss_task * y_mask[k]
-                    msk_loss_mean = msk_loss_task.sum() / y_mask[k].sum()
-                    loss += msk_loss_mean
-                    losses_trn[i] += msk_loss_task.detach().cpu().numpy().tolist()
-
-            # backward
-            if self._amp_enabled:
-                self.scaler.scale(loss).backward()
-            else:
-                loss.backward()
-
-            # print(len(grad_list), len(grad_list[-1]))
-            # print(f"Gradient at {n_iter}: {grad_list[-1][0]}")
-            
-            # update parameters
-            if n_iter != 0 and n_iter % self.batch_size_multiplier == 0:
-                if self._amp_enabled:
-                    self.scaler.step(self.optimizer)
-                    self.scaler.update()
-                    self.optimizer.zero_grad()
-                else: 
-                    self.optimizer.step()
-                    self.optimizer.zero_grad()
-            # set self.scheduler
-            self.scheduler.step(epoch + n_iter / iters)
-            # print(f"Weight: {self.net_.module.features[0].weight[0]}")
-
-            ''' TODO: change array to dictionary later '''
-            outputs = torch.stack(list(outputs.values()), dim=1)
-            y_batch = torch.stack(list(y_batch.values()), dim=1)
-            y_mask = torch.stack(list(y_mask.values()), dim=1)
-
-            # save outputs to evaluate performance later
-            scores_trn.append(outputs.detach().to(torch.float).cpu())
-            y_true_trn.append(y_batch.cpu())
-            y_mask_trn.append(y_mask.cpu())
-            
-            # log metrics to wandb
-
-            # update progress bar
-            if self.verbose > 1:
-                batch_size = len(x_batch)
-                pbr_batch.update(batch_size, {})
-                pbr_batch.refresh()
-
-            # clear cuda cache
-            if "cuda" in self.device:
-                torch.cuda.empty_cache()
-
-        # for better tqdm progress bar display
-        if self.verbose > 1:
-            pbr_batch.close()
-
-        # # set self.scheduler
-        # self.scheduler.step()
-
-        # calculate and print training performance metrics
-        scores_trn = torch.cat(scores_trn)
-        y_true_trn = torch.cat(y_true_trn)
-        y_mask_trn = torch.cat(y_mask_trn)
-        y_pred_trn = (scores_trn > 0).to(torch.int)
-        y_prob_trn = torch.sigmoid(scores_trn)
-        met_trn = get_metrics_multitask(
-            y_true_trn.numpy(),
-            y_pred_trn.numpy(),
-            y_prob_trn.numpy(),
-            y_mask_trn.numpy()
-        )
-
-        # add loss to metrics
-        for i in range(len(self.tgt_modalities)):
-            met_trn[i]['Loss'] = np.mean(losses_trn[i])
-        
-        wandb.log({f"Train loss {list(self.tgt_modalities)[i]}": met_trn[i]['Loss']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        wandb.log({f"Train Balanced Accuracy {list(self.tgt_modalities)[i]}": met_trn[i]['Balanced Accuracy']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        
-        wandb.log({f"Train AUC (ROC) {list(self.tgt_modalities)[i]}": met_trn[i]['AUC (ROC)']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        wandb.log({f"Train AUPR {list(self.tgt_modalities)[i]}": met_trn[i]['AUC (PR)']  for i in range(len(self.tgt_modalities))}, step=epoch)
-
-        if self.verbose > 2:
-            print_metrics_multitask(met_trn)
-        
-        return met_trn
-    
-    # @torch.no_grad()
-    def validate_one_epoch(self, ldr_vld, epoch):
-        # progress bar for validation
-        if self.verbose > 1:
-            pbr_batch = ProgressBar(len(ldr_vld.dataset), 'Epoch {:03d} (VLD)'.format(epoch))
-
-        # set model to validation mode
-        torch.set_grad_enabled(False)
-        self.net_.eval()
-
-        scores_vld, y_true_vld, y_mask_vld = [], [], []
-        losses_vld = [[] for _ in self.tgt_modalities]
-        for batch_data in ldr_vld:
-            # if len(batch_data["image"]) < self.batch_size:
-            #     continue
-            x_batch = batch_data["image"].to(self.device, non_blocking=True)
-            y_batch = {k: v.to(self.device, non_blocking=True) for k,v in batch_data["label"].items()}
-            y_mask = {k: v.to(self.device, non_blocking=True) for k,v in batch_data["mask"].items()}
-            
-            # forward
-            with torch.autocast(
-                device_type = 'cpu' if self.device == 'cpu' else 'cuda',
-                dtype = torch.bfloat16 if self.device == 'cpu' else torch.float16,
-                enabled = self._amp_enabled
-            ):
-                
-                outputs = self.net_(x_batch, shap=False)
-
-                # calculate multitask loss
-                for i, k in enumerate(self.tgt_modalities):
-                    loss_task = self.loss_fn[k](outputs[k], y_batch[k])
-                    msk_loss_task = loss_task * y_mask[k]
-                    losses_vld[i] += msk_loss_task.detach().cpu().numpy().tolist()
-
-            ''' TODO: change array to dictionary later '''
-            outputs = torch.stack(list(outputs.values()), dim=1)
-            y_batch = torch.stack(list(y_batch.values()), dim=1)
-            y_mask = torch.stack(list(y_mask.values()), dim=1)
-
-            # save outputs to evaluate performance later
-            scores_vld.append(outputs.detach().to(torch.float).cpu())
-            y_true_vld.append(y_batch.cpu())
-            y_mask_vld.append(y_mask.cpu())
-
-            # update progress bar
-            if self.verbose > 1:
-                batch_size = len(x_batch)
-                pbr_batch.update(batch_size, {})
-                pbr_batch.refresh()
-
-            # clear cuda cache
-            if "cuda" in self.device:
-                torch.cuda.empty_cache()
-
-        # for better tqdm progress bar display
-        if self.verbose > 1:
-            pbr_batch.close()
-
-        # calculate and print validation performance metrics
-        scores_vld = torch.cat(scores_vld)
-        y_true_vld = torch.cat(y_true_vld)
-        y_mask_vld = torch.cat(y_mask_vld)
-        y_pred_vld = (scores_vld > 0).to(torch.int)
-        y_prob_vld = torch.sigmoid(scores_vld)
-        met_vld = get_metrics_multitask(
-            y_true_vld.numpy(),
-            y_pred_vld.numpy(),
-            y_prob_vld.numpy(),
-            y_mask_vld.numpy()
-        )
-
-        # add loss to metrics
-        for i in range(len(self.tgt_modalities)):
-            met_vld[i]['Loss'] = np.mean(losses_vld[i])
-            
-        wandb.log({f"Validation loss {list(self.tgt_modalities)[i]}": met_vld[i]['Loss'] for i in range(len(self.tgt_modalities))}, step=epoch)
-        wandb.log({f"Validation Balanced Accuracy {list(self.tgt_modalities)[i]}": met_vld[i]['Balanced Accuracy']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        
-        wandb.log({f"Validation AUC (ROC) {list(self.tgt_modalities)[i]}": met_vld[i]['AUC (ROC)']  for i in range(len(self.tgt_modalities))}, step=epoch)
-        wandb.log({f"Validation AUPR {list(self.tgt_modalities)[i]}": met_vld[i]['AUC (PR)']  for i in range(len(self.tgt_modalities))}, step=epoch)
-
-        if self.verbose > 2:
-            print_metrics_multitask(met_vld)
-        
-        return met_vld
-    
-
-    def save(self, filepath: str, epoch: int = 0) -> None:
-        ''' ... '''
-        check_is_fitted(self)
-        if self.data_parallel:
-            state_dict = self.net_.module.state_dict()
-        else:
-            state_dict = self.net_.state_dict()
-
-        # attach model hyper parameters
-        state_dict['tgt_modalities'] = self.tgt_modalities
-        state_dict['optimizer'] = self.optimizer
-        state_dict['bn_size'] = self.bn_size
-        state_dict['growth_rate'] = self.growth_rate
-        state_dict['block_config'] = self.block_config
-        state_dict['compression'] = self.compression
-        state_dict['num_init_features'] = self.num_init_features
-        state_dict['drop_rate'] = self.drop_rate
-        state_dict['epoch'] = epoch
-
-        if self.scaler is not None:
-            state_dict['scaler'] = self.scaler.state_dict()
-        if self.label_distribution:
-            state_dict['label_distribution'] = self.label_distribution
-
-        torch.save(state_dict, filepath)
-
-    def load(self, filepath: str, map_location: str = 'cpu', how='latest') -> None:
-        ''' ... '''
-        # load state_dict
-        if how == 'latest':
-            if torch.load(filepath)['epoch'] > torch.load(f'{filepath[:-3]}_AUPR.pt')['epoch']:
-                print("Loading model saved using AUROC")
-                state_dict = torch.load(filepath, map_location=map_location)
-            else:
-                print("Loading model saved using AUPR")
-                state_dict = torch.load(f'{filepath[:-3]}_AUPR.pt', map_location=map_location)
-        else:
-            state_dict = torch.load(filepath, map_location=map_location)
-
-        # load data modalities
-        self.tgt_modalities: dict[str, dict[str, Any]] = state_dict.pop('tgt_modalities')
-        if 'label_distribution' in state_dict:
-            self.label_distribution: dict[str, dict[int, int]] = state_dict.pop('label_distribution')
-        if 'optimizer' in state_dict:
-            self.optimizer = state_dict.pop('optimizer')
-        if 'bn_size' in state_dict:
-            self.bn_size = state_dict.pop('bn_size')
-        if 'growth_rate' in state_dict:
-            self.growth_rate = state_dict.pop('growth_rate')
-        if 'block_config' in state_dict:
-            self.block_config = state_dict.pop('block_config')
-        if 'compression' in state_dict:
-            self.compression = state_dict.pop('compression')
-        if 'num_init_features' in state_dict:
-            self.num_init_features = state_dict.pop('num_init_features')
-        if 'drop_rate' in state_dict:
-            self.drop_rate = state_dict.pop('drop_rate')
-        if 'epoch' in state_dict:
-            self.start_epoch = state_dict.pop('epoch')
-            print(f'Epoch: {self.start_epoch}')
-
-        # initialize model
-
-        self.net_ = get_backend(self.img_backend)( 
-                tgt_modalities = self.tgt_modalities,
-                bn_size = self.bn_size,
-                growth_rate=self.growth_rate, 
-                block_config=self.block_config, 
-                compression=self.compression,
-                num_init_features=self.num_init_features, 
-                drop_rate=self.drop_rate,
-                load_from_ckpt=self.load_from_ckpt
-            )
-        print(self.net_)
-            
-        if 'scaler' in state_dict and state_dict['scaler']:
-            self.scaler.load_state_dict(state_dict.pop('scaler'))
-        self.net_.load_state_dict(state_dict)
-        check_is_fitted(self)
-        self.net_.to(self.device)
-
-    def to(self, device: str) -> Self:
-        ''' Mount model to the given device. '''
-        self.device = device
-        if hasattr(self, 'model'): self.net_ = self.net_.to(device)
-        return self
-    
-    @classmethod
-    def from_ckpt(cls, filepath: str, device='cpu', img_backend=None, load_from_ckpt=True, how='latest') -> Self:
-        ''' ... '''
-        obj = cls(None, None, None,device=device)
-        if device == 'cuda':
-            obj.device = "{}:{}".format(obj.device, str(obj.cuda_devices[0]))
-        print(obj.device)
-        obj.img_backend=img_backend
-        obj.load_from_ckpt = load_from_ckpt
-        obj.load(filepath, map_location=obj.device, how=how)
-        return obj
-    
-    def _init_net(self):
-        """ ... """
-        self.start_epoch = 0
-        # set the device for use
-        if self.device == 'cuda':
-            self.device = "{}:{}".format(self.device, str(self.cuda_devices[0]))
-        # self.load(self.ckpt_path, map_location=self.device)
-        # print("Loading model from checkpoint...")
-        # self.load(self.ckpt_path, map_location=self.device)
-
-        if self.load_from_ckpt:
-            try:
-                print("Loading model from checkpoint...")
-                self.load(self.ckpt_path, map_location=self.device)
-            except:
-                print("Cannot load from checkpoint. Initializing new model...")
-                self.load_from_ckpt = False
-
-        if not self.load_from_ckpt:
-            self.net_ = get_backend(self.img_backend)( 
-                tgt_modalities = self.tgt_modalities,
-                bn_size = self.bn_size,
-                growth_rate=self.growth_rate, 
-                block_config=self.block_config, 
-                compression=self.compression,
-                num_init_features=self.num_init_features, 
-                drop_rate=self.drop_rate,
-                load_from_ckpt=self.load_from_ckpt
-            ) 
-            
-            # # intialize model parameters using xavier_uniform
-            # for p in self.net_.parameters():
-            #     if p.dim() > 1:
-            #         torch.nn.init.xavier_uniform_(p)
-        
-        self.net_.to(self.device)
-
-        # Initialize the number of GPUs
-        if self.data_parallel and torch.cuda.device_count() > 1:
-            print("Available", torch.cuda.device_count(), "GPUs!")
-            self.net_ = torch.nn.DataParallel(self.net_, device_ids=self.cuda_devices)
-
-        # return net
-
-    def _init_dataloader(self, trn_list, vld_list, img_train_trans=None, img_vld_trans=None):
-    # def _init_dataloader(self, x, y):
-        """ ... """
-        # # split dataset
-        # x_trn, x_vld, y_trn, y_vld = train_test_split(
-        #     x, y, test_size = 0.2, random_state = 0,
-        # )
-
-        # # initialize dataset and dataloader
-        # dat_trn = CNNTrainingValidationDataset(
-        #     x_trn, y_trn,
-        #     self.tgt_modalities,
-        #     img_transform=img_train_trans,
-        # )
-
-        # dat_vld = CNNTrainingValidationDataset(
-        #     x_vld, y_vld,
-        #     self.tgt_modalities,
-        #     img_transform=img_vld_trans,
-        # )
-        
-        dat_trn = monai.data.Dataset(data=trn_list, transform=img_train_trans)
-        dat_vld = monai.data.Dataset(data=vld_list, transform=img_vld_trans)
-        collate_fn_trn = functools.partial(collate_handle_corrupted, dataset=dat_trn, dtype=torch.FloatTensor, labels=self.tgt_modalities)
-        collate_fn_vld = functools.partial(collate_handle_corrupted, dataset=dat_vld, dtype=torch.FloatTensor, labels=self.tgt_modalities)
-
-        ldr_trn = DataLoader(
-            dataset = dat_trn,
-            batch_size = self.batch_size,
-            shuffle = True,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = collate_fn_trn,
-            # pin_memory = True
-        )
-
-        ldr_vld = DataLoader(
-            dataset = dat_vld,
-            batch_size = self.batch_size,
-            shuffle = False,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = collate_fn_vld,
-            # pin_memory = True
-        )
-
-        return ldr_trn, ldr_vld
-    
-    def _init_optimizer(self):
-        """ ... """
-        params = list(self.net_.parameters()) 
-        # for p in params:
-        #     print(p.requires_grad)
-        return torch.optim.AdamW(
-            params,
-            lr = self.lr,
-            betas = (0.9, 0.98),
-            weight_decay = self.weight_decay
-        )
-    
-    def _init_scheduler(self, optimizer):
-        """ ... """
-        # return torch.optim.lr_scheduler.OneCycleLR(
-        #     optimizer = optimizer, 
-        #     max_lr = self.lr,
-        #     total_steps = self.num_epochs,
-        #     verbose = (self.verbose > 2)
-        # )
-
-        # return torch.optim.lr_scheduler.CosineAnnealingLR(
-        #     optimizer=optimizer, 
-        #     T_max=64, 
-        #     verbose=(self.verbose > 2)
-        # )
-
-        return torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
-            optimizer=optimizer,
-            T_0=64,
-            T_mult=2,
-            eta_min = 0,
-            verbose=(self.verbose > 2)
-        )
-    
-    def _init_loss_func(self, 
-        num_per_cls: dict[str, tuple[int, int]],
-    ) -> dict[str, Module]:
-        """ ... """
-        return {k: nn.SigmoidFocalLossBeta(
-            beta = self.beta,
-            gamma = self.gamma,
-            num_per_cls = num_per_cls[k],
-            reduction = 'none',
-        ) for k in self.tgt_modalities}
-    
-    def _proc_fit(self):
-        """ ... """
-        
-    def _init_test_dataloader(self, batch_size, tst_list, img_tst_trans=None):
-        # input validation
-        check_is_fitted(self)
-        print(self.device)
-        
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # set model to eval mode
-        torch.set_grad_enabled(False)
-        self.net_.eval()
-        
-        dat_tst = monai.data.Dataset(data=tst_list, transform=img_tst_trans)
-        collate_fn_tst = functools.partial(collate_handle_corrupted, dataset=dat_tst, dtype=torch.FloatTensor, labels=self.tgt_modalities)
-        # print(collate_fn_tst)
-
-        ldr_tst = DataLoader(
-            dataset = dat_tst,
-            batch_size = batch_size,
-            shuffle = False,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = collate_fn_tst,
-            # pin_memory = True
-        )
-        return ldr_tst
-
-    
-    def predict_logits(self,
-        ldr_tst: Any | None = None,
-    ) -> list[dict[str, float]]:
-        
-        # run model and collect results
-        logits: list[dict[str, float]] = []
-        for batch_data in tqdm(ldr_tst):
-            # print(batch_data["image"])
-            if len(batch_data) == 0:
-                continue
-            x_batch = batch_data["image"].to(self.device, non_blocking=True)
-            outputs = self.net_(x_batch, shap=False)
-            
-            # convert output from dict-of-list to list of dict, then append
-            tmp = {k: outputs[k].tolist() for k in self.tgt_modalities}
-            tmp = [{k: tmp[k][i] for k in self.tgt_modalities} for i in range(len(next(iter(tmp.values()))))]
-            logits += tmp
-
-        return logits
-        
-    def predict_proba(self,
-        ldr_tst: Any | None = None,
-        temperature: float = 1.0,
-    ) -> list[dict[str, float]]:
-        ''' ... '''
-        logits = self.predict_logits(ldr_tst)
-        print("got logits")
-        return logits, [{k: expit(smp[k] / temperature) for k in self.tgt_modalities} for smp in logits]
-
-    def predict(self,
-        ldr_tst: Any | None = None,
-    ) -> list[dict[str, int]]:
-        ''' ... '''
-        logits, proba = self.predict_proba(ldr_tst)
-        print("got proba")
-        return logits, proba, [{k: int(smp[k] > 0.5) for k in self.tgt_modalities} for smp in proba]

adrd/model/train_resnet.py

@@ -1,484 +0,0 @@
-import torch
-import numpy as np
-import tqdm
-from sklearn.base import BaseEstimator
-from sklearn.utils.validation import check_is_fitted
-from sklearn.model_selection import train_test_split
-from scipy.special import expit
-from copy import deepcopy
-from contextlib import suppress
-from typing import Any, Self
-from icecream import ic
-
-from .. import nn
-from ..utils import TransformerTrainingDataset
-from ..utils import TransformerValidationDataset
-from ..utils import MissingMasker
-from ..utils import ConstantImputer
-from ..utils import Formatter
-from ..utils.misc import ProgressBar
-from ..utils.misc import get_metrics_multitask, print_metrics_multitask
-
-
-class TrainResNet(BaseEstimator):
-    ''' ... '''
-    def __init__(self,
-        src_modalities: dict[str, dict[str, Any]],
-        tgt_modalities: dict[str, dict[str, Any]],
-        label_fractions: dict[str, float],
-        num_epochs: int = 32,
-        batch_size: int = 8,
-        lr: float = 1e-2,
-        weight_decay: float = 0.0,
-        gamma: float = 0.0,
-        criterion: str | None = None,
-        device: str = 'cpu',
-        cuda_devices: list = [1,2],
-        mri_feature: str = 'img_MRI_T1',
-        ckpt_path: str = '/home/skowshik/ADRD_repo/adrd_tool/adrd/dev/ckpt/ckpt.pt',
-        load_from_ckpt: bool = True,
-        save_intermediate_ckpts: bool = False,
-        data_parallel: bool = False,
-        verbose: int = 0,
-    ):  
-        ''' ... '''
-        # for multiprocessing
-        self._rank = 0
-        self._lock = None
-
-        # positional parameters
-        self.src_modalities = src_modalities
-        self.tgt_modalities = tgt_modalities
-
-        # training parameters
-        self.label_fractions = label_fractions
-        self.num_epochs = num_epochs
-        self.batch_size = batch_size
-        self.lr = lr
-        self.weight_decay = weight_decay
-        self.gamma = gamma
-        self.criterion = criterion
-        self.device = device
-        self.cuda_devices = cuda_devices
-        self.mri_feature = mri_feature
-        self.ckpt_path = ckpt_path
-        self.load_from_ckpt = load_from_ckpt
-        self.save_intermediate_ckpts = save_intermediate_ckpts
-        self.data_parallel = data_parallel
-        self.verbose = verbose
-
-    def fit(self, x, y):
-        ''' ... '''
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # set the device for use
-        if self.device == 'cuda':
-            self.device = "{}:{}".format(self.device, str(self.cuda_devices[0]))
-
-        # initialize model
-        if self.load_from_ckpt:
-            try:
-                print("Loading model from checkpoint...")
-                self.load(self.ckpt_path, map_location=self.device)
-            except:
-                print("Cannot load from checkpoint. Initializing new model...")
-                self.load_from_ckpt = False
-
-        # initialize model
-        if not self.load_from_ckpt:
-            self.net_ = nn.ResNetModel(
-                self.tgt_modalities,
-                mri_feature = self.mri_feature
-            )
-            # intialize model parameters using xavier_uniform
-            for p in self.net_.parameters():
-                if p.dim() > 1:
-                    torch.nn.init.xavier_uniform_(p)
-
-        self.net_.to(self.device)
-
-        # Initialize the number of GPUs
-        if self.data_parallel and torch.cuda.device_count() > 1:
-            print("Available", torch.cuda.device_count(), "GPUs!")
-            self.net_ = torch.nn.DataParallel(self.net_, device_ids=self.cuda_devices)
-
-
-        # split dataset
-        x_trn, x_vld, y_trn, y_vld = train_test_split(
-            x, y, test_size = 0.2, random_state = 0,
-        )
-
-        # initialize dataset and dataloader
-        dat_trn = TransformerTrainingDataset(
-            x_trn, y_trn,
-            self.src_modalities,
-            self.tgt_modalities,
-            dropout_rate = .5,
-            dropout_strategy = 'compensated',
-            mri_feature = self.mri_feature,
-        )
-
-        dat_vld = TransformerValidationDataset(
-            x_vld, y_vld,
-            self.src_modalities,
-            self.tgt_modalities,
-            mri_feature = self.mri_feature,
-        )
-
-        # ic(dat_trn[0])
-
-        ldr_trn = torch.utils.data.DataLoader(
-            dat_trn,
-            batch_size = self.batch_size,
-            shuffle = True,
-            drop_last = False,
-            num_workers = 0,
-            collate_fn = TransformerTrainingDataset.collate_fn,
-            # pin_memory = True
-        )
-
-        ldr_vld = torch.utils.data.DataLoader(
-            dat_vld,
-            batch_size = self.batch_size,
-            shuffle = False,
-            drop_last = False,
-            num_workers = 0,
-            collate_fn = TransformerTrainingDataset.collate_fn,
-            # pin_memory = True
-        )
-
-        # initialize optimizer
-        optimizer = torch.optim.AdamW(
-            self.net_.parameters(),
-            lr = self.lr,
-            betas = (0.9, 0.98),
-            weight_decay = self.weight_decay
-        )
-        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=64, verbose=(self.verbose > 2))
-        
-        # initialize loss function (binary cross entropy)
-        loss_fn = {}
-
-        for k in self.tgt_modalities:
-            alpha = pow((1 - self.label_fractions[k]), self.gamma)
-            # if alpha < 0.5:
-            #     alpha = -1
-            loss_fn[k] = nn.SigmoidFocalLoss(
-                alpha = alpha,
-                gamma = self.gamma,
-                reduction = 'none'
-            )
-
-        # to record the best validation performance criterion
-        if self.criterion is not None:
-            best_crit = None
-
-        # progress bar for epoch loops
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch = tqdm.tqdm(
-                    desc = 'Rank {:02d}'.format(self._rank),
-                    total = self.num_epochs,
-                    position = self._rank,
-                    ascii = True,
-                    leave = False,
-                    bar_format='{l_bar}{r_bar}'
-                )
-
-        # Define a hook function to print and store the gradient of a layer
-        def print_and_store_grad(grad, grad_list):
-            grad_list.append(grad)
-            # print(grad)
-
-        # grad_list = []
-        # self.net_.module.img_net_.featurizer.down_tr64.ops[0].conv1.weight.register_hook(lambda grad: print_and_store_grad(grad, grad_list))
-        # self.net_.module.modules_emb_src['gender'].weight.register_hook(lambda grad: print_and_store_grad(grad, grad_list))
-
-
-        # training loop
-        for epoch in range(self.num_epochs):
-            # progress bar for batch loops
-            if self.verbose > 1: 
-                pbr_batch = ProgressBar(len(dat_trn), 'Epoch {:03d} (TRN)'.format(epoch))
-
-            # set model to train mode
-            torch.set_grad_enabled(True)
-            self.net_.train()
-
-            scores_trn, y_true_trn = [], []
-            losses_trn = [[] for _ in self.tgt_modalities]
-            for x_batch, y_batch, mask in ldr_trn:
-
-                # mount data to the proper device
-                x_batch = {k: x_batch[k].to(self.device) for k in x_batch}
-                y_batch = {k: y_batch[k].to(torch.float).to(self.device) for k in y_batch}
-
-                # forward
-                outputs = self.net_(x_batch)
-
-                # calculate multitask loss
-                loss = 0
-                for i, k in enumerate(self.tgt_modalities):
-                    loss_task = loss_fn[k](outputs[k], y_batch[k])
-                    loss += loss_task.mean()
-                    losses_trn[i] += loss_task.detach().cpu().numpy().tolist()
-
-                # backward
-                optimizer.zero_grad(set_to_none=True)
-                loss.backward()
-                optimizer.step()
-
-                ''' TODO: change array to dictionary later '''
-                outputs = torch.stack(list(outputs.values()), dim=1)
-                y_batch = torch.stack(list(y_batch.values()), dim=1)
-
-                # save outputs to evaluate performance later
-                scores_trn.append(outputs.detach().to(torch.float).cpu())
-                y_true_trn.append(y_batch.cpu())
-
-                # update progress bar
-                if self.verbose > 1:
-                    batch_size = len(next(iter(x_batch.values())))
-                    pbr_batch.update(batch_size, {})
-                    pbr_batch.refresh()
-                
-                # clear cuda cache
-                if "cuda" in self.device:
-                    torch.cuda.empty_cache()
-
-            # for better tqdm progress bar display
-            if self.verbose > 1:
-                pbr_batch.close()
-
-            # set scheduler
-            scheduler.step()
-
-            # calculate and print training performance metrics
-            scores_trn = torch.cat(scores_trn)
-            y_true_trn = torch.cat(y_true_trn)
-            y_pred_trn = (scores_trn > 0).to(torch.int)
-            y_prob_trn = torch.sigmoid(scores_trn)
-            met_trn = get_metrics_multitask(
-                y_true_trn.numpy(),
-                y_pred_trn.numpy(),
-                y_prob_trn.numpy()
-            )
-
-            # add loss to metrics
-            for i in range(len(self.tgt_modalities)):
-                met_trn[i]['Loss'] = np.mean(losses_trn[i])
-
-            if self.verbose > 2:
-                print_metrics_multitask(met_trn)
-
-            # progress bar for validation
-            if self.verbose > 1:
-                pbr_batch = ProgressBar(len(dat_vld), 'Epoch {:03d} (VLD)'.format(epoch))
-
-            # set model to validation mode
-            torch.set_grad_enabled(False)
-            self.net_.eval()
-
-            scores_vld, y_true_vld = [], []
-            losses_vld = [[] for _ in self.tgt_modalities]
-            for x_batch, y_batch, mask in ldr_vld:
-                # mount data to the proper device
-                x_batch = {k: x_batch[k].to(self.device) for k in x_batch}
-                y_batch = {k: y_batch[k].to(torch.float).to(self.device) for k in y_batch}
-
-                # forward
-                outputs = self.net_(x_batch)
-
-                # calculate multitask loss
-                for i, k in enumerate(self.tgt_modalities):
-                    loss_task = loss_fn[k](outputs[k], y_batch[k])
-                    losses_vld[i] += loss_task.detach().cpu().numpy().tolist()
-
-                ''' TODO: change array to dictionary later '''
-                outputs = torch.stack(list(outputs.values()), dim=1)
-                y_batch = torch.stack(list(y_batch.values()), dim=1)
-
-                # save outputs to evaluate performance later
-                scores_vld.append(outputs.detach().to(torch.float).cpu())
-                y_true_vld.append(y_batch.cpu())
-
-                # update progress bar
-                if self.verbose > 1:
-                    batch_size = len(next(iter(x_batch.values())))
-                    pbr_batch.update(batch_size, {})
-                    pbr_batch.refresh()
-
-                # clear cuda cache
-                if "cuda" in self.device:
-                    torch.cuda.empty_cache()
-
-            # for better tqdm progress bar display
-            if self.verbose > 1:
-                pbr_batch.close()
-
-            # calculate and print validation performance metrics
-            scores_vld = torch.cat(scores_vld)
-            y_true_vld = torch.cat(y_true_vld)
-            y_pred_vld = (scores_vld > 0).to(torch.int)
-            y_prob_vld = torch.sigmoid(scores_vld)
-            met_vld = get_metrics_multitask(
-                y_true_vld.numpy(),
-                y_pred_vld.numpy(),
-                y_prob_vld.numpy()
-            )
-
-            # add loss to metrics
-            for i in range(len(self.tgt_modalities)):
-                met_vld[i]['Loss'] = np.mean(losses_vld[i])
-
-            if self.verbose > 2:
-                print_metrics_multitask(met_vld)
-
-            # save the model if it has the best validation performance criterion by far
-            if self.criterion is None: continue
-            
-            # is current criterion better than previous best?
-            curr_crit = np.mean([met_vld[i][self.criterion] for i in range(len(self.tgt_modalities))])
-            if best_crit is None or np.isnan(best_crit):
-                is_better = True
-            elif self.criterion == 'Loss' and best_crit >= curr_crit:
-                is_better = True
-            elif self.criterion != 'Loss' and best_crit <= curr_crit:
-                is_better = True
-            else:
-                is_better = False
-
-            # update best criterion
-            if is_better:
-                best_crit = curr_crit
-                best_state_dict = deepcopy(self.net_.state_dict())
-                if self.save_intermediate_ckpts:
-                    print("Saving the model...")
-                    self.save(self.ckpt_path)
-
-            if self.verbose > 2:
-                print('Best {}: {}'.format(self.criterion, best_crit))
-
-            if self.verbose == 1:
-                with self._lock if self._lock is not None else suppress():
-                    pbr_epoch.update(1)
-                    pbr_epoch.refresh()
-
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch.close()
-
-        # restore the model of the best validation performance across all epoches
-        if ldr_vld is not None and self.criterion is not None:
-            self.net_.load_state_dict(best_state_dict)
-
-        return self
-
-    def predict_logits(self,
-        x: list[dict[str, Any]],
-    ) -> list[dict[str, float]]:
-        '''
-        The input x can be a single sample or a list of samples.
-        '''
-        # input validation
-        check_is_fitted(self)
-        
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # set model to eval mode
-        torch.set_grad_enabled(False)
-        self.net_.eval()
-
-        # number of samples to evaluate
-        n_samples = len(x)
-
-        # format x
-        fmt = Formatter(self.src_modalities)
-        x = [fmt(smp) for smp in x]
-
-        # generate missing mask (BEFORE IMPUTATION)
-        msk = MissingMasker(self.src_modalities)
-        mask = [msk(smp) for smp in x]
-
-        # reformat x and then impute by 0s
-        imp = ConstantImputer(self.src_modalities)
-        x = [imp(smp) for smp in x]
-
-        # convert list-of-dict to dict-of-list
-        x = {k: [smp[k] for smp in x] for k in self.src_modalities}
-        mask = {k: [smp[k] for smp in mask] for k in self.src_modalities}
-        
-        # to tensor
-        x = {k: torch.as_tensor(np.array(v)).to(self.device) for k, v in x.items()}
-        mask = {k: torch.as_tensor(np.array(v)).to(self.device) for k, v in mask.items()}
-
-        # calculate logits
-        logits = self.net_(x)
-
-        # convert dict-of-list to list-of-dict
-        logits = {k: logits[k].tolist() for k in self.tgt_modalities}
-        logits = [{k: logits[k][i] for k in self.tgt_modalities} for i in range(n_samples)]
-
-        return logits
-        
-    def predict_proba(self,
-        x: list[dict[str, Any]],
-        temperature: float = 1.0
-    ) -> list[dict[str, float]]:
-        ''' ... '''
-        # calculate logits
-        logits = self.predict_logits(x)
-        
-        # convert logits to probabilities and 
-        proba = [{k: expit(smp[k] / temperature) for k in self.tgt_modalities} for smp in logits]
-        return proba
-
-    def predict(self,
-        x: list[dict[str, Any]],
-    ) -> list[dict[str, int]]:
-        ''' ... '''
-        proba = self.predict_proba(x)
-        return [{k: int(smp[k] > 0.5) for k in self.tgt_modalities} for smp in proba]
-
-    def save(self, filepath: str) -> None:
-        ''' ... '''
-        check_is_fitted(self)
-        if self.data_parallel:
-            state_dict = self.net_.module.state_dict()
-        else:
-            state_dict = self.net_.state_dict()
-
-        # attach model hyper parameters
-        state_dict['src_modalities'] = self.src_modalities
-        state_dict['tgt_modalities'] = self.tgt_modalities
-        state_dict['mri_feature'] = self.mri_feature
-
-        torch.save(state_dict, filepath)
-
-    def load(self, filepath: str, map_location: str='cpu') -> None:
-        ''' ... '''
-        # load state_dict
-        state_dict = torch.load(filepath, map_location=map_location)
-
-        # load data modalities
-        self.src_modalities = state_dict.pop('src_modalities')
-        self.tgt_modalities = state_dict.pop('tgt_modalities')
-
-        # initialize model
-        self.net_ = nn.ResNetModel(
-                self.tgt_modalities,
-                mri_feature = state_dict.pop('mri_feature')
-            )
-
-        # load model parameters
-        self.net_.load_state_dict(state_dict)
-        self.net_.to(self.device) 
-
-    @classmethod
-    def from_ckpt(cls, filepath: str, device='cpu') -> Self:
-        ''' ... '''
-        obj = cls(None, None, None,device=device)
-        obj.load(filepath)
-        return obj

adrd/model/transformer.py

@@ -1,600 +0,0 @@
-__all__ = ['Transformer']
-
-import torch
-from torch.utils.data import DataLoader
-import numpy as np
-import tqdm
-from sklearn.base import BaseEstimator
-from sklearn.utils.validation import check_is_fitted
-from sklearn.model_selection import train_test_split
-from scipy.special import expit
-from copy import deepcopy
-from contextlib import suppress
-from typing import Any, Self, Type
-from functools import wraps
-Tensor = Type[torch.Tensor]
-Module = Type[torch.nn.Module]
-
-from .. import nn
-from ..utils import TransformerTrainingDataset
-from ..utils import TransformerBalancedTrainingDataset
-from ..utils import Transformer2ndOrderBalancedTrainingDataset
-from ..utils import TransformerValidationDataset
-from ..utils import TransformerTestingDataset
-from ..utils.misc import ProgressBar
-from ..utils.misc import get_metrics_multitask, print_metrics_multitask
-from ..utils.misc import convert_args_kwargs_to_kwargs
-
-
-def _manage_ctx_fit(func):
-    ''' ... '''
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        # format arguments
-        kwargs = convert_args_kwargs_to_kwargs(func, args, kwargs)
-
-        if kwargs['self']._device_ids is None:
-            return func(**kwargs)
-        else:
-            # change primary device
-            default_device = kwargs['self'].device
-            kwargs['self'].device = kwargs['self']._device_ids[0]
-            rtn = func(**kwargs)
-            kwargs['self'].to(default_device)
-            return rtn
-    return wrapper
-
-
-class Transformer(BaseEstimator):
-    ''' ... '''
-    def __init__(self,
-        src_modalities: dict[str, dict[str, Any]],
-        tgt_modalities: dict[str, dict[str, Any]],
-        d_model: int = 32,
-        nhead: int = 1,
-        num_layers: int = 1,
-        num_epochs: int = 32,
-        batch_size: int = 8,
-        batch_size_multiplier: int = 1,
-        lr: float = 1e-2,
-        weight_decay: float = 0.0,
-        beta: float = 0.9999,
-        gamma: float = 2.0,
-        scale: float = 1.0,
-        lambd: float = 0.0,
-        criterion: str | None = None,
-        device: str = 'cpu',
-        verbose: int = 0,
-        _device_ids: list | None = None,
-        _dataloader_num_workers: int = 0,
-        _amp_enabled: bool = False,
-    ) -> None:  
-        ''' ... '''
-        # for multiprocessing
-        self._rank = 0
-        self._lock = None
-
-        # positional parameters
-        self.src_modalities = src_modalities
-        self.tgt_modalities = tgt_modalities
-
-        # training parameters
-        self.d_model = d_model
-        self.nhead = nhead
-        self.num_layers = num_layers
-        self.num_epochs = num_epochs
-        self.batch_size = batch_size
-        self.batch_size_multiplier = batch_size_multiplier
-        self.lr = lr
-        self.weight_decay = weight_decay
-        self.beta = beta
-        self.gamma = gamma
-        self.scale = scale
-        self.lambd = lambd
-        self.criterion = criterion
-        self.device = device
-        self.verbose = verbose
-        self._device_ids = _device_ids
-        self._dataloader_num_workers = _dataloader_num_workers
-        self._amp_enabled = _amp_enabled
-
-    @_manage_ctx_fit
-    def fit(self, 
-        x, y, 
-        is_embedding: dict[str, bool] | None = None,
-    ) -> Self:
-        ''' ... '''
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # initialize neural network
-        self.net_ = self._init_net()
-
-        # initialize dataloaders
-        ldr_trn, ldr_vld = self._init_dataloader(x, y, is_embedding)
-
-        # initialize optimizer and scheduler
-        optimizer = self._init_optimizer()
-        scheduler = self._init_scheduler(optimizer)
-
-        # gradient scaler for AMP
-        if self._amp_enabled: scaler = torch.cuda.amp.GradScaler()
-        
-        # initialize loss function (binary cross entropy)
-        loss_func = self._init_loss_func({
-            k: (
-                sum([_[k] == 0 for _ in ldr_trn.dataset.tgt]),
-                sum([_[k] == 1 for _ in ldr_trn.dataset.tgt]),
-            ) for k in self.tgt_modalities
-        })
-
-        # to record the best validation performance criterion
-        if self.criterion is not None: best_crit = None
-
-        # progress bar for epoch loops
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch = tqdm.tqdm(
-                    desc = 'Rank {:02d}'.format(self._rank),
-                    total = self.num_epochs,
-                    position = self._rank,
-                    ascii = True,
-                    leave = False,
-                    bar_format='{l_bar}{r_bar}'
-                )
-
-        # training loop
-        for epoch in range(self.num_epochs):
-            # progress bar for batch loops
-            if self.verbose > 1: 
-                pbr_batch = ProgressBar(len(ldr_trn.dataset), 'Epoch {:03d} (TRN)'.format(epoch))
-
-            # set model to train mode
-            torch.set_grad_enabled(True)
-            self.net_.train()
-
-            scores_trn: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            y_true_trn: dict[str, list[int]]   = {k: [] for k in self.tgt_modalities}
-            losses_trn: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            for n_iter, (x_batch, y_batch, mask_x, mask_y) in enumerate(ldr_trn):
-                # mount data to the proper device
-                x_batch = {k: x_batch[k].to(self.device) for k in self.src_modalities}
-                y_batch = {k: y_batch[k].to(torch.float).to(self.device) for k in self.tgt_modalities}
-                mask_x = {k: mask_x[k].to(self.device) for k in self.src_modalities}
-                mask_y = {k: mask_y[k].to(self.device) for k in self.tgt_modalities}
-
-                # forward
-                with torch.autocast(
-                    device_type = 'cpu' if self.device == 'cpu' else 'cuda',
-                    dtype = torch.bfloat16 if self.device == 'cpu' else torch.float16,
-                    enabled = self._amp_enabled,
-                ):
-                    outputs = self.net_(x_batch, mask_x, is_embedding)
-
-                    # calculate multitask loss
-                    loss = 0
-                    for i, tgt_k in enumerate(self.tgt_modalities):
-                        loss_k = loss_func[tgt_k](outputs[tgt_k], y_batch[tgt_k])
-                        loss_k = torch.masked_select(loss_k, torch.logical_not(mask_y[tgt_k].squeeze()))
-                        loss += loss_k.mean()
-                        losses_trn[tgt_k] += loss_k.detach().cpu().numpy().tolist()
-
-                        # if self.lambd != 0:
-                
-                # backward
-                if self._amp_enabled:
-                    scaler.scale(loss).backward()
-                else:
-                    loss.backward()
-                
-                # update parameters
-                if n_iter != 0 and n_iter % self.batch_size_multiplier == 0:
-                    if self._amp_enabled:
-                        scaler.step(optimizer)
-                        scaler.update()
-                        optimizer.zero_grad()
-                    else: 
-                        optimizer.step()
-                        optimizer.zero_grad()
-
-                # save outputs to evaluate performance later
-                for tgt_k in self.tgt_modalities:
-                    tmp = torch.masked_select(outputs[tgt_k], torch.logical_not(mask_y[tgt_k].squeeze()))
-                    scores_trn[tgt_k] += tmp.detach().cpu().numpy().tolist()
-                    tmp = torch.masked_select(y_batch[tgt_k], torch.logical_not(mask_y[tgt_k].squeeze()))
-                    y_true_trn[tgt_k] += tmp.cpu().numpy().tolist()
-
-                # update progress bar
-                if self.verbose > 1:
-                    batch_size = len(next(iter(x_batch.values())))
-                    pbr_batch.update(batch_size, {})
-                    pbr_batch.refresh()
-
-            # for better tqdm progress bar display
-            if self.verbose > 1:
-                pbr_batch.close()
-
-            # set scheduler
-            scheduler.step()
-
-            # calculate and print training performance metrics
-            y_pred_trn: dict[str, list[int]]   = {k: [] for k in self.tgt_modalities}
-            y_prob_trn: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            for tgt_k in self.tgt_modalities:
-                for i in range(len(scores_trn[tgt_k])):
-                    y_pred_trn[tgt_k].append(1 if scores_trn[tgt_k][i] > 0 else 0)
-                    y_prob_trn[tgt_k].append(expit(scores_trn[tgt_k][i]))
-            met_trn = get_metrics_multitask(y_true_trn, y_pred_trn, y_prob_trn)
-
-            # add loss to metrics
-            for tgt_k in self.tgt_modalities:
-                met_trn[tgt_k]['Loss'] = np.mean(losses_trn[tgt_k])
-
-            if self.verbose > 2:
-                print_metrics_multitask(met_trn)
-
-            # progress bar for validation
-            if self.verbose > 1:
-                pbr_batch = ProgressBar(len(ldr_vld.dataset), 'Epoch {:03d} (VLD)'.format(epoch))
-
-            # set model to validation mode
-            torch.set_grad_enabled(False)
-            self.net_.eval()
-
-            scores_vld: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            y_true_vld: dict[str, list[int]]   = {k: [] for k in self.tgt_modalities}
-            losses_vld: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            for x_batch, y_batch, mask_x, mask_y in ldr_vld:
-                # mount data to the proper device
-                x_batch = {k: x_batch[k].to(self.device) for k in self.src_modalities}
-                y_batch = {k: y_batch[k].to(torch.float).to(self.device) for k in self.tgt_modalities}
-                mask_x = {k: mask_x[k].to(self.device) for k in self.src_modalities}
-                mask_y = {k: mask_y[k].to(self.device) for k in self.tgt_modalities}
-
-                # forward
-                with torch.autocast(
-                    device_type = 'cpu' if self.device == 'cpu' else 'cuda',
-                    dtype = torch.bfloat16 if self.device == 'cpu' else torch.float16,
-                    enabled = self._amp_enabled
-                ):
-                    outputs = self.net_(x_batch, mask_x, is_embedding)
-
-                    # calculate multitask loss
-                    for i, tgt_k in enumerate(self.tgt_modalities):
-                        loss_k = loss_func[tgt_k](outputs[tgt_k], y_batch[tgt_k])
-                        loss_k = torch.masked_select(loss_k, torch.logical_not(mask_y[tgt_k].squeeze()))
-                        losses_vld[tgt_k] += loss_k.detach().cpu().numpy().tolist()
-
-                # save outputs to evaluate performance later
-                for tgt_k in self.tgt_modalities:
-                    tmp = torch.masked_select(outputs[tgt_k], torch.logical_not(mask_y[tgt_k].squeeze()))
-                    scores_vld[tgt_k] += tmp.detach().cpu().numpy().tolist()
-                    tmp = torch.masked_select(y_batch[tgt_k], torch.logical_not(mask_y[tgt_k].squeeze()))
-                    y_true_vld[tgt_k] += tmp.cpu().numpy().tolist()
-
-                # update progress bar
-                if self.verbose > 1:
-                    batch_size = len(next(iter(x_batch.values())))
-                    pbr_batch.update(batch_size, {})
-                    pbr_batch.refresh()
-
-            # for better tqdm progress bar display
-            if self.verbose > 1:
-                pbr_batch.close()
-
-            # calculate and print validation performance metrics
-            y_pred_vld: dict[str, list[int]]   = {k: [] for k in self.tgt_modalities}
-            y_prob_vld: dict[str, list[float]] = {k: [] for k in self.tgt_modalities}
-            for tgt_k in self.tgt_modalities:
-                for i in range(len(scores_vld[tgt_k])):
-                    y_pred_vld[tgt_k].append(1 if scores_vld[tgt_k][i] > 0 else 0)
-                    y_prob_vld[tgt_k].append(expit(scores_vld[tgt_k][i]))
-            met_vld = get_metrics_multitask(y_true_vld, y_pred_vld, y_prob_vld)
-
-            # add loss to metrics
-            for tgt_k in self.tgt_modalities:
-                met_vld[tgt_k]['Loss'] = np.mean(losses_vld[tgt_k])
-
-            if self.verbose > 2:
-                print_metrics_multitask(met_vld)
-
-            # save the model if it has the best validation performance criterion by far
-            if self.criterion is None: continue
-            
-            # is current criterion better than previous best?
-            curr_crit = np.mean([met_vld[k][self.criterion] for k in self.tgt_modalities])
-            if best_crit is None or np.isnan(best_crit):
-                is_better = True
-            elif self.criterion == 'Loss' and best_crit >= curr_crit:
-                is_better = True
-            elif self.criterion != 'Loss' and best_crit <= curr_crit:
-                is_better = True
-            else:
-                is_better = False
-
-            # update best criterion
-            if is_better:
-                best_crit = curr_crit
-                best_state_dict = deepcopy(self.net_.state_dict())
-
-            if self.verbose > 2:
-                print('Best {}: {}'.format(self.criterion, best_crit))
-
-            if self.verbose == 1:
-                with self._lock if self._lock is not None else suppress():
-                    pbr_epoch.update(1)
-                    pbr_epoch.refresh()
-
-        if self.verbose == 1:
-            with self._lock if self._lock is not None else suppress():
-                pbr_epoch.close()
-
-        # restore the model of the best validation performance across all epoches
-        if ldr_vld is not None and self.criterion is not None:
-            self.net_.load_state_dict(best_state_dict)
-
-        return self
-
-    def predict_logits(self,
-        x: list[dict[str, Any]],
-        is_embedding: dict[str, bool] | None = None,
-        _batch_size: int | None = None,
-    ) -> list[dict[str, float]]:
-        '''
-        The input x can be a single sample or a list of samples.
-        '''
-        # input validation
-        check_is_fitted(self)
-        
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # set model to eval mode
-        torch.set_grad_enabled(False)
-        self.net_.eval()
-
-        # intialize dataset and dataloader object
-        dat = TransformerTestingDataset(x, self.src_modalities, is_embedding)
-        ldr = DataLoader(
-            dataset = dat,
-            batch_size = _batch_size if _batch_size is not None else len(x),
-            shuffle = False,
-            drop_last = False,
-            num_workers = 0,
-            collate_fn = TransformerTestingDataset.collate_fn,
-        )
-
-        # run model and collect results
-        logits: list[dict[str, float]] = []
-        for x_batch, mask_x in ldr:
-            # mount data to the proper device
-            x_batch = {k: x_batch[k].to(self.device) for k in self.src_modalities}
-            mask_x = {k: mask_x[k].to(self.device) for k in self.src_modalities}
-
-            # forward
-            output: dict[str, Tensor] = self.net_(x_batch, mask_x, is_embedding)
-            
-            # convert output from dict-of-list to list of dict, then append
-            tmp = {k: output[k].tolist() for k in self.tgt_modalities}
-            tmp = [{k: tmp[k][i] for k in self.tgt_modalities} for i in range(len(next(iter(tmp.values()))))]
-            logits += tmp
-
-        return logits
-        
-    def predict_proba(self,
-        x: list[dict[str, Any]],
-        is_embedding: dict[str, bool] | None = None,
-        temperature: float = 1.0,
-        _batch_size: int | None = None,
-    ) -> list[dict[str, float]]:
-        ''' ... '''
-        logits = self.predict_logits(x, is_embedding, _batch_size)
-        return [{k: expit(smp[k] / temperature) for k in self.tgt_modalities} for smp in logits]
-
-    def predict(self,
-        x: list[dict[str, Any]],
-        is_embedding: dict[str, bool] | None = None,
-        _batch_size: int | None = None,
-    ) -> list[dict[str, int]]:
-        ''' ... '''
-        logits = self.predict_logits(x, is_embedding, _batch_size)
-        return [{k: int(smp[k] > 0.0) for k in self.tgt_modalities} for smp in logits]
-
-    def save(self, filepath: str) -> None:
-        ''' ... '''
-        check_is_fitted(self)
-        state_dict = self.net_.state_dict()
-
-        # attach model hyper parameters
-        state_dict['src_modalities'] = self.src_modalities
-        state_dict['tgt_modalities'] = self.tgt_modalities
-        state_dict['d_model'] = self.d_model
-        state_dict['nhead'] = self.nhead
-        state_dict['num_layers'] = self.num_layers
-        torch.save(state_dict, filepath)
-
-    def load(self, filepath: str) -> None:
-        ''' ... '''
-        # load state_dict
-        state_dict = torch.load(filepath, map_location='cpu')
-
-        # load essential parameters
-        self.src_modalities: dict[str, dict[str, Any]] = state_dict.pop('src_modalities')
-        self.tgt_modalities: dict[str, dict[str, Any]] = state_dict.pop('tgt_modalities')
-        self.d_model = state_dict.pop('d_model')
-        self.nhead = state_dict.pop('nhead')
-        self.num_layers = state_dict.pop('num_layers')
-
-        # initialize model
-        self.net_ = nn.Transformer(
-            self.src_modalities,
-            self.tgt_modalities,
-            self.d_model,
-            self.nhead,
-            self.num_layers,
-        )
-
-        # load model parameters
-        self.net_.load_state_dict(state_dict)
-        self.to(self.device)
-
-    def to(self, device: str) -> Self:
-        ''' Mount model to the given device. '''
-        self.device = device
-        if hasattr(self, 'net_'): self.net_ = self.net_.to(device)
-        return self
-    
-    @classmethod
-    def from_ckpt(cls, filepath: str) -> Self:
-        ''' ... '''
-        obj = cls(None, None)
-        obj.load(filepath)
-        return obj
-    
-    def _init_net(self):
-        """ ... """
-        net = nn.Transformer(
-            self.src_modalities,
-            self.tgt_modalities,
-            self.d_model,
-            self.nhead,
-            self.num_layers,
-        ).to(self.device)
-
-        # train on multiple GPUs using torch.nn.DataParallel
-        if self._device_ids is not None:
-            net = torch.nn.DataParallel(net, device_ids=self._device_ids)
-
-        # intialize model parameters using xavier_uniform
-        for p in net.parameters():
-            if p.dim() > 1:
-                torch.nn.init.xavier_uniform_(p)
-        
-        return net
-    
-    def _init_dataloader(self, x, y, is_embedding):
-        """ ... """
-        # split dataset
-        x_trn, x_vld, y_trn, y_vld = train_test_split(
-            x, y, test_size = 0.2, random_state = 0,
-        )
-
-        # initialize dataset and dataloader
-        # dat_trn = TransformerTrainingDataset(
-        # dat_trn = TransformerBalancedTrainingDataset(
-        dat_trn = Transformer2ndOrderBalancedTrainingDataset(
-            x_trn, y_trn,
-            self.src_modalities,
-            self.tgt_modalities,
-            dropout_rate = .5,
-            # dropout_strategy = 'compensated',
-            dropout_strategy = 'permutation',
-        )
-
-        dat_vld = TransformerValidationDataset(
-            x_vld, y_vld,
-            self.src_modalities,
-            self.tgt_modalities,
-            is_embedding,
-        )
-
-        ldr_trn = DataLoader(
-            dataset = dat_trn,
-            batch_size = self.batch_size,
-            shuffle = True,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = TransformerTrainingDataset.collate_fn,
-            # pin_memory = True
-        )
-
-        ldr_vld = DataLoader(
-            dataset = dat_vld,
-            batch_size = self.batch_size,
-            shuffle = False,
-            drop_last = False,
-            num_workers = self._dataloader_num_workers,
-            collate_fn = TransformerValidationDataset.collate_fn,
-            # pin_memory = True
-        )
-
-        return ldr_trn, ldr_vld
-    
-    def _init_optimizer(self):
-        """ ... """
-        return torch.optim.AdamW(
-            self.net_.parameters(),
-            lr = self.lr,
-            betas = (0.9, 0.98),
-            weight_decay = self.weight_decay
-        )
-    
-    def _init_scheduler(self, optimizer):
-        """ ... """
-        return torch.optim.lr_scheduler.OneCycleLR(
-            optimizer = optimizer, 
-            max_lr = self.lr,
-            total_steps = self.num_epochs,
-            verbose = (self.verbose > 2)
-        )
-    
-    def _init_loss_func(self, 
-        num_per_cls: dict[str, tuple[int, int]],
-    ) -> dict[str, Module]:
-        """ ... """
-        return {k: nn.SigmoidFocalLoss(
-            beta = self.beta,
-            gamma = self.gamma,
-            scale = self.scale,
-            num_per_cls = num_per_cls[k],
-            reduction = 'none',
-        ) for k in self.tgt_modalities}
-    
-    def _extract_embedding(self,
-        x: list[dict[str, Any]],
-        is_embedding: dict[str, bool] | None = None,
-        _batch_size: int | None = None,
-    ) -> list[dict[str, Any]]:
-        """ ... """
-        # input validation
-        check_is_fitted(self)
-        
-        # for PyTorch computational efficiency
-        torch.set_num_threads(1)
-
-        # set model to eval mode
-        torch.set_grad_enabled(False)
-        self.net_.eval()
-
-        # intialize dataset and dataloader object
-        dat = TransformerTestingDataset(x, self.src_modalities, is_embedding)
-        ldr = DataLoader(
-            dataset = dat,
-            batch_size = _batch_size if _batch_size is not None else len(x),
-            shuffle = False,
-            drop_last = False,
-            num_workers = 0,
-            collate_fn = TransformerTestingDataset.collate_fn,
-        )
-
-        # run model and extract embeddings
-        embeddings: list[dict[str, Any]] = []
-        for x_batch, _ in ldr:
-            # mount data to the proper device
-            x_batch = {k: x_batch[k].to(self.device) for k in self.src_modalities}
-
-            # forward
-            out: dict[str, Tensor] = self.net_.forward_emb(x_batch, is_embedding)
-            
-            # convert output from dict-of-list to list of dict, then append
-            tmp = {k: out[k].detach().cpu().numpy() for k in self.src_modalities}
-            tmp = [{k: tmp[k][i] for k in self.src_modalities} for i in range(len(next(iter(tmp.values()))))]
-            embeddings += tmp
-
-        # remove imputed embeddings
-        for i in range(len(x)):
-            avail = [k for k, v in x[i].items() if v is not None]
-            embeddings[i] = {k: embeddings[i][k] for k in avail}
-
-        return embeddings
-

adrd/nn/__init__.py

@@ -1,12 +0,0 @@
-from .transformer import Transformer
-from .vitautoenc import ViTAutoEnc
-from .unet import UNet3D
-from .unet_3d import UNet3DBase
-from .focal_loss import SigmoidFocalLoss
-from .unet_img_model import ImageModel
-from .img_model_wrapper import ImagingModelWrapper
-from .resnet_img_model import ResNetModel
-from .c3d import C3D
-from .dense_net import DenseNet
-from .cnn_resnet3d import CNNResNet3D
-from .cnn_resnet3d_with_linear_classifier import CNNResNet3DWithLinearClassifier

adrd/nn/blocks.py

@@ -1,57 +0,0 @@
-# Copyright (c) MONAI Consortium
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#     http://www.apache.org/licenses/LICENSE-2.0
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from monai.networks.blocks.mlp import MLPBlock
-from typing import Sequence, Union
-import torch
-import torch.nn as nn
-
-from ..nn.selfattention import SABlock
-
-class TransformerBlock(nn.Module):
-    """
-    A transformer block, based on: "Dosovitskiy et al.,
-    An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale <https://arxiv.org/abs/2010.11929>"
-    """
-
-    def __init__(
-        self, hidden_size: int, mlp_dim: int, num_heads: int, dropout_rate: float = 0.0, qkv_bias: bool = False
-    ) -> None:
-        """
-        Args:
-            hidden_size: dimension of hidden layer.
-            mlp_dim: dimension of feedforward layer.
-            num_heads: number of attention heads.
-            dropout_rate: faction of the input units to drop.
-            qkv_bias: apply bias term for the qkv linear layer
-
-        """
-
-        super().__init__()
-
-        if not (0 <= dropout_rate <= 1):
-            raise ValueError("dropout_rate should be between 0 and 1.")
-
-        if hidden_size % num_heads != 0:
-            raise ValueError("hidden_size should be divisible by num_heads.")
-
-        self.mlp = MLPBlock(hidden_size, mlp_dim, dropout_rate)
-        self.norm1 = nn.LayerNorm(hidden_size)
-        self.attn = SABlock(hidden_size, num_heads, dropout_rate, qkv_bias)
-        self.norm2 = nn.LayerNorm(hidden_size)
-
-    def forward(self, x, return_attention=False):
-        y, attn = self.attn(self.norm1(x))
-        if return_attention:
-            return attn
-        x = x + y
-        x = x + self.mlp(self.norm2(x))
-        return x

adrd/nn/c3d.py

@@ -1,99 +0,0 @@
-# From https://github.com/xmuyzz/3D-CNN-PyTorch/blob/master/models/C3DNet.py
-
-import torch
-import torch.nn as nn
-import sys
-# from icecream import ic
-import math
-
-class C3D(torch.nn.Module):
-    
-    def __init__(self, tgt_modalities, in_channels=1, load_from_ckpt=None):
-        
-        super(C3D, self).__init__()
-        self.conv_group1 = nn.Sequential(
-            nn.Conv3d(in_channels, 64, kernel_size=3, padding=1),
-            nn.BatchNorm3d(64),
-            nn.ReLU(),
-            nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(1, 2, 2)))
-        self.conv_group2 = nn.Sequential(
-            nn.Conv3d(64, 128, kernel_size=3, padding=1),
-            nn.BatchNorm3d(128),
-            nn.ReLU(),
-            nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2)))
-        self.conv_group3 = nn.Sequential(
-            nn.Conv3d(128, 256, kernel_size=3, padding=1),
-            nn.BatchNorm3d(256),
-            nn.ReLU(),
-            nn.Conv3d(256, 256, kernel_size=3, padding=1),
-            nn.BatchNorm3d(256),
-            nn.ReLU(),
-            nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
-            )
-        self.conv_group4 = nn.Sequential(
-            nn.Conv3d(256, 512, kernel_size=3, padding=1),
-            nn.BatchNorm3d(512),
-            nn.ReLU(),
-            nn.Conv3d(512, 512, kernel_size=3, padding=1),
-            nn.BatchNorm3d(512),
-            nn.ReLU(),
-            nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2), padding=(0, 1, 1))
-            )
-
-        # last_duration = int(math.floor(128 / 16))
-        # last_size = int(math.ceil(128 / 32))
-        self.fc1 = nn.Sequential(
-            nn.Linear((512 * 15 * 9 * 9) , 512),
-            nn.ReLU(),
-            nn.Dropout(0.5))
-        self.fc2 = nn.Sequential(
-            nn.Linear(512, 256),
-            nn.ReLU(),
-            nn.Dropout(0.5))
-        # self.fc = nn.Sequential(
-        #     nn.Linear(4096, num_classes))
-        
-        self.fc = torch.nn.ModuleDict()
-        for k in tgt_modalities:
-            self.fc[k] = torch.nn.Linear(256, 1)
-            
-    def forward(self, x):
-        # for k in x.keys():
-        #     x[k] = x[k].to(torch.float32)
-        
-        # x = torch.stack([o for o in x.values()], dim=0)[0]
-        # print(x.shape)
-        
-        out = self.conv_group1(x)
-        out = self.conv_group2(out)
-        out = self.conv_group3(out)
-        out = self.conv_group4(out)
-        out = out.view(out.size(0), -1)
-        # print(out.shape)
-        out = self.fc1(out)
-        out = self.fc2(out)
-        # out = self.fc(out)
-        
-        tgt_iter = self.fc.keys()
-        out_tgt = {k: self.fc[k](out).squeeze(1) for k in tgt_iter}
-        return out_tgt
-    
-
-if __name__ == "__main__":
-    model = C3D(tgt_modalities=['NC', 'MCI', 'DE'])
-    print(model)
-    x = torch.rand((1, 1, 128, 128, 128))
-    # layers = list(model.features.named_children())
-    # features = nn.Sequential(*list(model.features.children()))(x)
-    # print(features.shape)
-    print(sum(p.numel() for p in model.parameters()))
-    # layer_found = False
-    # features = None
-    # desired_layer_name = 'transition3'
-
-    # for name, layer in layers:
-    #     if name == desired_layer_name:
-    #         x = layer(x)
-    #         print(x)
-    # model(x)
-    # print(features)

adrd/nn/cnn_resnet3d.py

@@ -1,81 +0,0 @@
-import torch
-import torch.nn as nn
-from typing import Any, Type
-Tensor = Type[torch.Tensor]
-
-from .resnet3d import r3d_18
-
-
-class CNNResNet3D(nn.Module):
-
-    def __init__(self,
-        src_modalities: dict[str, dict[str, Any]],
-        tgt_modalities: dict[str, dict[str, Any]]
-    ) -> None:
-        """ ... """
-        super().__init__()
-
-        # resnet
-        # embedding modules for source
-        self.modules_emb_src = nn.ModuleDict()
-        for k, info in src_modalities.items():
-            if info['type'] == 'imaging' and len(info['img_shape']) == 4:
-                self.modules_emb_src[k] = nn.Sequential(
-                    r3d_18(),
-                    nn.Dropout(0.5)
-                )
-            else:
-                # unrecognized
-                raise ValueError('{} is an unrecognized data modality'.format(k))
-
-        # classifiers (binary only)
-        self.modules_cls = nn.ModuleDict()
-        for k, info in tgt_modalities.items():
-            if info['type'] == 'categorical' and info['num_categories'] == 2:
-                # categorical
-                self.modules_cls[k] = nn.Linear(256, 1)
-            else:
-                # unrecognized
-                raise ValueError
-            
-    def forward(self,
-        x: dict[str, Tensor],
-    ) -> dict[str, Tensor]:
-        """ ... """
-        out_emb = self.forward_emb(x)
-        out_emb = out_emb[list(out_emb.keys())[0]]
-        out_cls = self.forward_cls(out_emb)
-        return out_cls
-
-    def forward_emb(self,
-        x: dict[str, Tensor],
-    ) -> dict[str, Tensor]:
-        """ ... """
-        out_emb = dict()
-        for k in self.modules_emb_src.keys():
-            out_emb[k] = self.modules_emb_src[k](x[k])
-        return out_emb
-    
-    def forward_cls(self,
-        out_emb: dict[str, Tensor]
-    ) -> dict[str, Tensor]:
-        """ ... """
-        out_cls = dict()
-        for k in self.modules_cls.keys():
-            out_cls[k] = self.modules_cls[k](out_emb).squeeze(1)
-        return out_cls
-    
-
-# for testing purpose only
-if __name__ == '__main__':
-    src_modalities = {
-        'img_MRI_T1': {'type': 'imaging', 'img_shape': [1, 182, 218, 182]}
-    }
-    tgt_modalities = {
-        'AD': {'type': 'categorical', 'num_categories': 2},
-        'PD': {'type': 'categorical', 'num_categories': 2}
-    }
-    net = CNNResNet3D(src_modalities, tgt_modalities)
-    net.eval()
-    x = {'img_MRI_T1': torch.zeros(2, 1, 182, 218, 182)}
-    print(net(x))

adrd/nn/cnn_resnet3d_with_linear_classifier.py

@@ -1,56 +0,0 @@
-import torch
-import torch.nn as nn
-from typing import Any, Type
-Tensor = Type[torch.Tensor]
-
-from .resnet3d import r3d_18
-
-class CNNResNet3DWithLinearClassifier(nn.Module):
-
-    def __init__(self,
-        src_modalities: dict[str, dict[str, Any]],
-        tgt_modalities: dict[str, dict[str, Any]]
-    ) -> None:
-        """ ... """
-        super().__init__()
-        self.core = _CNNResNet3DWithLinearClassifier(len(tgt_modalities))
-        self.src_modalities = src_modalities
-        self.tgt_modalities = tgt_modalities
-
-    def forward(self,
-        x: dict[str, Tensor],
-    ) -> dict[str, Tensor]:
-        """ x is expected to be a singleton dictionary """
-        src_k = list(x.keys())[0]
-        x = x[src_k]
-        out = self.core(x)
-        out = {tgt_k: out[:, i] for i, tgt_k in enumerate(self.tgt_modalities)}
-        return out
-
-
-class _CNNResNet3DWithLinearClassifier(nn.Module):
-
-    def __init__(self,
-        len_tgt_modalities: int,
-    ) -> None:
-        """ ... """
-        super().__init__()
-        self.cnn = r3d_18()
-        self.cls = nn.Sequential(
-            nn.Dropout(0.5),
-            nn.Linear(256, len_tgt_modalities),
-        )
-
-    def forward(self, x: Tensor) -> Tensor:
-        """ ... """
-        out_emb = self.forward_emb(x)
-        out_cls = self.forward_cls(out_emb)
-        return out_cls
-    
-    def forward_emb(self, x: Tensor) -> Tensor:
-        """ ... """
-        return self.cnn(x)
-    
-    def forward_cls(self, out_emb: Tensor) -> Tensor:
-        """ ... """
-        return self.cls(out_emb)

adrd/nn/dense_net.py

@@ -1,211 +0,0 @@
-# This implementation is based on the DenseNet-BC implementation in torchvision
-# https://github.com/pytorch/vision/blob/master/torchvision/models/densenet.py
-# https://github.com/gpleiss/efficient_densenet_pytorch/blob/master/models/densenet.py
-
-
-import math
-import torch
-import numpy as np
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint as cp
-from collections import OrderedDict
-
-
-def _bn_function_factory(norm, relu, conv):
-    def bn_function(*inputs):
-        concated_features = torch.cat(inputs, 1)
-        bottleneck_output = conv(relu(norm(concated_features)))
-        return bottleneck_output
-
-    return bn_function
-
-
-class _DenseLayer(nn.Module):
-    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, efficient=False):
-        super(_DenseLayer, self).__init__()
-        self.add_module('norm1', nn.BatchNorm3d(num_input_features)),
-        self.add_module('relu1', nn.ReLU(inplace=True)),
-        self.add_module('conv1', nn.Conv3d(num_input_features, bn_size * growth_rate, kernel_size=1, stride=1, bias=False)),
-        self.add_module('norm2', nn.BatchNorm3d(bn_size * growth_rate)),
-        self.add_module('relu2', nn.ReLU(inplace=True)),
-        self.add_module('conv2', nn.Conv3d(bn_size * growth_rate, growth_rate, kernel_size=3, stride=1, padding=1, bias=False)),
-        self.drop_rate = drop_rate
-        self.efficient = efficient
-
-    def forward(self, *prev_features):
-        bn_function = _bn_function_factory(self.norm1, self.relu1, self.conv1)
-        if self.efficient and any(prev_feature.requires_grad for prev_feature in prev_features):
-            bottleneck_output = cp.checkpoint(bn_function, *prev_features)
-        else:
-            bottleneck_output = bn_function(*prev_features)
-        new_features = self.conv2(self.relu2(self.norm2(bottleneck_output)))
-        if self.drop_rate > 0:
-            new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)
-        return new_features
-
-
-class _Transition(nn.Sequential):
-    def __init__(self, num_input_features, num_output_features):
-        super(_Transition, self).__init__()
-        self.add_module('norm', nn.BatchNorm3d(num_input_features))
-        self.add_module('relu', nn.ReLU(inplace=True))
-        self.add_module('conv', nn.Conv3d(num_input_features, num_output_features, kernel_size=1, stride=1, bias=False))
-        self.add_module('pool', nn.AvgPool3d(kernel_size=2, stride=2))
-
-
-class _DenseBlock(nn.Module):
-    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, efficient=False):
-        super(_DenseBlock, self).__init__()
-        for i in range(num_layers):
-            layer = _DenseLayer(
-                num_input_features + i * growth_rate,
-                growth_rate=growth_rate,
-                bn_size=bn_size,
-                drop_rate=drop_rate,
-                efficient=efficient,
-            )
-            self.add_module('denselayer%d' % (i + 1), layer)
-
-    def forward(self, init_features):
-        features = [init_features]
-        for name, layer in self.named_children():
-            new_features = layer(*features)
-            features.append(new_features)
-        return torch.cat(features, 1)
-
-
-class DenseNet(nn.Module):
-    r"""Densenet-BC model class, based on
-    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
-    Args:
-        growth_rate (int) - how many filters to add each layer (`k` in paper)
-        block_config (list of 3 or 4 ints) - how many layers in each pooling block
-        num_init_features (int) - the number of filters to learn in the first convolution layer
-        bn_size (int) - multiplicative factor for number of bottle neck layers
-            (i.e. bn_size * k features in the bottleneck layer)
-        drop_rate (float) - dropout rate after each dense layer
-        tgt_modalities (list) - list of target modalities
-        efficient (bool) - set to True to use checkpointing. Much more memory efficient, but slower.
-    """
-    # def __init__(self, tgt_modalities, growth_rate=12, block_config=(3, 3, 3), compression=0.5,
-    #              num_init_features=16, bn_size=4, drop_rate=0, efficient=False, load_from_ckpt=False): # config 1
-    
-    def __init__(self, tgt_modalities, growth_rate=12, block_config=(3, 3, 3), compression=0.5,
-                 num_init_features=16, bn_size=4, drop_rate=0, efficient=False, load_from_ckpt=False): # config 2
-        
-        super(DenseNet, self).__init__()
-
-        # First convolution
-        self.features = nn.Sequential(OrderedDict([('conv0', nn.Conv3d(1, num_init_features, kernel_size=7, stride=2, padding=0, bias=False)),]))
-        self.features.add_module('norm0', nn.BatchNorm3d(num_init_features))
-        self.features.add_module('relu0', nn.ReLU(inplace=True))
-        self.features.add_module('pool0', nn.MaxPool3d(kernel_size=3, stride=2, padding=0, ceil_mode=False))
-        self.tgt_modalities = tgt_modalities
-
-        # Each denseblock
-        num_features = num_init_features
-        for i, num_layers in enumerate(block_config):
-            block = _DenseBlock(
-                num_layers=num_layers,
-                num_input_features=num_features,
-                bn_size=bn_size,
-                growth_rate=growth_rate,
-                drop_rate=drop_rate,
-                efficient=efficient,
-            )
-            self.features.add_module('denseblock%d' % (i + 1), block)
-            num_features = num_features + num_layers * growth_rate
-            if i != len(block_config):
-                trans = _Transition(num_input_features=num_features,
-                                    num_output_features=int(num_features * compression))
-                self.features.add_module('transition%d' % (i + 1), trans)
-                num_features = int(num_features * compression)
-
-        # Final batch norm
-        self.features.add_module('norm_final', nn.BatchNorm3d(num_features))
-        
-        # Classification heads
-        self.tgt = torch.nn.ModuleDict()
-        for k in tgt_modalities:
-            # self.tgt[k] = torch.nn.Linear(621, 1) # config 2
-            self.tgt[k] = torch.nn.Sequential(
-                    torch.nn.Linear(self.test_size(), 256),
-                    torch.nn.ReLU(),
-                    torch.nn.Linear(256, 1)
-                )
-
-        print(f'load_from_ckpt: {load_from_ckpt}')
-        # Initialization
-        if not load_from_ckpt:
-            for name, param in self.named_parameters():
-                if 'conv' in name and 'weight' in name:
-                    n = param.size(0) * param.size(2) * param.size(3) * param.size(4)
-                    param.data.normal_().mul_(math.sqrt(2. / n))
-                elif 'norm' in name and 'weight' in name:
-                    param.data.fill_(1)
-                elif 'norm' in name and 'bias' in name:
-                    param.data.fill_(0)
-                elif ('classifier' in name or 'tgt' in name) and 'bias' in name:
-                    param.data.fill_(0)
-
-        # self.size = self.test_size()
-
-    def forward(self, x, shap=True):
-        # print(x.shape)
-        features = self.features(x)
-        # print(features.shape)
-        out = F.relu(features, inplace=True)
-        # out = F.adaptive_avg_pool3d(out, (1, 1, 1))
-        out = torch.flatten(out, 1)
-        
-        # print(out.shape)
-        
-        # out_tgt = self.tgt(out).squeeze(1)
-        # print(out_tgt)
-        # return F.softmax(out_tgt)
-        
-        tgt_iter = self.tgt.keys()
-        out_tgt = {k: self.tgt[k](out).squeeze(1) for k in tgt_iter}
-        if shap:
-            out_tgt = torch.stack(list(out_tgt.values()))
-            return out_tgt.T
-        else: 
-            return out_tgt
-
-    def test_size(self):
-        case = torch.ones((1, 1, 182, 218, 182))
-        output = self.features(case).view(-1).size(0)
-        return output
-
-
-if __name__ == "__main__":
-    model = DenseNet(
-        tgt_modalities=['NC', 'MCI', 'DE'], 
-        growth_rate=12, 
-        block_config=(2, 3, 2), 
-        compression=0.5,
-        num_init_features=16, 
-        drop_rate=0.2)
-    print(model)
-    torch.manual_seed(42)
-    x = torch.rand((1, 1, 182, 218, 182))
-    # layers = list(model.features.named_children())
-    features = nn.Sequential(*list(model.features.children()))(x)
-    print(features.shape)
-    print(sum(p.numel() for p in model.parameters()))
-    # out = mdl.net_(x, shap=False)
-    # print(out)
-
-    out = model(x, shap=False)
-    print(out)
-    # layer_found = False
-    # features = None
-    # desired_layer_name = 'transition3'
-
-    # for name, layer in layers:
-    #     if name == desired_layer_name:
-    #         x = layer(x)
-    #         print(x)
-    # model(x)
-    # print(features)

adrd/nn/focal_loss.py

@@ -1,120 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import sys
-
-class SigmoidFocalLoss(nn.Module):
-    ''' ... '''
-    def __init__(
-        self,
-        alpha: float = -1,
-        gamma: float = 2.0,
-        reduction: str = 'mean',
-    ):
-        ''' ... '''
-        super().__init__()
-        self.alpha = alpha
-        self.gamma = gamma
-        self.reduction = reduction
-
-    def forward(self, input, target):
-        ''' ... '''
-        p = torch.sigmoid(input)
-        ce_loss = F.binary_cross_entropy_with_logits(input, target, reduction='none')
-        p_t = p * target + (1 - p) * (1 - target)
-        loss = ce_loss * ((1 - p_t) ** self.gamma)
-
-        if self.alpha >= 0:
-            alpha_t = self.alpha * target + (1 - self.alpha) * (1 - target)
-            loss = alpha_t * loss
-
-        if self.reduction == 'mean':
-            loss = loss.mean()
-        elif self.reduction == 'sum':
-            loss = loss.sum()
-
-        return loss
-
-
-class SigmoidFocalLossBeta(nn.Module):
-    ''' ... '''
-    def __init__(
-        self,
-        beta: float = 0.9999,
-        gamma: float = 2.0,
-        num_per_cls = (1, 1),
-        reduction: str = 'mean',
-    ):
-        ''' ... '''
-        super().__init__()
-        eps = sys.float_info.epsilon
-        self.gamma = gamma
-        self.reduction = reduction
-
-        # weights to balance loss
-        self.weight_neg = ((1 - beta) / (1 - beta ** num_per_cls[0] + eps))
-        self.weight_pos = ((1 - beta) / (1 - beta ** num_per_cls[1] + eps))
-        # weight_neg = (1 - beta) / (1 - beta ** num_per_cls[0])
-        # weight_pos = (1 - beta) / (1 - beta ** num_per_cls[1])
-        # self.weight_neg = weight_neg / (weight_neg + weight_pos)
-        # self.weight_pos = weight_pos / (weight_neg + weight_pos)
-
-    def forward(self, input, target):
-        ''' ... '''
-        p = torch.sigmoid(input)
-        p_t = p * target + (1 - p) * (1 - target)
-        ce_loss = F.binary_cross_entropy_with_logits(input, target, reduction='none')
-        loss = ce_loss * ((1 - p_t) ** self.gamma)
-
-        alpha_t = self.weight_pos * target + self.weight_neg * (1 - target)
-        loss = alpha_t * loss
-
-        if self.reduction == 'mean':
-            loss = loss.mean()
-        elif self.reduction == 'sum':
-            loss = loss.sum()
-
-        return loss
-
-class AsymmetricLoss(nn.Module):
-    def __init__(self, gamma_neg=4, gamma_pos=1, alpha=0.5, clip=0.05, eps=1e-8, disable_torch_grad_focal_loss=True):
-        super(AsymmetricLoss, self).__init__()
-        self.alpha = alpha
-        self.gamma_neg = gamma_neg
-        self.gamma_pos = gamma_pos
-        self.clip = clip
-        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
-        self.eps = eps
-
-
-    def forward(self, x, y):
-        """"
-        Parameters
-        ----------
-        x: input logits
-        y: targets (multi-label binarized vector)
-        """
-        # Calculating Probabilities
-        x_sigmoid = torch.sigmoid(x)
-        xs_pos = x_sigmoid
-        xs_neg = 1 - x_sigmoid
-        # Asymmetric Clipping
-        if self.clip is not None and self.clip > 0:
-            xs_neg = (xs_neg + self.clip).clamp(max=1)
-        # Basic CE calculation
-        los_pos = y * torch.log(xs_pos.clamp(min=self.eps))
-        los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))
-        loss = self.alpha*los_pos + (1-self.alpha)*los_neg
-        # Asymmetric Focusing
-        if self.gamma_neg > 0 or self.gamma_pos > 0:
-            if self.disable_torch_grad_focal_loss:
-                torch.set_grad_enabled(False)
-            pt0 = xs_pos * y
-            pt1 = xs_neg * (1 - y)  # pt = p if t > 0 else 1-p
-            pt = pt0 + pt1
-            one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)
-            one_sided_w = torch.pow(1 - pt, one_sided_gamma)
-            if self.disable_torch_grad_focal_loss:
-                torch.set_grad_enabled(True)
-            loss *= one_sided_w
-        return -loss#.sum()

adrd/nn/img_model_wrapper.py

@@ -1,174 +0,0 @@
-import torch
-from .. import nn
-from .. import model
-import numpy as np
-from icecream import ic
-from monai.networks.nets.swin_unetr import SwinUNETR
-from typing import Any
-
-class ImagingModelWrapper(torch.nn.Module):
-    def __init__(
-            self,
-            arch: str = 'ViTAutoEnc',
-            tgt_modalities: dict | None = {},
-            img_size: int | None = 128,
-            patch_size: int | None = 16,
-            ckpt_path: str | None = None,
-            train_backbone: bool = False,
-            out_dim: int = 128,
-            layers: int | None = 1,
-            device: str = 'cpu',
-            fusion_stage: str = 'middle',
-            ):
-        super(ImagingModelWrapper, self).__init__()
-
-        self.arch = arch
-        self.tgt_modalities = tgt_modalities
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.train_backbone = train_backbone
-        self.ckpt_path = ckpt_path
-        self.device = device
-        self.out_dim = out_dim
-        self.layers = layers
-        self.fusion_stage = fusion_stage
-        
-        
-        if "swinunetr" in self.arch.lower():
-            if "emb" not in self.arch.lower():
-                ckpt_path = '/projectnb/ivc-ml/dlteif/pretrained_models/model_swinvit.pt'
-                ckpt = torch.load(ckpt_path, map_location='cpu')
-                self.img_model = SwinUNETR(
-                    in_channels=1,
-                    out_channels=1,
-                    img_size=128,
-                    feature_size=48,
-                    use_checkpoint=True,
-                )
-                ckpt["state_dict"] = {k.replace("swinViT.", "module."): v for k, v in ckpt["state_dict"].items()}
-                ic(ckpt["state_dict"].keys())
-                self.img_model.load_from(ckpt)
-            self.dim = 768
-
-        elif "vit" in self.arch.lower():    
-            if "emb" not in self.arch.lower():
-                # Initialize image model
-                self.img_model = nn.__dict__[self.arch](
-                    in_channels = 1,
-                    img_size = self.img_size,
-                    patch_size = self.patch_size,
-                )
-
-                if self.ckpt_path:
-                    self.img_model.load(self.ckpt_path, map_location=self.device)
-                self.dim = self.img_model.hidden_size
-            else:
-                self.dim = 768
-
-        if "vit" in self.arch.lower() or "swinunetr" in self.arch.lower():    
-            dim = self.dim
-            if self.fusion_stage == 'middle':
-                downsample = torch.nn.ModuleList()
-                # print('Number of layers: ', self.layers)
-                for i in range(self.layers):
-                    if i == self.layers - 1:
-                        dim_out = self.out_dim
-                        # print(layers)
-                        ks = 2
-                        stride = 2
-                    else:
-                        dim_out = dim // 2
-                        ks = 2
-                        stride = 2
-
-                    downsample.append(
-                        torch.nn.Conv1d(in_channels=dim, out_channels=dim_out, kernel_size=ks, stride=stride)
-                    )
-                    
-                    dim = dim_out
-                    
-                    downsample.append(
-                        torch.nn.BatchNorm1d(dim)
-                    )
-                    downsample.append(
-                        torch.nn.ReLU()
-                    )
-                    
-                    
-                self.downsample = torch.nn.Sequential(*downsample)
-            elif self.fusion_stage == 'late':
-                self.downsample = torch.nn.Identity()
-            else:
-                pass
-            
-            # print('Downsample layers: ', self.downsample)
-                
-        elif "densenet" in self.arch.lower():
-            if "emb" not in self.arch.lower():
-                self.img_model = model.ImagingModel.from_ckpt(self.ckpt_path, device=self.device, img_backend=self.arch, load_from_ckpt=True).net_
-            
-            self.downsample = torch.nn.Linear(3900, self.out_dim)
-
-        # randomly initialize weights for downsample block
-        for p in self.downsample.parameters():
-            if p.dim() > 1:
-                torch.nn.init.xavier_uniform_(p)
-            p.requires_grad = True
-            
-        if "emb" not in self.arch.lower():
-            # freeze imaging model parameters
-            if "densenet" in self.arch.lower():
-                for n, p in self.img_model.features.named_parameters():
-                    if not self.train_backbone:
-                        p.requires_grad = False
-                    else:
-                        p.requires_grad = True
-                for n, p in self.img_model.tgt.named_parameters():
-                    p.requires_grad = False
-            else:
-                for n, p in self.img_model.named_parameters():
-                    # print(n, p.requires_grad)
-                    if not self.train_backbone:
-                        p.requires_grad = False
-                    else:
-                        p.requires_grad = True
-    
-    def forward(self, x):
-        # print("--------ImagingModelWrapper forward--------")
-        if "emb" not in self.arch.lower():
-            if "swinunetr" in self.arch.lower():
-                # print(x.size())
-                out = self.img_model(x)
-                # print(out.size())
-                out = self.downsample(out)
-                # print(out.size())
-                out = torch.mean(out, dim=-1)
-                # print(out.size())
-            elif "vit" in self.arch.lower():
-                out = self.img_model(x, return_emb=True)
-                ic(out.size())
-                out = self.downsample(out)
-                out = torch.mean(out, dim=-1)
-            elif "densenet" in self.arch.lower():
-                out = torch.nn.Sequential(*list(self.img_model.features.children()))(x)
-                # print(out.size())
-                out = torch.flatten(out, 1)
-                out = self.downsample(out)
-        else:
-            # print(x.size())
-            if "swinunetr" in self.arch.lower():
-                x = torch.squeeze(x, dim=1)
-                x = x.view(x.size(0),self.dim, -1)
-            # print('x: ', x.size())    
-            out = self.downsample(x)
-            # print('out: ', out.size())
-            if self.fusion_stage == 'middle':
-                if "vit" in self.arch.lower() or "swinunetr" in self.arch.lower():
-                    out = torch.mean(out, dim=-1)
-                else:
-                    out = torch.squeeze(out, dim=1)
-            elif self.fusion_stage == 'late':
-                pass
-
-        return out
-

adrd/nn/net_resnet3d.py

@@ -1,338 +0,0 @@
-"""
-Created on Sat Nov 21 10:49:39 2021
-
-@author: cxue2
-"""
-
-import torch.nn as nn
-
-
-__all__ = ['r3d_18', 'mc3_18', 'r2plus1d_18']
-
-
-class Conv3DSimple(nn.Conv3d):
-    def __init__(self,
-                 in_planes,
-                 out_planes,
-                 midplanes=None,
-                 stride=1,
-                 padding=1):
-
-        super(Conv3DSimple, self).__init__(
-            in_channels=in_planes,
-            out_channels=out_planes,
-            kernel_size=(3, 3, 3),
-            stride=stride,
-            padding=padding,
-            bias=False)
-
-    @staticmethod
-    def get_downsample_stride(stride):
-        return stride, stride, stride
-
-
-class Conv2Plus1D(nn.Sequential):
-
-    def __init__(self,
-                 in_planes,
-                 out_planes,
-                 midplanes,
-                 stride=1,
-                 padding=1):
-        super(Conv2Plus1D, self).__init__(
-            nn.Conv3d(in_planes, midplanes, kernel_size=(1, 3, 3),
-                      stride=(1, stride, stride), padding=(0, padding, padding),
-                      bias=False),
-            nn.BatchNorm3d(midplanes),
-            nn.ReLU(inplace=True),
-            nn.Conv3d(midplanes, out_planes, kernel_size=(3, 1, 1),
-                      stride=(stride, 1, 1), padding=(padding, 0, 0),
-                      bias=False))
-
-    @staticmethod
-    def get_downsample_stride(stride):
-        return stride, stride, stride
-
-
-class Conv3DNoTemporal(nn.Conv3d):
-
-    def __init__(self,
-                 in_planes,
-                 out_planes,
-                 midplanes=None,
-                 stride=1,
-                 padding=1):
-
-        super(Conv3DNoTemporal, self).__init__(
-            in_channels=in_planes,
-            out_channels=out_planes,
-            kernel_size=(1, 3, 3),
-            stride=(1, stride, stride),
-            padding=(0, padding, padding),
-            bias=False)
-
-    @staticmethod
-    def get_downsample_stride(stride):
-        return 1, stride, stride
-
-
-class BasicBlock(nn.Module):
-
-    expansion = 1
-
-    def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None):
-        midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes)
-
-        super(BasicBlock, self).__init__()
-        self.conv1 = nn.Sequential(
-            conv_builder(inplanes, planes, midplanes, stride),
-            nn.BatchNorm3d(planes),
-            nn.ReLU(inplace=True)
-        )
-        self.conv2 = nn.Sequential(
-            conv_builder(planes, planes, midplanes),
-            nn.BatchNorm3d(planes)
-        )
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.conv2(out)
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None):
-
-        super(Bottleneck, self).__init__()
-        midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes)
-
-        # 1x1x1
-        self.conv1 = nn.Sequential(
-            nn.Conv3d(inplanes, planes, kernel_size=1, bias=False),
-            nn.BatchNorm3d(planes),
-            nn.ReLU(inplace=True)
-        )
-        # Second kernel
-        self.conv2 = nn.Sequential(
-            conv_builder(planes, planes, midplanes, stride),
-            nn.BatchNorm3d(planes),
-            nn.ReLU(inplace=True)
-        )
-
-        # 1x1x1
-        self.conv3 = nn.Sequential(
-            nn.Conv3d(planes, planes * self.expansion, kernel_size=1, bias=False),
-            nn.BatchNorm3d(planes * self.expansion)
-        )
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.conv2(out)
-        out = self.conv3(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class BasicStem(nn.Sequential):
-    """The default conv-batchnorm-relu stem
-    """
-    def __init__(self):
-        super(BasicStem, self).__init__(
-            nn.Conv3d(1, 64, kernel_size=(7, 7, 7), stride=(2, 2, 2),
-                      padding=(3, 3, 3), bias=False),
-            nn.BatchNorm3d(64),
-            nn.ReLU(inplace=True))
-
-
-class R2Plus1dStem(nn.Sequential):
-    """R(2+1)D stem is different than the default one as it uses separated 3D convolution
-    """
-    def __init__(self):
-        super(R2Plus1dStem, self).__init__(
-            nn.Conv3d(3, 45, kernel_size=(1, 7, 7),
-                      stride=(1, 2, 2), padding=(0, 3, 3),
-                      bias=False),
-            nn.BatchNorm3d(45),
-            nn.ReLU(inplace=True),
-            nn.Conv3d(45, 64, kernel_size=(3, 1, 1),
-                      stride=(1, 1, 1), padding=(1, 0, 0),
-                      bias=False),
-            nn.BatchNorm3d(64),
-            nn.ReLU(inplace=True))
-
-
-class VideoResNet(nn.Module):
-
-    def __init__(self, block, conv_makers, layers,
-                 stem, num_classes=16,
-                 zero_init_residual=False):
-        """Generic resnet video generator.
-        Args:
-            block (nn.Module): resnet building block
-            conv_makers (list(functions)): generator function for each layer
-            layers (List[int]): number of blocks per layer
-            stem (nn.Module, optional): Resnet stem, if None, defaults to conv-bn-relu. Defaults to None.
-            num_classes (int, optional): Dimension of the final FC layer. Defaults to 400.
-            zero_init_residual (bool, optional): Zero init bottleneck residual BN. Defaults to False.
-        """
-        super(VideoResNet, self).__init__()
-        self.inplanes = 64
-
-        self.stem = stem()
-
-        self.layer1 = self._make_layer(block, conv_makers[0], 64, layers[0], stride=1)
-        self.layer2 = self._make_layer(block, conv_makers[1], 128, layers[1], stride=2)
-        self.layer3 = self._make_layer(block, conv_makers[2], 192, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, conv_makers[3], 256, layers[3], stride=2)
-
-        self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1))
-        self.fc = nn.Linear(256 * block.expansion, num_classes)
-
-        # init weights
-        self._initialize_weights()
-
-        if zero_init_residual:
-            for m in self.modules():
-                if isinstance(m, Bottleneck):
-                    nn.init.constant_(m.bn3.weight, 0)
-
-    def forward(self, x):
-        x = self.stem(x)
-
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-
-        x = self.avgpool(x)
-        # Flatten the layer to fc
-        x = x.flatten(1)
-        x = self.fc(x)
-
-        return x
-
-    def _make_layer(self, block, conv_builder, planes, blocks, stride=1):
-        downsample = None
-
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            ds_stride = conv_builder.get_downsample_stride(stride)
-            downsample = nn.Sequential(
-                nn.Conv3d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=ds_stride, bias=False),
-                nn.BatchNorm3d(planes * block.expansion)
-            )
-        layers = []
-        layers.append(block(self.inplanes, planes, conv_builder, stride, downsample))
-
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes, conv_builder))
-
-        return nn.Sequential(*layers)
-
-    def _initialize_weights(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv3d):
-                nn.init.kaiming_normal_(m.weight, mode='fan_out',
-                                        nonlinearity='relu')
-                if m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm3d):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, 0, 0.01)
-                nn.init.constant_(m.bias, 0)
-
-
-def _video_resnet(arch, pretrained=False, progress=True, **kwargs):
-    model = VideoResNet(**kwargs)
-
-    return model
-
-
-def r3d_18(pretrained=False, progress=True, **kwargs):
-    """Construct 18 layer Resnet3D model as in
-    https://arxiv.org/abs/1711.11248
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on Kinetics-400
-        progress (bool): If True, displays a progress bar of the download to stderr
-    Returns:
-        nn.Module: R3D-18 network
-    """
-
-    return _video_resnet('r3d_18',
-                         pretrained, progress,
-                         block=BasicBlock,
-                         conv_makers=[Conv3DSimple] * 4,
-                         layers=[2, 2, 2, 2],
-                         stem=BasicStem, **kwargs)
-
-
-def mc3_18(pretrained=False, progress=True, **kwargs):
-    """Constructor for 18 layer Mixed Convolution network as in
-    https://arxiv.org/abs/1711.11248
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on Kinetics-400
-        progress (bool): If True, displays a progress bar of the download to stderr
-    Returns:
-        nn.Module: MC3 Network definition
-    """
-    return _video_resnet('mc3_18',
-                         pretrained, progress,
-                         block=BasicBlock,
-                         conv_makers=[Conv3DSimple] + [Conv3DNoTemporal] * 3,
-                         layers=[2, 2, 2, 2],
-                         stem=BasicStem, **kwargs)
-
-
-def r2plus1d_18(pretrained=False, progress=True, **kwargs):
-    """Constructor for the 18 layer deep R(2+1)D network as in
-    https://arxiv.org/abs/1711.11248
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on Kinetics-400
-        progress (bool): If True, displays a progress bar of the download to stderr
-    Returns:
-        nn.Module: R(2+1)D-18 network
-    """
-    return _video_resnet('r2plus1d_18',
-                         pretrained, progress,
-                         block=BasicBlock,
-                         conv_makers=[Conv2Plus1D] * 4,
-                         layers=[2, 2, 2, 2],
-                         stem=R2Plus1dStem, **kwargs)
-
-
-if __name__ == '__main__':
-
-    import torch
-
-    net = r3d_18().to(0)
-    x = torch.zeros(3, 1, 182, 218, 182).to(0)
-
-    print(net(x).shape)
-    print(net)

adrd/nn/resnet3d.py

@@ -1,256 +0,0 @@
-"""
-Simplified from torchvision.models.video.r3d_18. The citation information is
-shown below.
-
-@article{DBLP:journals/corr/abs-1711-11248,
-  author    = {Du Tran and
-               Heng Wang and
-               Lorenzo Torresani and
-               Jamie Ray and
-               Yann LeCun and
-               Manohar Paluri},
-  title     = {A Closer Look at Spatiotemporal Convolutions for Action Recognition},
-  journal   = {CoRR},
-  volume    = {abs/1711.11248},
-  year      = {2017},
-  url       = {http://arxiv.org/abs/1711.11248},
-  archivePrefix = {arXiv},
-  eprint    = {1711.11248},
-  timestamp = {Mon, 13 Aug 2018 16:46:39 +0200},
-  biburl    = {https://dblp.org/rec/journals/corr/abs-1711-11248.bib},
-  bibsource = {dblp computer science bibliography, https://dblp.org}
-}
-"""
-
-import torch.nn as nn
-
-
-class Conv3DSimple(nn.Conv3d):
-    def __init__(self,
-                 in_planes,
-                 out_planes,
-                 midplanes=None,
-                 stride=1,
-                 padding=1):
-
-        super().__init__(
-            in_channels=in_planes,
-            out_channels=out_planes,
-            kernel_size=(3, 3, 3),
-            stride=stride,
-            padding=padding,
-            bias=False)
-
-    @staticmethod
-    def get_downsample_stride(stride):
-        return stride, stride, stride
-    
-
-class BasicBlock(nn.Module):
-
-    expansion = 1
-
-    def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None):
-        midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes)
-
-        super(BasicBlock, self).__init__()
-        self.conv1 = nn.Sequential(
-            conv_builder(inplanes, planes, midplanes, stride),
-            nn.BatchNorm3d(planes),
-            nn.ReLU(inplace=True)
-        )
-        self.conv2 = nn.Sequential(
-            conv_builder(planes, planes, midplanes),
-            nn.BatchNorm3d(planes)
-        )
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.conv2(out)
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None):
-
-        super(Bottleneck, self).__init__()
-        midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes)
-
-        # 1x1x1
-        self.conv1 = nn.Sequential(
-            nn.Conv3d(inplanes, planes, kernel_size=1, bias=False),
-            nn.BatchNorm3d(planes),
-            nn.ReLU(inplace=True)
-        )
-        # Second kernel
-        self.conv2 = nn.Sequential(
-            conv_builder(planes, planes, midplanes, stride),
-            nn.BatchNorm3d(planes),
-            nn.ReLU(inplace=True)
-        )
-
-        # 1x1x1
-        self.conv3 = nn.Sequential(
-            nn.Conv3d(planes, planes * self.expansion, kernel_size=1, bias=False),
-            nn.BatchNorm3d(planes * self.expansion)
-        )
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.conv2(out)
-        out = self.conv3(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class BasicStem(nn.Sequential):
-    """The default conv-batchnorm-relu stem
-    """
-    def __init__(self):
-        super(BasicStem, self).__init__(
-            nn.Conv3d(1, 64, kernel_size=(7, 7, 7), stride=(2, 2, 2),
-                      padding=(3, 3, 3), bias=False),
-            nn.BatchNorm3d(64),
-            nn.ReLU(inplace=True))
-
-
-class VideoResNet(nn.Module):
-
-    def __init__(self, block, conv_makers, layers,
-                 stem, num_classes=2,
-                 zero_init_residual=False):
-        """Generic resnet video generator.
-        Args:
-            block (nn.Module): resnet building block
-            conv_makers (list(functions)): generator function for each layer
-            layers (List[int]): number of blocks per layer
-            stem (nn.Module, optional): Resnet stem, if None, defaults to conv-bn-relu. Defaults to None.
-            num_classes (int, optional): Dimension of the final FC layer. Defaults to 400.
-            zero_init_residual (bool, optional): Zero init bottleneck residual BN. Defaults to False.
-        """
-        super(VideoResNet, self).__init__()
-        self.inplanes = 64
-
-        self.stem = stem()
-
-        self.layer1 = self._make_layer(block, conv_makers[0], 64, layers[0], stride=1)
-        self.layer2 = self._make_layer(block, conv_makers[1], 128, layers[1], stride=2)
-        self.layer3 = self._make_layer(block, conv_makers[2], 192, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, conv_makers[3], 256, layers[3], stride=2)
-
-        self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1))
-        # self.fc = nn.Linear(256 * block.expansion, num_classes)
-
-        # init weights
-        self._initialize_weights()
-
-        if zero_init_residual:
-            for m in self.modules():
-                if isinstance(m, Bottleneck):
-                    nn.init.constant_(m.bn3.weight, 0)
-
-    def forward(self, x):
-        x = self.stem(x)
-
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-
-        x = self.avgpool(x)
-        # Flatten the layer to fc
-        x = x.flatten(1)
-        # x = self.fc(x)
-
-        return x
-
-    def _make_layer(self, block, conv_builder, planes, blocks, stride=1):
-        downsample = None
-
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            ds_stride = conv_builder.get_downsample_stride(stride)
-            downsample = nn.Sequential(
-                nn.Conv3d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=ds_stride, bias=False),
-                nn.BatchNorm3d(planes * block.expansion)
-            )
-        layers = []
-        layers.append(block(self.inplanes, planes, conv_builder, stride, downsample))
-
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes, conv_builder))
-
-        return nn.Sequential(*layers)
-
-    def _initialize_weights(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv3d):
-                nn.init.kaiming_normal_(m.weight, mode='fan_out',
-                                        nonlinearity='relu')
-                if m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.BatchNorm3d):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                nn.init.normal_(m.weight, 0, 0.01)
-                nn.init.constant_(m.bias, 0)
-
-
-def _video_resnet(arch, pretrained=False, progress=True, **kwargs):
-    model = VideoResNet(**kwargs)
-
-    return model
-
-
-def r3d_18(pretrained=False, progress=True, **kwargs):
-    """Construct 18 layer Resnet3D model as in
-    https://arxiv.org/abs/1711.11248
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on Kinetics-400
-        progress (bool): If True, displays a progress bar of the download to stderr
-    Returns:
-        nn.Module: R3D-18 network
-    """
-
-    return _video_resnet('r3d_18',
-                         pretrained, progress,
-                         block=BasicBlock,
-                         conv_makers=[Conv3DSimple] * 4,
-                         layers=[2, 2, 2, 2],
-                         stem=BasicStem, **kwargs)
-
-
-if __name__ == '__main__':
-    """ ... """
-    import torch
-
-    net = r3d_18().to('cuda:1')
-    x = torch.zeros(8, 1, 182, 218, 182).to('cuda:1')
-
-    print(net(x).shape)

adrd/nn/resnet_img_model.py

@@ -1,81 +0,0 @@
-import torch
-import torch.nn as nn
-import sys
-from icecream import ic
-# sys.path.append('/home/skowshik/ADRD_repo/adrd_tool/adrd/')
-from .net_resnet3d import r3d_18
-# from dev.data.dataset_csv import CSVDataset
-
-
-class ResNetModel(nn.Module):
-    ''' ... '''
-    def __init__(
-        self, 
-        tgt_modalities,
-        mri_feature = 'img_MRI_T1',
-    ):
-        ''' ... '''
-        super().__init__()
-
-        self.mri_feature = mri_feature
-
-        self.img_net_ = r3d_18()
-
-        # self.modules_emb_src = nn.Sequential(
-        #         nn.BatchNorm1d(9),
-        #         nn.Linear(9, d_model)
-        #     )
-
-        # classifiers (binary only)
-        self.modules_cls = nn.ModuleDict()
-        for k, info in tgt_modalities.items():
-            if info['type'] == 'categorical' and info['num_categories'] == 2:
-                # categorical
-                self.modules_cls[k] = nn.Linear(64, 1)
-
-            else:
-                # unrecognized
-                raise ValueError
-
-    def forward(self, x):
-        ''' ... '''
-        tgt_iter = self.modules_cls.keys()
-
-        img_x_batch = x[self.mri_feature]
-        img_out = self.img_net_(img_x_batch)
-
-        # ic(img_out.shape)
-
-        # run linear classifiers
-        out = [self.modules_cls[k](img_out).squeeze(1) for i, k in enumerate(tgt_iter)]
-        out = torch.stack(out, dim=1)
-
-        # ic(out.shape)
-
-        # out to dict
-        out = {k: out[:, i] for i, k in enumerate(tgt_iter)}
-            
-        return out
-
-
-if __name__ == '__main__':
-    ''' for testing purpose only '''
-    # import torch
-    # import numpy as np
-
-    # seed = 0
-    # print('Loading training dataset ... ')
-    # dat_trn = CSVDataset(mode=0, split=[1, 700], seed=seed)
-    # print(len(dat_trn))
-    # tgt_modalities = dat_trn.label_modalities
-    # net = ResNetModel(tgt_modalities).to('cuda')
-    # x = dat_trn.features
-    # x = {k: torch.as_tensor(np.array([x[i][k] for i in range(len(x))])).to('cuda') for k in x[0]}
-    # ic(x)
-    
-
-    # # print(net(x).shape)
-    # print(net(x))
-
-
-

adrd/nn/selfattention.py

@@ -1,62 +0,0 @@
-# Copyright (c) MONAI Consortium
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#     http://www.apache.org/licenses/LICENSE-2.0
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from monai.utils import optional_import
-import torch
-import torch.nn as nn
-
-
-Rearrange, _ = optional_import("einops.layers.torch", name="Rearrange")
-
-
-class SABlock(nn.Module):
-    """
-    A self-attention block, based on: "Dosovitskiy et al.,
-    An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale <https://arxiv.org/abs/2010.11929>"
-    """
-
-    def __init__(self, hidden_size: int, num_heads: int, dropout_rate: float = 0.0, qkv_bias: bool = False) -> None:
-        """
-        Args:
-            hidden_size: dimension of hidden layer.
-            num_heads: number of attention heads.
-            dropout_rate: faction of the input units to drop.
-            qkv_bias: bias term for the qkv linear layer.
-
-        """
-
-        super().__init__()
-
-        if not (0 <= dropout_rate <= 1):
-            raise ValueError("dropout_rate should be between 0 and 1.")
-
-        if hidden_size % num_heads != 0:
-            raise ValueError("hidden size should be divisible by num_heads.")
-
-        self.num_heads = num_heads
-        self.out_proj = nn.Linear(hidden_size, hidden_size)
-        self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias)
-        self.input_rearrange = Rearrange("b h (qkv l d) -> qkv b l h d", qkv=3, l=num_heads)
-        self.out_rearrange = Rearrange("b h l d -> b l (h d)")
-        self.drop_output = nn.Dropout(dropout_rate)
-        self.drop_weights = nn.Dropout(dropout_rate)
-        self.head_dim = hidden_size // num_heads
-        self.scale = self.head_dim**-0.5
-
-    def forward(self, x):
-        output = self.input_rearrange(self.qkv(x))
-        q, k, v = output[0], output[1], output[2]
-        att_mat = (torch.einsum("blxd,blyd->blxy", q, k) * self.scale).softmax(dim=-1)
-        att_mat = self.drop_weights(att_mat)
-        x = torch.einsum("bhxy,bhyd->bhxd", att_mat, v)
-        x = self.out_rearrange(x)
-        x = self.out_proj(x)
-        x = self.drop_output(x)
-        return x, att_mat

adrd/nn/transformer.py

@@ -1,268 +0,0 @@
-import torch
-import numpy as np
-from .. import nn
-# from ..nn import ImagingModelWrapper
-from .net_resnet3d import r3d_18
-from typing import Any, Type
-import math
-Tensor = Type[torch.Tensor]
-from icecream import ic
-ic.disable()
-
-class Transformer(torch.nn.Module):
-    ''' ... '''
-    def __init__(self,
-        src_modalities: dict[str, dict[str, Any]],
-        tgt_modalities: dict[str, dict[str, Any]],
-        d_model: int,
-        nhead: int,
-        num_encoder_layers: int = 1,
-        num_decoder_layers: int = 1,
-        device: str = 'cpu',
-        cuda_devices: list = [3],
-        img_net: str | None = None,
-        layers: int = 3,
-        img_size: int | None = 128,
-        patch_size: int | None = 16,
-        imgnet_ckpt: str | None = None,
-        train_imgnet: bool = False,
-        fusion_stage: str = 'middle',
-    ) -> None:
-        ''' ... '''
-        super().__init__()
-
-        self.d_model = d_model
-        self.nhead = nhead
-        self.num_encoder_layers = num_encoder_layers
-        self.num_decoder_layers = num_decoder_layers
-        self.img_net = img_net
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.imgnet_ckpt = imgnet_ckpt
-        self.train_imgnet = train_imgnet
-        self.layers = layers
-        self.src_modalities = src_modalities
-        self.tgt_modalities = tgt_modalities
-        self.device = device
-        self.fusion_stage = fusion_stage
-
-        # embedding modules for source
-    
-        self.modules_emb_src = torch.nn.ModuleDict()
-        print('Downsample layers: ', self.layers)
-        self.img_model = nn.ImagingModelWrapper(arch=self.img_net, img_size=self.img_size, patch_size=self.patch_size, ckpt_path=self.imgnet_ckpt, train_backbone=self.train_imgnet, layers=self.layers, out_dim=self.d_model, device=self.device, fusion_stage=self.fusion_stage)
-            
-        for k, info in src_modalities.items():
-            # ic(k)
-            # for key, val in info.items():
-                # ic(key, val)
-            if info['type'] == 'categorical':
-                self.modules_emb_src[k] = torch.nn.Embedding(info['num_categories'], d_model)
-            elif info['type'] == 'numerical':
-                self.modules_emb_src[k] = torch.nn.Sequential(
-                    torch.nn.BatchNorm1d(info['shape'][0]),
-                    torch.nn.Linear(info['shape'][0], d_model)
-                )
-            elif info['type'] == 'imaging':
-                # print(info['shape'], info['img_shape'])
-                if self.img_net:
-                    self.modules_emb_src[k] = self.img_model
-                    
-            else:
-                # unrecognized
-                raise ValueError('{} is an unrecognized data modality'.format(k))
-            
-        # positional encoding
-        self.pe = PositionalEncoding(d_model)
-
-        # auxiliary embedding vectors for targets
-        self.emb_aux = torch.nn.Parameter(
-            torch.zeros(len(tgt_modalities), 1, d_model),
-            requires_grad = True,
-        )
-        
-        # transformer
-        enc = torch.nn.TransformerEncoderLayer(
-            self.d_model, self.nhead,
-            dim_feedforward = self.d_model,
-            activation = 'gelu',
-            dropout = 0.3,
-        )
-        self.transformer = torch.nn.TransformerEncoder(enc, self.num_encoder_layers)
-
-
-        # classifiers (binary only)
-        self.modules_cls = torch.nn.ModuleDict()
-        for k, info in tgt_modalities.items():
-            if info['type'] == 'categorical' and info['num_categories'] == 2:
-                self.modules_cls[k] = torch.nn.Linear(d_model, 1)
-            else:
-                # unrecognized
-                raise ValueError
-            
-        # for n,p in self.named_parameters():
-        #     print(n, p.requires_grad)
-
-    def forward(self,
-        x: dict[str, Tensor],
-        mask: dict[str, Tensor],
-        # x_img: dict[str, Tensor] | Any = None,
-        skip_embedding: dict[str, bool] | None = None,
-        return_out_emb: bool = False,
-    ) -> dict[str, Tensor]:
-        """ ... """
-
-        out_emb = self.forward_emb(x, mask, skip_embedding)
-        if self.fusion_stage == "late":
-            out_emb = {k: v for k,v in out_emb.items() if "img_MRI" not in k}
-            img_out_emb = {k: v for k,v in out_emb.items() if "img_MRI" in k}
-            # for k,v in out_emb.items():
-                # print(k, v.size())
-            mask_nonimg = {k: v for k,v in mask.items() if "img_MRI" not in k}
-            out_trf = self.forward_trf(out_emb, mask_nonimg) # (8,128) + (8,50,128)
-            # print("out_trf: ", out_trf.size())
-            out_trf = torch.concatenate()
-        else:
-            out_trf = self.forward_trf(out_emb, mask)
-
-        out_cls = self.forward_cls(out_trf)
-            
-        if return_out_emb:
-            return out_emb, out_cls
-        return out_cls
-
-    def forward_emb(self,
-                    x: dict[str, Tensor],
-                    mask: dict[str, Tensor],
-                    skip_embedding: dict[str, bool] | None = None,
-                    # x_img: dict[str, Tensor] | Any = None,
-                ) -> dict[str, Tensor]:
-        """ ... """
-        # print("-------forward_emb--------")
-        out_emb = dict()
-        for k in self.modules_emb_src.keys():
-            if skip_embedding is None or k not in skip_embedding or not skip_embedding[k]:
-                if "img_MRI" in k:
-                    # print("img_MRI in ", k)
-                    if torch.all(mask[k]):
-                        if "swinunetr" in self.img_net.lower() and self.fusion_stage == 'late':
-                            out_emb[k] = torch.zeros((1,768,4,4,4))
-                        else:
-                            if 'cuda' in self.device:
-                                device = x[k].device
-                                # print(device)
-                            else:
-                                device = self.device
-                            out_emb[k] = torch.zeros((mask[k].shape[0], self.d_model)).to(device, non_blocking=True)
-                        # print("mask is True, out_emb[k]: ", out_emb[k].size())
-                    else:
-                        # print("calling modules_emb_src...")
-                        out_emb[k] = self.modules_emb_src[k](x[k])
-                        # print("mask is False, out_emb[k]: ", out_emb[k].size())
-                    
-                else:
-                    out_emb[k] = self.modules_emb_src[k](x[k])
-                    
-                # out_emb[k] = self.modules_emb_src[k](x[k])
-            else:
-                out_emb[k] = x[k]
-        return out_emb
-
-    def forward_trf(self,
-        out_emb: dict[str, Tensor],
-        mask: dict[str, Tensor],
-    ) -> dict[str, Tensor]:
-        """ ... """
-        # print('-----------forward_trf----------')
-        N = len(next(iter(out_emb.values())))  # batch size
-        S = len(self.modules_emb_src)  # number of sources
-        T = len(self.modules_cls)  # number of targets
-        if self.fusion_stage == 'late':
-            src_iter = [k for k in self.modules_emb_src.keys() if "img_MRI" not in k]
-            S = len(src_iter)  # number of sources
-
-        else:
-            src_iter = self.modules_emb_src.keys()
-        tgt_iter = self.modules_cls.keys()
-        
-        emb_src = torch.stack([o for o in out_emb.values()], dim=0)
-        # print('emb_src: ', emb_src.size())
-
-        self.pe.index = -1
-        emb_src = self.pe(emb_src)
-        # print('emb_src + pe: ', emb_src.size())
-        
-        # target embedding
-        # print('emb_aux: ', self.emb_aux.size())
-        emb_tgt = self.emb_aux.repeat(1, N, 1)
-        # print('emb_tgt: ', emb_tgt.size())
-        
-        # concatenate source embeddings and target embeddings
-        emb_all = torch.concatenate((emb_tgt, emb_src), dim=0)
-
-        # combine masks
-        mask_src = [mask[k] for k in src_iter]
-        mask_src = torch.stack(mask_src, dim=1)
-        
-        # target masks
-        mask_tgt = torch.zeros((N, T), dtype=torch.bool, device=self.emb_aux.device)
-        
-        # concatenate source masks and target masks
-        mask_all = torch.concatenate((mask_tgt, mask_src), dim=1)
-        
-        # repeat mask_all to fit transformer
-        mask_all = mask_all.unsqueeze(1).expand(-1, S + T, -1).repeat(self.nhead, 1, 1)
-
-        # run transformer
-        out_trf = self.transformer(
-            src = emb_all,
-            mask = mask_all,
-        )[0]
-        # print('out_trf: ', out_trf.size())
-        # out_trf = {k: out_trf[i] for i, k in enumerate(tgt_iter)}
-        return out_trf
-
-    def forward_cls(self,
-        out_trf: dict[str, Tensor],
-    ) -> dict[str, Tensor]:
-        """ ... """
-        tgt_iter = self.modules_cls.keys()
-        out_cls = {k: self.modules_cls[k](out_trf).squeeze(1) for k in tgt_iter}
-        return out_cls
-    
-class PositionalEncoding(torch.nn.Module):
-
-    def __init__(self, 
-        d_model: int, 
-        max_len: int = 512
-    ):
-        """ ... """
-        super().__init__()
-        position = torch.arange(max_len).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
-        pe = torch.zeros(max_len, 1, d_model)
-        pe[:, 0, 0::2] = torch.sin(position * div_term)
-        pe[:, 0, 1::2] = torch.cos(position * div_term)
-        self.register_buffer('pe', pe)
-        self.index = -1
-
-    def forward(self, x: Tensor, pe_type: str = 'non_img') -> Tensor:
-        """
-        Arguments:
-            x: Tensor, shape ``[seq_len, batch_size, embedding_dim]``
-        """
-        # print('pe: ', self.pe.size())
-        # print('x: ', x.size())
-        if pe_type == 'img':
-            self.index += 1
-            return x + self.pe[self.index]
-        else:
-            self.index += 1
-            return x + self.pe[self.index:x.size(0)+self.index]
-
-
-if __name__ == '__main__':
-    ''' for testing purpose only '''
-    pass
-
-

adrd/nn/unet.py

@@ -1,232 +0,0 @@
-import numpy as np
-import torch
-import torch.nn as nn
-import torchvision
-from torchvision import models
-from torch.nn import init
-import torch.nn.functional as F
-from icecream import ic
-
-
-class ContBatchNorm3d(nn.modules.batchnorm._BatchNorm):
-    def _check_input_dim(self, input):
-
-        if input.dim() != 5:
-            raise ValueError('expected 5D input (got {}D input)'.format(input.dim()))
-        #super(ContBatchNorm3d, self)._check_input_dim(input)
-
-    def forward(self, input):
-        self._check_input_dim(input)
-        return F.batch_norm(
-            input, self.running_mean, self.running_var, self.weight, self.bias,
-            True, self.momentum, self.eps)
-
-
-class LUConv(nn.Module):
-    def __init__(self, in_chan, out_chan, act):
-        super(LUConv, self).__init__()
-        self.conv1 = nn.Conv3d(in_chan, out_chan, kernel_size=3, padding=1)
-        self.bn1 = ContBatchNorm3d(out_chan)
-
-        if act == 'relu':
-            self.activation = nn.ReLU(out_chan)
-        elif act == 'prelu':
-            self.activation = nn.PReLU(out_chan)
-        elif act == 'elu':
-            self.activation = nn.ELU(inplace=True)
-        else:
-            raise
-
-    def forward(self, x):
-        out = self.activation(self.bn1(self.conv1(x)))
-        return out
-
-
-def _make_nConv(in_channel, depth, act, double_chnnel=False):
-    if double_chnnel:
-        layer1 = LUConv(in_channel, 32 * (2 ** (depth+1)),act)
-        layer2 = LUConv(32 * (2 ** (depth+1)), 32 * (2 ** (depth+1)),act)
-    else:
-        layer1 = LUConv(in_channel, 32*(2**depth),act)
-        layer2 = LUConv(32*(2**depth), 32*(2**depth)*2,act)
-
-    return nn.Sequential(layer1,layer2)
-
-
-class DownTransition(nn.Module):
-    def __init__(self, in_channel,depth, act):
-        super(DownTransition, self).__init__()
-        self.ops = _make_nConv(in_channel, depth,act)
-        self.maxpool = nn.MaxPool3d(2)
-        self.current_depth = depth
-
-    def forward(self, x):
-        if self.current_depth == 3:
-            out = self.ops(x)
-            out_before_pool = out
-        else:
-            out_before_pool = self.ops(x)
-            out = self.maxpool(out_before_pool)
-        return out, out_before_pool
-
-class UpTransition(nn.Module):
-    def __init__(self, inChans, outChans, depth,act):
-        super(UpTransition, self).__init__()
-        self.depth = depth
-        self.up_conv = nn.ConvTranspose3d(inChans, outChans, kernel_size=2, stride=2)
-        self.ops = _make_nConv(inChans+ outChans//2,depth, act, double_chnnel=True)
-
-    def forward(self, x, skip_x):
-        out_up_conv = self.up_conv(x)
-        concat = torch.cat((out_up_conv,skip_x),1)
-        out = self.ops(concat)
-        return out
-
-class OutputTransition(nn.Module):
-    def __init__(self, inChans, n_labels):
-
-        super(OutputTransition, self).__init__()
-        self.final_conv = nn.Conv3d(inChans, n_labels, kernel_size=1)
-        self.sigmoid = nn.Sigmoid()
-
-    def forward(self, x):
-        out = self.sigmoid(self.final_conv(x))
-        return out
-
-class ConvLayer(nn.Module):
-    def __init__(self, in_channels, out_channels, drop_rate, kernel, pooling, BN=True, relu_type='leaky'):
-        super().__init__()
-        kernel_size, kernel_stride, kernel_padding = kernel
-        pool_kernel, pool_stride, pool_padding = pooling
-        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, kernel_stride, kernel_padding, bias=False)
-        self.pooling = nn.MaxPool3d(pool_kernel, pool_stride, pool_padding)
-        self.BN = nn.BatchNorm3d(out_channels)
-        self.relu = nn.LeakyReLU(inplace=False) if relu_type=='leaky' else nn.ReLU(inplace=False)
-        self.dropout = nn.Dropout(drop_rate, inplace=False) 
-       
-    def forward(self, x):
-        x = self.conv(x)
-        x = self.pooling(x)
-        x = self.BN(x)
-        x = self.relu(x)
-        x = self.dropout(x)
-        return x
-    
-class AttentionModule(nn.Module):
-    def __init__(self, in_channels, out_channels, drop_rate=0.1):
-        super(AttentionModule, self).__init__()
-        self.conv = nn.Conv3d(in_channels, out_channels, 1, 1, 0, bias=False)
-        self.attention = ConvLayer(in_channels, out_channels, drop_rate, (1, 1, 0), (1, 1, 0))
-
-    def forward(self, x, return_attention=True):
-        feats = self.conv(x)
-        att = F.softmax(self.attention(x))
-
-        out = feats * att
-
-        if return_attention:
-            return att, out
-        
-        return out 
-
-class UNet3D(nn.Module):
-    # the number of convolutions in each layer corresponds
-    # to what is in the actual prototxt, not the intent
-    def __init__(self, n_class=1, act='relu', pretrained=False, input_size=(1,1,182,218,182), attention=False, drop_rate=0.1, blocks=4):
-        super(UNet3D, self).__init__()
-
-        self.blocks = blocks
-        self.down_tr64 = DownTransition(1,0,act)
-        self.down_tr128 = DownTransition(64,1,act)
-        self.down_tr256 = DownTransition(128,2,act)
-        self.down_tr512 = DownTransition(256,3,act)
-
-        self.up_tr256 = UpTransition(512, 512,2,act)
-        self.up_tr128 = UpTransition(256,256, 1,act)
-        self.up_tr64 = UpTransition(128,128,0,act)
-        self.out_tr = OutputTransition(64, 1)
-
-        self.pretrained = pretrained
-        self.attention = attention
-        if pretrained:
-            print("Using image pretrained model checkpoint")
-            weight_dir = '/home/skowshik/ADRD_repo/img_pretrained_ckpt/Genesis_Chest_CT.pt'
-            checkpoint = torch.load(weight_dir)
-            state_dict = checkpoint['state_dict']
-            unParalled_state_dict = {}
-            for key in state_dict.keys():
-                unParalled_state_dict[key.replace("module.", "")] = state_dict[key]
-            self.load_state_dict(unParalled_state_dict)
-            del self.up_tr256
-            del self.up_tr128
-            del self.up_tr64
-            del self.out_tr
-        
-        if self.blocks == 5:
-            self.down_tr1024 = DownTransition(512,4,act)
-            
-
-        # self.conv1 = nn.Conv3d(512, 256, 1, 1, 0, bias=False)
-        # self.conv2 = nn.Conv3d(256, 128, 1, 1, 0, bias=False)
-        # self.conv3 = nn.Conv3d(128, 64, 1, 1, 0, bias=False)
-
-        if attention:
-            self.attention_module = AttentionModule(1024 if self.blocks==5 else 512, n_class, drop_rate=drop_rate)
-        # Output.
-        self.avgpool = nn.AvgPool3d((6,7,6), stride=(6,6,6))
-
-        dummy_inp = torch.rand(input_size)
-        dummy_feats = self.forward(dummy_inp, stage='get_features')
-        dummy_feats = dummy_feats[0]
-        self.in_features = list(dummy_feats.shape)
-        ic(self.in_features)
-
-        self._init_weights()
-
-    def _init_weights(self):
-        if not self.pretrained:
-            for m in self.modules():
-                if isinstance(m, nn.Conv3d):
-                    init.kaiming_normal_(m.weight)
-                elif isinstance(m, ContBatchNorm3d):
-                    init.constant_(m.weight, 1)
-                    init.constant_(m.bias, 0)
-                elif isinstance(m, nn.Linear):
-                    init.kaiming_normal_(m.weight)
-                    init.constant_(m.bias, 0)
-        elif self.attention:
-            for m in self.attention_module.modules():
-                if isinstance(m, nn.Conv3d):
-                    init.kaiming_normal_(m.weight)
-                elif isinstance(m, nn.BatchNorm3d):
-                    init.constant_(m.weight, 1)
-                    init.constant_(m.bias, 0)
-        else:
-            pass
-        # Zero initialize the last batchnorm in each residual branch.
-        # for m in self.modules():
-        #     if isinstance(m, BottleneckBlock):
-        #         init.constant_(m.out_conv.bn.weight, 0)
-    
-    def forward(self, x, stage='normal', attention=False):
-        ic('backbone forward')
-        self.out64, self.skip_out64 = self.down_tr64(x)
-        self.out128,self.skip_out128 = self.down_tr128(self.out64)
-        self.out256,self.skip_out256 = self.down_tr256(self.out128)
-        self.out512,self.skip_out512 = self.down_tr512(self.out256)
-        if self.blocks == 5:
-            self.out1024,self.skip_out1024 = self.down_tr1024(self.out512)
-            ic(self.out1024.shape)
-        # self.out = self.conv1(self.out512)
-        # self.out = self.conv2(self.out)
-        # self.out = self.conv3(self.out)
-        # self.out = self.conv(self.out)
-        ic(hasattr(self, 'attention_module'))
-        if hasattr(self, 'attention_module'):
-            att, feats = self.attention_module(self.out1024 if self.blocks==5 else self.out512)
-        else:
-            feats = self.out1024 if self.blocks==5 else self.out512
-        ic(feats.shape)
-        if attention:
-            return att, feats
-        return feats

adrd/nn/unet_3d.py

@@ -1,63 +0,0 @@
-import sys
-sys.path.append('..')
-# from feature_extractor.for_image_data.backbone import CNN_GAP, ResNet3D, UNet3D
-import torch
-import torch.nn as nn
-from torchvision import models
-import torch.nn.functional as F
-# from . import UNet3D
-from .unet import UNet3D
-from icecream import ic
-
-
-class UNet3DBase(nn.Module):
-    def __init__(self, n_class=1, act='relu', attention=False, pretrained=False, drop_rate=0.1, blocks=4):
-        super(UNet3DBase, self).__init__()
-        model = UNet3D(n_class=n_class, attention=attention, pretrained=pretrained, blocks=blocks)
-
-        self.blocks = blocks
-
-        self.down_tr64 = model.down_tr64
-        self.down_tr128 = model.down_tr128
-        self.down_tr256 = model.down_tr256
-        self.down_tr512 = model.down_tr512
-        if self.blocks == 5:
-            self.down_tr1024 = model.down_tr1024
-        # self.block_modules = nn.ModuleList([self.down_tr64, self.down_tr128, self.down_tr256, self.down_tr512])
-
-        self.in_features = model.in_features
-        # ic(attention)
-        if attention:
-            self.attention_module = model.attention_module
-        #     self.attention_module = AttentionModule(512, n_class, drop_rate=drop_rate)
-        # self.avgpool = nn.AvgPool3d((6,7,6), stride=(6,6,6))
-
-    def forward(self, x, stage='normal', attention=False):
-        # ic('UNet3DBase forward')
-        self.out64, self.skip_out64 = self.down_tr64(x)
-        # ic(self.out64.shape, self.skip_out64.shape)
-        self.out128,self.skip_out128 = self.down_tr128(self.out64)
-        # ic(self.out128.shape, self.skip_out128.shape)
-        self.out256,self.skip_out256 = self.down_tr256(self.out128)
-        # ic(self.out256.shape, self.skip_out256.shape)
-        self.out512,self.skip_out512 = self.down_tr512(self.out256)
-        # ic(self.out512.shape, self.skip_out512.shape)
-        if self.blocks == 5:
-            self.out1024,self.skip_out1024 = self.down_tr1024(self.out512)
-        # ic(self.out1024.shape, self.skip_out1024.shape)
-        # ic(hasattr(self, 'attention_module'))
-        if hasattr(self, 'attention_module'):
-            att, feats = self.attention_module(self.out1024 if self.blocks == 5 else self.out512)
-        else:
-            feats = self.out1024 if self.blocks == 5 else self.out512
-        # ic(feats.shape)
-        if attention:
-            return att, feats
-        return feats
-
-        # self.out_up_256 = self.up_tr256(self.out512,self.skip_out256)
-        # self.out_up_128 = self.up_tr128(self.out_up_256, self.skip_out128)
-        # self.out_up_64 = self.up_tr64(self.out_up_128, self.skip_out64)
-        # self.out = self.out_tr(self.out_up_64)
-
-        # return self.out

adrd/nn/unet_img_model.py

@@ -1,211 +0,0 @@
-from pyexpat import features
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.cuda.amp import autocast
-import numpy as np
-import re
-from icecream import ic
-import math
-import torch.nn.utils.weight_norm as weightNorm
-
-# from . import UNet3DBase
-from .unet_3d import UNet3DBase
-
-
-def init_weights(m):
-    classname = m.__class__.__name__
-    if classname.find('Conv2d') != -1 or classname.find('ConvTranspose2d') != -1:
-        nn.init.kaiming_uniform_(m.weight)
-        nn.init.zeros_(m.bias)
-    elif classname.find('BatchNorm') != -1:
-        nn.init.normal_(m.weight, 1.0, 0.02)
-        nn.init.zeros_(m.bias)
-    elif classname.find('Linear') != -1:
-        nn.init.xavier_normal_(m.weight)
-        nn.init.zeros_(m.bias)
-
-class feat_classifier(nn.Module):
-    def __init__(self, class_num, bottleneck_dim=256, type="linear"):
-        super(feat_classifier, self).__init__()
-        self.type = type
-        # if type in ['conv', 'gap'] and len(bottleneck_dim) > 3:
-            # bottleneck_dim = bottleneck_dim[-3:]
-        ic(bottleneck_dim)
-        if type == 'wn':
-            self.layer = weightNorm(
-                nn.Linear(bottleneck_dim[1:], class_num), name="weight")
-            # self.fc.apply(init_weights)
-        elif type == 'gap':
-            if len(bottleneck_dim) > 3:
-                bottleneck_dim = bottleneck_dim[-3:]
-            self.layer = nn.AvgPool3d(bottleneck_dim, stride=(1,1,1))
-        elif type == 'conv':
-            if len(bottleneck_dim) > 3:
-                bottleneck_dim = bottleneck_dim[-4:]
-            ic(bottleneck_dim)
-            self.layer = nn.Conv3d(bottleneck_dim[0], class_num, kernel_size=bottleneck_dim[1:])
-            ic(self.layer)
-        else:
-            print('bottleneck dim: ', bottleneck_dim)
-            self.layer = nn.Sequential(
-                            torch.nn.Flatten(start_dim=1, end_dim=-1),
-                            nn.Linear(math.prod(bottleneck_dim), class_num)
-            )
-        self.layer.apply(init_weights)
-
-    def forward(self, x):
-        # print('=> feat_classifier forward')
-        # ic(x.size())
-        x = self.layer(x)
-        # ic(x.size())
-        if self.type in ['gap','conv']:
-            x = torch.squeeze(x)
-            if len(x.shape) < 2:
-                x = torch.unsqueeze(x,0)
-        # print('returning x: ', x.size())
-        return x
-
-class ImageModel(nn.Module):
-    """
-    Empirical Risk Minimization (ERM)
-    """
-
-    def __init__(
-            self, 
-            counts=None,
-            classifier='gap',
-            accum_iter=8,
-            save_emb=False,
-            # ssl,
-            num_classes=1,
-            load_img_ckpt=False,
-        ):
-        super(ImageModel, self).__init__()
-        if counts is not None:
-            if isinstance(counts[0], list):
-                counts = np.stack(counts, axis=0).sum(axis=0)
-                print('counts: ', counts)
-                total = np.sum(counts)
-                print(total/counts)
-                self.weight = total/torch.FloatTensor(counts)
-            else:
-                total = sum(counts)
-                self.weight = torch.FloatTensor([total/c for c in counts])
-        else:
-            self.weight = None
-        print('weight: ', self.weight)
-        # device = torch.device(f'cuda:{args.gpu_id}' if args.gpu_id is not None else 'cpu')
-        self.criterion = nn.CrossEntropyLoss(weight=self.weight)
-        # if ssl:
-        #     # add contrastive loss
-        #     # self.ssl_criterion = 
-        #     pass
-
-        self.featurizer = UNet3DBase(n_class=num_classes, attention=True, pretrained=load_img_ckpt)
-        self.classifier = feat_classifier(
-            num_classes, self.featurizer.in_features, classifier)
-
-        self.network = nn.Sequential(
-            self.featurizer, self.classifier)
-        self.accum_iter = accum_iter
-        self.acc_steps = 0
-        self.save_embedding = save_emb
-
-    def update(self, minibatches, opt, sch, scaler):
-        print('--------------def update----------------')
-        device = list(self.parameters())[0].device
-        all_x = torch.cat([data[1].to(device).float() for data in minibatches])
-        all_y = torch.cat([data[2].to(device).long() for data in minibatches])
-        print('all_x: ', all_x.size())
-        # all_p = self.predict(all_x)
-        # all_probs =  
-        label_list = all_y.tolist()
-        count = float(len(label_list))
-        ic(count)
-            
-        uniques = sorted(list(set(label_list)))
-        ic(uniques)
-        counts = [float(label_list.count(i)) for i in uniques]
-        ic(counts)
-        
-        weights = [count / c for c in counts]
-        ic(weights)
-        
-        with autocast():
-            loss = self.criterion(self.predict(all_x), all_y)
-        self.acc_steps += 1
-        print('class: ', loss.item())
-
-        scaler.scale(loss / self.accum_iter).backward()
-        
-        if self.acc_steps == self.accum_iter:
-            scaler.step(opt)
-            if sch:
-                sch.step()
-            scaler.update()
-            self.zero_grad()
-            self.acc_steps = 0
-            torch.cuda.empty_cache()
-            
-        del all_x
-        del all_y
-        return {'class': loss.item()}, sch
-
-    def forward(self, *args, **kwargs):
-        return self.network(*args, **kwargs)
-    
-    def predict(self, x, stage='normal', attention=False):
-        # print('network device: ', list(self.network.parameters())[0].device)
-        # print('x device: ', x.device)
-        if stage == 'get_features' or self.save_embedding:
-            feats = self.network[0](x, attention=attention)
-            output = self.network[1](feats[-1] if attention else feats)
-            return feats, output
-        else:
-            return self.network(x)
-
-    def extract_features(self, x, attention=False):
-        feats = self.network[0](x, attention=attention)
-        return feats
-
-    def load_checkpoint(self, state_dict):
-        try:
-            self.load_checkpoint_helper(state_dict)
-        except:
-            featurizer_dict = {}
-            net_dict = {}
-            for key,val in state_dict.items():
-                if 'featurizer' in key:
-                    featurizer_dict[key] = val
-                elif 'network' in key:
-                    net_dict[key] = val
-            self.featurizer.load_state_dict(featurizer_dict)
-            self.classifier.load_state_dict(net_dict)
-
-    def load_checkpoint_helper(self, state_dict):
-        try:
-            self.load_state_dict(state_dict)
-            print('try: loaded')
-        except RuntimeError as e:
-            print('--> except')
-            if 'Missing key(s) in state_dict:' in str(e):
-                state_dict = {
-                    key.replace('module.', '', 1): value
-                    for key, value in state_dict.items()
-                }
-                state_dict = {
-                    key.replace('featurizer.', '', 1).replace('classifier.','',1): value
-                    for key, value in state_dict.items()
-                }
-                state_dict = {
-                    re.sub('network.[0-9].', '', key): value
-                    for key, value in state_dict.items()
-                }
-                try:
-                    del state_dict['criterion.weight']
-                except:
-                    pass
-                self.load_state_dict(state_dict)
-                
-                print('except: loaded')

adrd/nn/vitautoenc.py

@@ -1,163 +0,0 @@
-# Copyright (c) MONAI Consortium
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#     http://www.apache.org/licenses/LICENSE-2.0
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-from monai.networks.blocks.patchembedding import PatchEmbeddingBlock
-from monai.networks.layers import Conv
-from monai.utils import ensure_tuple_rep
-
-from typing import Sequence, Union
-import torch
-import torch.nn as nn
-
-from ..nn.blocks import TransformerBlock
-from icecream import ic
-ic.disable()
-
-__all__ = ["ViTAutoEnc"]
-
-
-class ViTAutoEnc(nn.Module):
-    """
-    Vision Transformer (ViT), based on: "Dosovitskiy et al.,
-    An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale <https://arxiv.org/abs/2010.11929>"
-
-    Modified to also give same dimension outputs as the input size of the image
-    """
-
-    def __init__(
-        self,
-        in_channels: int,
-        img_size: Union[Sequence[int], int],
-        patch_size: Union[Sequence[int], int],
-        out_channels: int = 1,
-        deconv_chns: int = 16,
-        hidden_size: int = 768,
-        mlp_dim: int = 3072,
-        num_layers: int = 12,
-        num_heads: int = 12,
-        pos_embed: str = "conv",
-        dropout_rate: float = 0.0,
-        spatial_dims: int = 3,
-    ) -> None:
-        """
-        Args:
-            in_channels: dimension of input channels or the number of channels for input
-            img_size: dimension of input image.
-            patch_size: dimension of patch size.
-            hidden_size: dimension of hidden layer.
-            out_channels: number of output channels.
-            deconv_chns: number of channels for the deconvolution layers.
-            mlp_dim: dimension of feedforward layer.
-            num_layers: number of transformer blocks.
-            num_heads: number of attention heads.
-            pos_embed: position embedding layer type.
-            dropout_rate: faction of the input units to drop.
-            spatial_dims: number of spatial dimensions.
-
-        Examples::
-
-            # for single channel input with image size of (96,96,96), conv position embedding and segmentation backbone
-            # It will provide an output of same size as that of the input
-            >>> net = ViTAutoEnc(in_channels=1, patch_size=(16,16,16), img_size=(96,96,96), pos_embed='conv')
-
-            # for 3-channel with image size of (128,128,128), output will be same size as of input
-            >>> net = ViTAutoEnc(in_channels=3, patch_size=(16,16,16), img_size=(128,128,128), pos_embed='conv')
-
-        """
-
-        super().__init__()
-
-        self.patch_size = ensure_tuple_rep(patch_size, spatial_dims)
-        self.spatial_dims = spatial_dims
-        self.hidden_size = hidden_size
-        
-        self.patch_embedding = PatchEmbeddingBlock(
-            in_channels=in_channels,
-            img_size=img_size,
-            patch_size=patch_size,
-            hidden_size=hidden_size,
-            num_heads=num_heads,
-            pos_embed=pos_embed,
-            dropout_rate=dropout_rate,
-            spatial_dims=self.spatial_dims,
-        )
-        self.blocks = nn.ModuleList(
-            [TransformerBlock(hidden_size, mlp_dim, num_heads, dropout_rate) for i in range(num_layers)]
-        )
-        self.norm = nn.LayerNorm(hidden_size)
-
-        new_patch_size = [4] * self.spatial_dims
-        conv_trans = Conv[Conv.CONVTRANS, self.spatial_dims]
-        # self.conv3d_transpose* is to be compatible with existing 3d model weights.
-        self.conv3d_transpose = conv_trans(hidden_size, deconv_chns, kernel_size=new_patch_size, stride=new_patch_size)
-        self.conv3d_transpose_1 = conv_trans(
-            in_channels=deconv_chns, out_channels=out_channels, kernel_size=new_patch_size, stride=new_patch_size
-        )
-
-    def forward(self, x, return_emb=False, return_hiddens=False):
-        """
-        Args:
-            x: input tensor must have isotropic spatial dimensions,
-                such as ``[batch_size, channels, sp_size, sp_size[, sp_size]]``.
-        """
-        spatial_size = x.shape[2:]
-        x = self.patch_embedding(x)
-        hidden_states_out = []
-        for blk in self.blocks:
-            x = blk(x)
-            hidden_states_out.append(x)
-        x = self.norm(x)
-        x = x.transpose(1, 2)
-        if return_emb:
-            return x
-        d = [s // p for s, p in zip(spatial_size, self.patch_size)]
-        x = torch.reshape(x, [x.shape[0], x.shape[1], *d])
-        x = self.conv3d_transpose(x)
-        x = self.conv3d_transpose_1(x)
-        if return_hiddens:
-            return x, hidden_states_out
-        return x
-    
-    def get_last_selfattention(self, x):
-        """
-        Args:
-            x: input tensor must have isotropic spatial dimensions,
-                such as ``[batch_size, channels, sp_size, sp_size[, sp_size]]``.
-        """
-        x = self.patch_embedding(x)
-        ic(x.size())
-        for i, blk in enumerate(self.blocks):
-            if i < len(self.blocks) - 1:
-                x = blk(x)
-                x.size()
-            else:
-                return blk(x, return_attention=True)
-            
-    def load(self, ckpt_path, map_location='cpu', checkpoint_key='state_dict'):
-        """
-        Args:
-            ckpt_path: path to the pretrained weights
-            map_location: device to load the checkpoint on
-        """
-        state_dict = torch.load(ckpt_path, map_location=map_location)
-        ic(state_dict['epoch'], state_dict['train_loss'])
-        if checkpoint_key in state_dict:
-            print(f"Take key {checkpoint_key} in provided checkpoint dict")
-            state_dict = state_dict[checkpoint_key]
-        # remove `module.` prefix
-        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
-        # remove `backbone.` prefix induced by multicrop wrapper
-        state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
-        msg = self.load_state_dict(state_dict, strict=False)
-        print('Pretrained weights found at {} and loaded with msg: {}'.format(ckpt_path, msg))
-
-