configs.py

import os
from dataclasses import dataclass

from jsonargparse import CLI
from jsonargparse import set_docstring_parse_options

INF = 2147483647

from docstring_parser import DocstringStyle

set_docstring_parse_options(style=DocstringStyle.GOOGLE, attribute_docstrings=True)


@dataclass
class BasicConfig:
    """Basic configurations
    Args:
        seed: seed for random number generator
        epochs: number of training epochs
    """
    fold_idx: int = 0
    """fold index [0-4]"""
    seed: int = 0
    """seed number"""
    epochs: int = 100
    """number of training epochs"""
    criterion: str = 'ce'
    """loss function"""
    optim: str = 'adam'
    suffix: str = ''
    wandb: bool = True
    wdb_group: str = None
    debug: bool = False
    print_freq: int = 100
    """print frequency"""
    save_freq: int = 50
    """save frequency"""
    workers: int = 10
    """number of CPUs"""
    ckpt: str = None
    """path to the checkpoint file"""
    resumed: bool = False
    """resume from checkpoint"""
    freeze_backbone: bool = False
    """either train or freeze the backbone"""
    exp_dir: str = ''
    """path to experiment directory"""
    metadata_path: str = 'data/metadata.csv'
    """path to metadata file containing the split info"""
    entropic_scale: int = 10
    """entropic scale in isomax loss"""
    use_coteaching: bool = False
    """apply co-teaching training"""
    forget_rate: float = 0.0
    """default forget rate"""
    eval_freq: int = 5
    """number of training epochs between two consecutive evaluations"""
    loss_cfg_path: str = 'utils/loss_functions/config'
    """path to the config for the loss functions"""
    oversampling: bool = False
    """Whether to oversample the minority class (only implemented for n-classes = 2)"""
    start_eval_epoch: int = 10
    """Epoch in which the evaluation starts"""


@dataclass
class DatasetConfig:
    """Basic dataset configurations"""
    input_size: tuple = (32, 32)
    """image input size"""
    num_classes: int = 2
    """number of classes"""
    batch_size: int = 100
    """batch size for both training and test"""
    in_channel: int = 1
    """number of input channels"""
    oversampling_cancer: int = -1
    """rate of cancer patch duplication"""
    min_inv: float = 0.4
    """minimum involvement of cancer in the training set"""
    min_inv_eval: float = 0.4
    """minimum involvement of cancer in the evaluation set"""
    min_gs: int = 7
    """minimum Gleason Score in the dataset"""
    patch_dim: tuple = (48, 32, 0)
    """patch-size, vertical stride, horizontal stride"""
    in_prostate: bool = True
    """limit the patches and loss calculations to inside prostate only"""
    patch_gen_config: dict = None
    """configuration for patch generation"""
    # Options when using co-teaching
    forget_rate: float = 0.0
    """percentage of training data that will be filtered"""
    exponent: int = 1
    """exponent for co-teaching"""
    num_gradual: int = 10
    """number of epochs to reach the forget-rate"""
    patch_save_dir: str = None
    """director for storing extracted patches, subfolder 'patches/{patch_dim}' will be created"""
    save_patches: bool = False
    """whether to save patches"""
    model_name: str = 'inception'
    """name of the model [inception, timesnet, transformer]"""
    patch_downsampling_factor: float = 0.25
    """downsampling factor for patches when training multivariate models"""
    # iLR
    c_thr: float = 0.9
    """Confidence threshold for label correction"""
    c_margin: float = 0.15
    """error margin between true and predicted inv for label correction"""
    c_retain: int = 5
    """number of epochs that the corrected labels are retained"""
    seq_len: int = 100
    """input sequence length"""
    aug_1d: bool = False
    """perform 1d augmentation on the timeseries"""
    expand_amount: int = 20
    """number of patches that can be added to the training set"""
    ilr: bool = True
    """whether to do label refinement"""

    def __post_init__(self):
        patch_dim = self.patch_dim
        if self.patch_save_dir is not None:
            if patch_dim is None:
                raise ValueError('patch_dim cannot be None when saving patches')
            self.patch_save_dir: str = (f'{self.patch_save_dir}/patches/'
                                        f'{patch_dim[0]}_str{patch_dim[1]}x{patch_dim[2]}')


@dataclass
class OptimizerConfig:
    """Basic optimizer configurations"""
    learning_rate: float = 1e-3
    """learning rate"""
    weight_decay: float = 0.
    """weight decay"""
    gamma: float = 0.7
    """learning rate decay"""
    momentum: float = 0.9
    """momentum"""
    scheduler: str = None
    """learning rate scheduler, could be none or onecycle"""


@dataclass
class TimeSeriesModelConfig:
    dropout: float = 0.1
    """dropout probability"""
    d_model: int = 16  # 512
    """dimension of model"""
    d_ff: int = 64  # 2048
    """dimension of fcn"""
    top_k: int = 5
    """for TimesBlock"""
    num_kernels: int = 6
    """for Inception"""
    embed: str = 'timeF'
    """time features encoding, options:[timeF, fixed, learned]"""
    e_layers: int = 2
    """num of encoder layers"""
    d_layers: int = 1
    """num of decoder layers"""
    output_attention: bool = False
    """whether to output attention in encoder"""
    n_heads: int = 8
    """number of attention heads"""
    task_name = 'classification'
    """task name for the model"""
    activation: str = 'gelu'
    """activation"""
    label_len = 0
    """start token length"""
    pred_len = 0
    """prediction uence length"""
    enc_in: int = 1
    """encoder input size"""
    dec_in: int = 1
    """decoder input size"""
    freq = 'h'
    """freq for time features encoding, options:[s:secondly, t:minutely, h:hourly, d:daily, b:business days, w:weekly, 
    m:monthly], you can also use more detailed freq like 15min or 3h"""
    c_out = 0
    """output size"""
    factor = 1
    """attention factor"""
    # inception1d
    num_blocks: int = 10
    """number of inception blocks"""
    out_channels: int = 16
    """number of output channels"""
    bottleneck_channels: int = 12
    """number of output bottleneck channels"""
    kernel_sizes: int = 15
    """kernel sizes"""
    stride: int = 1
    """stride"""


@dataclass
class BasicTraining(BasicConfig, DatasetConfig, OptimizerConfig, TimeSeriesModelConfig):
    """Configuration for training with cross entropy loss and one 2D 32x32 patch per core per epoch."""


@dataclass
class PatchTo1DTraining(BasicTraining):
    """options for training 1D model with one patch per core each epoch"""
    in_channel: int = 100
    dataset_name = 'patch_1d'


@dataclass
class PatchTo1DMTraining(BasicTraining):
    """options for training 1D model with one patch per core each epoch"""
    in_channel: int = 100
    dataset_name = 'patch_1dm'


@dataclass
class PatchTo2DTraining(BasicTraining):
    """options for training 2D model with one patch per core each epoch"""
    in_channel: int = 3
    dataset_name = 'patch_2d'


@dataclass
class PatchTo3DTraining(BasicTraining):
    """options for training 2D model with one patch per core each epoch"""
    in_channel: int = 100
    dataset_name = 'patch_3d'


@dataclass
class PatchTeUSTraining(BasicTraining):
    """options for training Multi-SWAV with one patch per core each epoch"""
    in_channel: int = 3
    dataset_name = 'patch_teus'


@dataclass
class NeedleTeUSTraining(BasicTraining):
    """options for training Multi-SWAV with one patch per core each epoch"""
    in_channel: int = 3
    dataset_name = 'needle_teus'
    input_size: tuple = (96, 96)


@dataclass
class PatchesTo1DTraining(PatchTo1DTraining):
    """options for training 1D model with all patches per core each epoch"""
    dataset_name = 'patches_1d'


@dataclass
class PatchesTo1DMTraining(PatchTo1DMTraining):
    """options for training 1D model with all patches per core each epoch"""
    dataset_name = 'patches_1dm'


@dataclass
class PatchesTo2DTraining(PatchTo2DTraining):
    """options for training 2D model with all patches per core each epoch"""
    dataset_name = 'patches_2d'


@dataclass
class PatchesTo3DTraining(PatchTo3DTraining):
    """options for training 2D model with all patches per core each epoch"""
    dataset_name = 'patches_3d'


@dataclass
class PatchesTeUSTraining(PatchTeUSTraining):
    """options for training TeUSFormer with all patches per core each epoch"""
    dataset_name = 'patches_teus'


@dataclass
class RfAatPTraining(BasicTraining):
    """options for training analytical RF image centered at the prostate"""
    input_size: tuple = (224, 224)
    """Input size"""
    in_channel: int = 3
    """number of input channels"""
    dataset_name = 'rf_a'


DATASET_OPTIONS = {
    'patch_1d': PatchTo1DTraining,
    'patch_1dm': PatchTo1DMTraining,
    'patch_2d': PatchTo2DTraining,
    'patch_3d': PatchTo3DTraining,
    'patch_teus': PatchTeUSTraining,
    'needle_teus': NeedleTeUSTraining,
    'patches_1d': PatchesTo1DTraining,
    'patches_1dm': PatchesTo1DMTraining,
    'patches_2d': PatchesTo2DTraining,
    'patches_3d': PatchesTo3DTraining,
    'patches_teus': PatchesTeUSTraining,
    'rf_a': RfAatPTraining,
}

CONFIG_OPTIONS = {
    'bc': BasicConfig,
    'dc': DatasetConfig,
    'optim': OptimizerConfig,
}


def unit_test():
    # print(CLI(Dummy))
    # from jsonargparse import compose_dataclasses
    # dtc = compose_dataclasses(RfAatPTraining)

    # from jsonargparse import ArgumentParser
    # parser = ArgumentParser()
    # parser.add_class_arguments(BasicConfig)
    # parser.add_class_arguments(DatasetConfig)
    # args = parser.parse_args()
    print(CLI(dict(DATASET_OPTIONS, **CONFIG_OPTIONS), as_positional=False))


if __name__ == '__main__':
    unit_test()