lits_balanced.yaml

---
trainer:
  class_path: eva.Trainer
  init_args:
    n_runs: &N_RUNS ${oc.env:N_RUNS, 1}
    default_root_dir: &OUTPUT_ROOT ${oc.env:OUTPUT_ROOT, logs/${oc.env:MODEL_NAME, vit_small_patch16_224_dino}/lits_balanced}
    max_steps: &MAX_STEPS ${oc.env:MAX_STEPS, 500000}
    plugins:
      class_path: eva.core.plugins.SubmoduleTorchCheckpointIO
      init_args:
        submodule: decoder
    callbacks:
      - class_path: eva.callbacks.ConfigurationLogger
      - class_path: lightning.pytorch.callbacks.TQDMProgressBar
        init_args:
          refresh_rate: ${oc.env:TQDM_REFRESH_RATE, 1}
      - class_path: eva.vision.callbacks.SemanticSegmentationLogger
        init_args:
          log_every_n_epochs: 1
          log_images: false
      - class_path: lightning.pytorch.callbacks.ModelCheckpoint
        init_args:
          filename: best
          save_last: true
          save_top_k: 1
          monitor: &MONITOR_METRIC ${oc.env:MONITOR_METRIC, val/GeneralizedDiceScore}
          mode: &MONITOR_METRIC_MODE ${oc.env:MONITOR_METRIC_MODE, max}
      - class_path: lightning.pytorch.callbacks.EarlyStopping
        init_args:
          min_delta: 0
          patience: 100
          monitor: *MONITOR_METRIC
          mode: *MONITOR_METRIC_MODE
      - class_path: eva.callbacks.SegmentationEmbeddingsWriter
        init_args:
          output_dir: &DATASET_EMBEDDINGS_ROOT ${oc.env:EMBEDDINGS_ROOT, ./data/embeddings}/${oc.env:MODEL_NAME, vit_small_patch16_224_dino}/lits_balanced
          dataloader_idx_map:
            0: train
            1: val
            2: test
          metadata_keys: ["slice_index"]
          overwrite: false
          backbone:
            class_path: eva.vision.models.ModelFromRegistry
            init_args:
              model_name: ${oc.env:MODEL_NAME, universal/vit_small_patch16_224_dino}
              model_kwargs:
                out_indices: ${oc.env:OUT_INDICES, 1}
              model_extra_kwargs: ${oc.env:MODEL_EXTRA_KWARGS, null}
    logger:
      - class_path: lightning.pytorch.loggers.TensorBoardLogger
        init_args:
          save_dir: *OUTPUT_ROOT
          name: ""
model:
  class_path: eva.vision.models.modules.SemanticSegmentationModule
  init_args:
    decoder:
      class_path: eva.vision.models.networks.decoders.segmentation.ConvDecoderMS
      init_args:
        in_features: ${oc.env:IN_FEATURES, 384}
        num_classes:  &NUM_CLASSES 3
    criterion:
      class_path: eva.core.losses.CrossEntropyLoss
      init_args:
        weight: [0.05, 0.1, 1.5]
    optimizer:
      class_path: torch.optim.AdamW
      init_args:
        lr: ${oc.env:LR_VALUE, 0.002}
    lr_scheduler:
      class_path: torch.optim.lr_scheduler.PolynomialLR
      init_args:
        total_iters: *MAX_STEPS
        power: 0.9
    postprocess:
      predictions_transforms:
        - class_path: torch.argmax
          init_args:
            dim: 1
    metrics:
      common:
        - class_path: eva.metrics.AverageLoss
      evaluation:
        - class_path: eva.vision.metrics.defaults.MulticlassSegmentationMetrics
          init_args:
            num_classes: *NUM_CLASSES
        - class_path: torchmetrics.ClasswiseWrapper
          init_args:
            metric:
              class_path: eva.vision.metrics.GeneralizedDiceScore
              init_args:
                num_classes: *NUM_CLASSES
                weight_type: linear
                per_class: true
data:
  class_path: eva.DataModule
  init_args:
    datasets:
      train:
        class_path: eva.vision.datasets.EmbeddingsSegmentationDataset
        init_args: &DATASET_ARGS
          root: *DATASET_EMBEDDINGS_ROOT
          manifest_file: manifest.csv
          split: train
      val:
        class_path: eva.vision.datasets.EmbeddingsSegmentationDataset
        init_args:
          <<: *DATASET_ARGS
          split: val
      test:
        class_path: eva.vision.datasets.EmbeddingsSegmentationDataset
        init_args:
          <<: *DATASET_ARGS
          split: test
      predict:
        - class_path: eva.vision.datasets.LiTSBalanced
          init_args: &PREDICT_DATASET_ARGS
            root: ${oc.env:DATA_ROOT, ./data/lits}
            split: train
            transforms:
              class_path: eva.vision.data.transforms.common.ResizeAndClamp
              init_args:
                size: ${oc.env:RESIZE_DIM, 224}
                clamp_range: [-1008, 822]
                mean: &NORMALIZE_MEAN ${oc.env:NORMALIZE_MEAN, [0.485, 0.456, 0.406]} 
                std: &NORMALIZE_STD ${oc.env:NORMALIZE_STD, [0.229, 0.224, 0.225]}
        - class_path: eva.vision.datasets.LiTSBalanced
          init_args:
            <<: *PREDICT_DATASET_ARGS
            split: val
        - class_path: eva.vision.datasets.LiTSBalanced
          init_args:
            <<: *PREDICT_DATASET_ARGS
            split: test
    dataloaders:
      train:
        batch_size: &BATCH_SIZE ${oc.env:BATCH_SIZE, 64}
        num_workers: &N_DATA_WORKERS ${oc.env:N_DATA_WORKERS, 4}
        shuffle: true
      val:
        batch_size: *BATCH_SIZE
        num_workers: *N_DATA_WORKERS
      test:
        batch_size: *BATCH_SIZE
        num_workers: *N_DATA_WORKERS
      predict:
        batch_size: &PREDICT_BATCH_SIZE ${oc.env:PREDICT_BATCH_SIZE, 64}
        num_workers: *N_DATA_WORKERS