# @package _global_

defaults:
  - override /criterion: null
  - override /datamodule: null
  - override /network: null
  - override /model: null
  - override /lr_schedulers: null
  - override /optimizers: null

epochs: &epochs 1000
summary: [[1, 1, 56, 1024]]

criterion:
  _target_: text_recognizer.criterions.barlow_twins.BarlowTwinsLoss
  dim: 512
  lambda_: 3.9e-3

callbacks:
  stochastic_weight_averaging:
    _target_: pytorch_lightning.callbacks.StochasticWeightAveraging
    swa_epoch_start: 0.75
    swa_lrs: 1.0e-5
    annealing_epochs: 10
    annealing_strategy: cos
    device: null

optimizers:
  madgrad:
    _target_: madgrad.MADGRAD
    lr: 1.0e-3
    momentum: 0.9
    weight_decay: 1.0e-6
    eps: 1.0e-6
    parameters: network

lr_schedulers:
  network:
    _target_: torch.optim.lr_scheduler.OneCycleLR
    max_lr: 3.0e-4
    total_steps: null
    epochs: *epochs
    steps_per_epoch: 45
    pct_start: 0.03
    anneal_strategy: cos
    cycle_momentum: true
    base_momentum: 0.85
    max_momentum: 0.95
    div_factor: 25
    final_div_factor: 1.0e4
    three_phase: false
    last_epoch: -1
    verbose: false
    # Non-class arguments
    interval: step
    monitor: val/loss

datamodule:
    _target_: text_recognizer.data.iam_lines.IAMLines
    batch_size: 16
    num_workers: 12
    train_fraction: 0.9
    pin_memory: false
    transform: transform/iam_lines_barlow.yaml
    test_transform: transform/iam_lines_barlow.yaml
    mapping: 
      _target_: text_recognizer.data.mappings.emnist_mapping.EmnistMapping

network:
  _target_: text_recognizer.networks.barlow_twins.network.BarlowTwins
  encoder:
    _target_: text_recognizer.networks.encoders.efficientnet.EfficientNet
    arch: b0
    out_channels: 1280
    stochastic_dropout_rate: 0.2
    bn_momentum: 0.99
    bn_eps: 1.0e-3
  projector:
    _target_: text_recognizer.networks.barlow_twins.projector.Projector
    dims: [1280, 512, 512, 512]

model:
  _target_: text_recognizer.models.barlow_twins.BarlowTwinsLitModel

trainer:
  _target_: pytorch_lightning.Trainer
  stochastic_weight_avg: true
  auto_scale_batch_size: binsearch
  auto_lr_find: false
  gradient_clip_val: 0.0
  fast_dev_run: false
  gpus: 1
  precision: 16
  max_epochs: *epochs
  terminate_on_nan: true
  weights_summary: null
  limit_train_batches: 1.0 
  limit_val_batches: 1.0
  limit_test_batches: 1.0
  resume_from_checkpoint: null
  accumulate_grad_batches: 32
  overfit_batches: 0