astroclip.yaml

seed_everything: 42
trainer:
  default_root_dir: "{ASTROCLIP_ROOT}/outputs"
  enable_checkpointing: true
  gradient_clip_val: 1.
  max_epochs: 100
  precision: null
  callbacks:
    - class_path: LearningRateMonitor
      init_args:
        logging_interval: "step"
    - class_path: ModelCheckpoint
      init_args:
        save_last: link
        save_top_k: 2
        every_n_epochs: 1
        monitor: "val_loss_nologit"
  logger:
    class_path: CustomWandbLogger
    init_args:
      project: "astroclip-alignment"
      entity: "{WANDB_ENTITY_NAME}"
      save_dir: ${trainer.default_root_dir}
  detect_anomaly: True
model:
  class_path: AstroClipModel
  init_args:
    image_encoder:
      class_path: astroclip.models.astroclip.ImageHead
      init_args:
        config: "astroclip/astrodino/config.yaml"
        model_weights: "{ASTROCLIP_ROOT}/pretrained/astrodino.ckpt"
        save_directory: "{ASTROCLIP_ROOT}/outputs/astrodino"
    spectrum_encoder:
      class_path: astroclip.models.astroclip.SpectrumHead
      init_args:
        model_path: "{ASTROCLIP_ROOT}/pretrained/specformer.ckpt"
data:
  class_path: AstroClipDataloader
  init_args:
    path: "{ASTROCLIP_ROOT}/datasets/astroclip_file/"
    columns:
      - image
      - spectrum
    batch_size: 256
    num_workers: 8
    collate_fn:
      class_path: astroclip.data.AstroClipCollator
      init_args:
        center_crop: 144
optimizer:
  class_path: torch.optim.adamw.AdamW
  init_args:
    lr: 1e-4
    weight_decay: 0.05
lr_scheduler:
  class_path: astroclip.CosineAnnealingWithWarmupLR
  init_args:
    T_max: 10_000
    T_warmup: 1_000
    eta_min: ${eval:'${optimizer.init_args.lr}//500'}