seed_everything: 42
trainer:
  accelerator: gpu
  strategy: auto
  devices: 4
  num_nodes: 1
  precision: 32-true
  callbacks:
    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
      init_args:
        logging_interval: step
    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
      init_args:
        # save every epoch
        monitor: loss/validate
        save_top_k: 3
        every_n_epochs: 1
  fast_dev_run: false
  max_epochs: 50
  max_steps: -1
  check_val_every_n_epoch: 1
  log_every_n_steps: 10
  enable_checkpointing: true
  inference_mode: true
  use_distributed_sampler: true
model:
  architecture: UNeXt2
  model_config:
    in_channels: 1
    out_channels: 2
    in_stack_depth: 5
    backbone: convnextv2_tiny
    pretrained: false
    stem_kernel_size:
      - 5
      - 4
      - 4
    decoder_mode: pixelshuffle
    decoder_conv_blocks: 2
    head_pool: true
    head_expansion_ratio: 4
    drop_path_rate: 0.0
  loss_function:
    class_path: viscy.light.engine.MixedLoss
    init_args:
      l1_alpha: 0.5
      l2_alpha: 0.0
      ms_dssim_alpha: 0.5
  lr: 0.0002
  schedule: WarmupCosine
  log_batches_per_epoch: 8
  log_samples_per_batch: 1
data:
  # modify this line
  data_path: training_and_validation_dataset.zarr
  source_channel: &s reconstructed-labelfree
  target_channel:
    - &t1 reconstructed-nucleus
    - &t2 reconstructed-membrane
  z_window_size: 5
  split_ratio: 0.8
  batch_size: 32
  num_workers: 16
  yx_patch_size: [384, 384]
  normalizations:
    - class_path: viscy.transforms.NormalizeSampled
      init_args:
        keys: [*s, *t1, *t2]
        level: fov_statistics
        subtrahend: median
        divisor: iqr
  augmentations:
    - class_path: viscy.transforms.RandWeightedCropd
      init_args:
        keys: [*s, *t1, *t2]
        w_key: *t1
        spatial_size: [-1, 512, 512]
        num_samples: 2
    - class_path: viscy.transforms.RandAffined
      init_args:
        keys: [*s, *t1, *t2]
        prob: 0.5
        rotate_range:
          - 3.14
          - 0.0
          - 0.0
        shear_range:
          - 0.0
          - 0.05
          - 0.05
        scale_range:
          - 0.5
          - 0.5
          - 0.5
    - class_path: viscy.transforms.RandAdjustContrastd
      init_args:
        keys: *s
        prob: 0.3
        gamma:
          - 0.75
          - 1.5
    - class_path: viscy.transforms.RandScaleIntensityd
      init_args:
        keys: *s
        factors: 0.5
        prob: 0.5
    - class_path: viscy.transforms.RandGaussianNoised
      init_args:
        keys: *s
        prob: 0.5
        mean: 0.0
        std: 5.0
    - class_path: viscy.transforms.RandGaussianSmoothd
      init_args:
        keys: *s
        prob: 0.5
        sigma_x:
          - 0.25
          - 1.5
        sigma_y:
          - 0.25
          - 1.5
        sigma_z:
          - 0.25
          - 1.5
  caching: false
