Update configs

21 files changed, 298 insertions, 0 deletions

diff --git a/training/conf/callbacks/htr.yaml b/training/conf/callbacks/htr.yaml
new file mode 100644
index 0000000..51c68c5
--- /dev/null
+++ b/training/conf/callbacks/htr.yaml
@@ -0,0 +1,6 @@
+defaults:
+  - default
+  - wandb: watch
+  - wandb: config 
+  - wandb: checkpoints
+  - wandb: htr_predictions
diff --git a/training/conf/callbacks/lightning/checkpoint.yaml b/training/conf/callbacks/lightning/checkpoint.yaml
new file mode 100644
index 0000000..b4101d8
--- /dev/null
+++ b/training/conf/callbacks/lightning/checkpoint.yaml
@@ -0,0 +1,9 @@
+model_checkpoint:
+  _target_: pytorch_lightning.callbacks.ModelCheckpoint
+  monitor: val/loss # name of the logged metric which determines when model is improving
+  save_top_k: 1 # save k best models (determined by above metric)
+  save_last: true # additionaly always save model from last epoch
+  mode: min # can be "max" or "min"
+  verbose: false
+  dirpath: checkpoints/
+  filename: "{epoch:02d}"
diff --git a/training/conf/callbacks/lightning/early_stopping.yaml b/training/conf/callbacks/lightning/early_stopping.yaml
new file mode 100644
index 0000000..a188df3
--- /dev/null
+++ b/training/conf/callbacks/lightning/early_stopping.yaml
@@ -0,0 +1,6 @@
+early_stopping:
+  _target_: pytorch_lightning.callbacks.EarlyStopping
+  monitor: val/loss # name of the logged metric which determines when model is improving
+  patience: 16 # how many epochs of not improving until training stops
+  mode: min # can be "max" or "min"
+  min_delta: 0 # minimum change in the monitored metric needed to qualify as an improvement
diff --git a/training/conf/callbacks/lightning/learning_rate_monitor.yaml b/training/conf/callbacks/lightning/learning_rate_monitor.yaml
new file mode 100644
index 0000000..4a14e1f
--- /dev/null
+++ b/training/conf/callbacks/lightning/learning_rate_monitor.yaml
@@ -0,0 +1,4 @@
+learning_rate_monitor:
+  _target_: pytorch_lightning.callbacks.LearningRateMonitor
+  logging_interval: step
+  log_momentum: false
diff --git a/training/conf/callbacks/vae.yaml b/training/conf/callbacks/vae.yaml
new file mode 100644
index 0000000..eec2c1f
--- /dev/null
+++ b/training/conf/callbacks/vae.yaml
@@ -0,0 +1,6 @@
+defaults:
+  - default
+  - wandb: watch
+  - wandb: checkpoints
+  - wandb: reconstructions
+  - wandb: config 
diff --git a/training/conf/callbacks/wandb/checkpoints.yaml b/training/conf/callbacks/wandb/checkpoints.yaml
new file mode 100644
index 0000000..a4a16ff
--- /dev/null
+++ b/training/conf/callbacks/wandb/checkpoints.yaml
@@ -0,0 +1,4 @@
+upload_ckpts_as_artifact:
+  _target_: callbacks.wandb_callbacks.UploadCheckpointsAsArtifact
+  ckpt_dir: checkpoints/
+  upload_best_only: true
diff --git a/training/conf/callbacks/wandb/config.yaml b/training/conf/callbacks/wandb/config.yaml
new file mode 100644
index 0000000..747a7c6
--- /dev/null
+++ b/training/conf/callbacks/wandb/config.yaml
@@ -0,0 +1,2 @@
+upload_code_as_artifact:
+  _target_: callbacks.wandb_callbacks.UploadConfigAsArtifact
diff --git a/training/conf/callbacks/wandb/predictions.yaml b/training/conf/callbacks/wandb/predictions.yaml
new file mode 100644
index 0000000..573fa96
--- /dev/null
+++ b/training/conf/callbacks/wandb/predictions.yaml
@@ -0,0 +1,3 @@
+log_text_predictions:
+  _target_: callbacks.wandb_callbacks.LogTextPredictions
+  num_samples: 8
diff --git a/training/conf/callbacks/wandb/reconstructions.yaml b/training/conf/callbacks/wandb/reconstructions.yaml
new file mode 100644
index 0000000..92f2d12
--- /dev/null
+++ b/training/conf/callbacks/wandb/reconstructions.yaml
@@ -0,0 +1,4 @@
+log_image_reconstruction:
+  _target_: callbacks.wandb_callbacks.LogReconstuctedImages
+  num_samples: 8
+  use_sigmoid: true
diff --git a/training/conf/callbacks/wandb/swa.yaml b/training/conf/callbacks/wandb/swa.yaml
new file mode 100644
index 0000000..73f8c66
--- /dev/null
+++ b/training/conf/callbacks/wandb/swa.yaml
@@ -0,0 +1,7 @@
+stochastic_weight_averaging:
+  _target_: pytorch_lightning.callbacks.StochasticWeightAveraging
+  swa_epoch_start: 0.8
+  swa_lrs: 0.05
+  annealing_epochs: 10
+  annealing_strategy: cos
+  device: null
diff --git a/training/conf/callbacks/wandb/watch.yaml b/training/conf/callbacks/wandb/watch.yaml
new file mode 100644
index 0000000..511608c
--- /dev/null
+++ b/training/conf/callbacks/wandb/watch.yaml
@@ -0,0 +1,4 @@
+watch_model:
+  _target_: callbacks.wandb_callbacks.WatchModel
+  log: all
+  log_freq: 100
diff --git a/training/conf/datamodule/target_transform/word_pieces.yaml b/training/conf/datamodule/target_transform/word_pieces.yaml
new file mode 100644
index 0000000..8ace2af
--- /dev/null
+++ b/training/conf/datamodule/target_transform/word_pieces.yaml
@@ -0,0 +1,3 @@
+word_pieces:
+  _target_: text_recognizer.data.transforms.word_pieces.WordPieces
+  max_len: 451
diff --git a/training/conf/datamodule/test_transform/default.yaml b/training/conf/datamodule/test_transform/default.yaml
new file mode 100644
index 0000000..d6f6d78
--- /dev/null
+++ b/training/conf/datamodule/test_transform/default.yaml
@@ -0,0 +1,2 @@
+to_tensor:
+  _target_: torchvision.transforms.ToTensor
diff --git a/training/conf/datamodule/test_transform/iam_lines.yaml b/training/conf/datamodule/test_transform/iam_lines.yaml
new file mode 100644
index 0000000..a14ec82
--- /dev/null
+++ b/training/conf/datamodule/test_transform/iam_lines.yaml
@@ -0,0 +1,6 @@
+embed_crop:
+  _target_: text_recognizer.data.transforms.embed_crop.EmbedCrop
+  augment: true
+
+to_tensor:
+  _target_: torchvision.transforms.ToTensor
diff --git a/training/conf/datamodule/test_transform/paragraphs.yaml b/training/conf/datamodule/test_transform/paragraphs.yaml
new file mode 100644
index 0000000..a69eae9
--- /dev/null
+++ b/training/conf/datamodule/test_transform/paragraphs.yaml
@@ -0,0 +1,6 @@
+center_crop:
+  _target_: torchvision.transforms.CenterCrop
+  size: [576, 640]
+
+to_tensor:
+  _target_: torchvision.transforms.ToTensor
diff --git a/training/conf/datamodule/transform/default.yaml b/training/conf/datamodule/transform/default.yaml
new file mode 100644
index 0000000..d6f6d78
--- /dev/null
+++ b/training/conf/datamodule/transform/default.yaml
@@ -0,0 +1,2 @@
+to_tensor:
+  _target_: torchvision.transforms.ToTensor
diff --git a/training/conf/datamodule/transform/emnist_lines.yaml b/training/conf/datamodule/transform/emnist_lines.yaml
new file mode 100644
index 0000000..986ba20
--- /dev/null
+++ b/training/conf/datamodule/transform/emnist_lines.yaml
@@ -0,0 +1,14 @@
+to_tensor:
+  _target_: torchvision.transforms.ToTensor
+
+color_jitter:
+  _target_: torchvision.transforms.ColorJitter
+  brightness: [0.5, 1.0]
+
+random_affine:
+  _target_: torchvision.transforms.RandomAffine
+  degrees: 3
+  translate: [0.0, 0.05]
+  scale: [0.4, 1.1]
+  shear: [-40, 50]
+  fill: 0
diff --git a/training/conf/datamodule/transform/iam_lines.yaml b/training/conf/datamodule/transform/iam_lines.yaml
new file mode 100644
index 0000000..710f6e3
--- /dev/null
+++ b/training/conf/datamodule/transform/iam_lines.yaml
@@ -0,0 +1,17 @@
+embed_crop:
+  _target_: text_recognizer.data.transforms.embed_crop.EmbedCrop
+  augment: false
+
+color_jitter:
+  _target_: torchvision.transforms.ColorJitter
+  brightness: [0.8, 1.6]
+
+random_affine:
+  _target_: torchvision.transforms.RandomAffine
+  degrees: 1
+  shear: [-30, 20]
+  interpolation: BILINEAR
+  fill: 0
+
+to_tensor:
+  _target_: torchvision.transforms.ToTensor
diff --git a/training/conf/datamodule/transform/iam_lines_barlow.yaml b/training/conf/datamodule/transform/iam_lines_barlow.yaml
new file mode 100644
index 0000000..b44ca4e
--- /dev/null
+++ b/training/conf/datamodule/transform/iam_lines_barlow.yaml
@@ -0,0 +1,39 @@
+barlow:
+  _target_: text_recognizer.data.transforms.barlow.BarlowTransform
+  prim:
+    embed_crop:
+      _target_: text_recognizer.data.transforms.embed_crop.EmbedCrop
+      augment: true
+
+    color_jitter:
+      _target_: torchvision.transforms.ColorJitter
+      brightness: [0.8, 1.6]
+
+    random_affine:
+      _target_: torchvision.transforms.RandomAffine
+      degrees: 1
+      shear: [-30, 20]
+      interpolation: BILINEAR
+      fill: 0
+
+    to_tensor:
+      _target_: torchvision.transforms.ToTensor
+
+  bis:
+    embed_crop:
+      _target_: text_recognizer.data.transforms.embed_crop.EmbedCrop
+      augment: true
+
+    color_jitter:
+      _target_: torchvision.transforms.ColorJitter
+      brightness: [0.5, 1.4]
+
+    random_affine:
+      _target_: torchvision.transforms.RandomAffine
+      degrees: 1.5
+      shear: [-30, 20]
+      interpolation: BILINEAR
+      fill: 0
+
+    to_tensor:
+      _target_: torchvision.transforms.ToTensor
diff --git a/training/conf/datamodule/transform/paragraphs.yaml b/training/conf/datamodule/transform/paragraphs.yaml
new file mode 100644
index 0000000..d33a4c3
--- /dev/null
+++ b/training/conf/datamodule/transform/paragraphs.yaml
@@ -0,0 +1,20 @@
+random_crop:
+  _target_: torchvision.transforms.RandomCrop
+  size: [576, 640]
+  padding: null
+  pad_if_needed: true
+  fill: 0
+  padding_mode: constant
+
+color_jitter:
+  _target_: torchvision.transforms.ColorJitter
+  brightness: [0.8, 1.6]
+
+random_affine:
+  _target_: torchvision.transforms.RandomAffine
+  degrees: 1
+  shear: [-10, 10]
+  interpolation: BILINEAR
+
+to_tensor:
+  _target_: torchvision.transforms.ToTensor
diff --git a/training/conf/experiment/barlow_twins.yaml b/training/conf/experiment/barlow_twins.yaml
new file mode 100644
index 0000000..4901e18
--- /dev/null
+++ b/training/conf/experiment/barlow_twins.yaml
@@ -0,0 +1,134 @@
+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]]
+
+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: 2.0e-4
+    momentum: 0.9
+    weight_decay: 0
+    eps: 1.0e-6
+
+    parameters: network
+
+lr_schedulers:
+  network:
+    _target_: torch.optim.lr_scheduler.OneCycleLR
+    max_lr: 2.0e-4
+    total_steps: null
+    epochs: *epochs
+    steps_per_epoch: 632
+    pct_start: 0.3
+    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_extended_paragraphs.IAMExtendedParagraphs
+  batch_size: 4
+  num_workers: 12
+  train_fraction: 0.8
+  augment: true
+  pin_memory: true
+  word_pieces: false
+  resize: null
+
+network:
+  _target_: text_recognizer.networks.conv_transformer.ConvTransformer
+  input_dims: [1, 576, 640]
+  hidden_dim: &hidden_dim 128
+  encoder_dim: 1280
+  dropout_rate: 0.2
+  num_classes: *num_classes
+  pad_index: *ignore_index
+  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
+  decoder:
+    _target_: text_recognizer.networks.transformer.Decoder
+    dim: *hidden_dim
+    depth: 3 
+    num_heads: 4
+    attn_fn: text_recognizer.networks.transformer.attention.Attention
+    attn_kwargs:
+      dim_head: 32
+      dropout_rate: 0.2
+    norm_fn: text_recognizer.networks.transformer.norm.ScaleNorm
+    ff_fn: text_recognizer.networks.transformer.mlp.FeedForward
+    ff_kwargs:
+      dim_out: null
+      expansion_factor: 4
+      glu: true
+      dropout_rate: 0.2
+    cross_attend: true
+    pre_norm: true
+    rotary_emb:
+      _target_: text_recognizer.networks.transformer.positional_encodings.rotary_embedding.RotaryEmbedding
+      dim: 32
+  pixel_pos_embedding:
+    _target_: text_recognizer.networks.transformer.positional_encodings.PositionalEncoding2D
+    hidden_dim: *hidden_dim 
+    max_h: 18
+    max_w: 20
+  token_pos_embedding:
+    _target_: text_recognizer.networks.transformer.positional_encodings.PositionalEncoding
+    hidden_dim: *hidden_dim 
+    dropout_rate: 0.2
+    max_len: *max_output_len
+
+model:
+  _target_: text_recognizer.models.transformer.TransformerLitModel
+  max_output_len: *max_output_len
+  start_token: <s>
+  end_token: <e>
+  pad_token: <p>
+
+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: 16
+  overfit_batches: 0