Remove lr args from model, add Cosine lr, fix to vqvae stack

8 files changed, 22 insertions, 12 deletions

diff --git a/training/conf/experiment/vqvae.yaml b/training/conf/experiment/vqvae.yaml
index eb40f3b..7a9e643 100644
--- a/training/conf/experiment/vqvae.yaml
+++ b/training/conf/experiment/vqvae.yaml
@@ -5,6 +5,7 @@ defaults:
   - override /criterion: mse
   - override /model: lit_vqvae
   - override /callbacks: wandb_vae
+  - override /lr_scheduler: cosine_annealing
 
 trainer:
   max_epochs: 64
@@ -13,11 +14,11 @@ trainer:
 datamodule:
   batch_size: 16
 
-lr_scheduler:
-  epochs: 64
-  steps_per_epoch: 1245
+# lr_scheduler:
+  # epochs: 64
+  # steps_per_epoch: 1245
 
-optimizer:
-  lr: 1.0e-3
+# optimizer:
+  # lr: 1.0e-3
 
 summary: [1, 576, 640]
diff --git a/training/conf/lr_scheduler/cosine_annealing.yaml b/training/conf/lr_scheduler/cosine_annealing.yaml
new file mode 100644
index 0000000..62667bb
--- /dev/null
+++ b/training/conf/lr_scheduler/cosine_annealing.yaml
@@ -0,0 +1,7 @@
+_target_: torch.optim.lr_scheduler.CosineAnnealingLR
+T_max: 64
+eta_min: 0.0
+last_epoch: -1
+
+interval: epoch
+monitor: val/loss
diff --git a/training/conf/lr_scheduler/one_cycle.yaml b/training/conf/lr_scheduler/one_cycle.yaml
index eecee8a..fb5987a 100644
--- a/training/conf/lr_scheduler/one_cycle.yaml
+++ b/training/conf/lr_scheduler/one_cycle.yaml
@@ -13,3 +13,7 @@ final_div_factor: 10000.0
 three_phase: true
 last_epoch: -1
 verbose: false
+
+# Non-class arguments
+interval: step
+monitor: val/loss
diff --git a/training/conf/model/lit_vqvae.yaml b/training/conf/model/lit_vqvae.yaml
index 409fa0d..632668b 100644
--- a/training/conf/model/lit_vqvae.yaml
+++ b/training/conf/model/lit_vqvae.yaml
@@ -1,4 +1,2 @@
 _target_: text_recognizer.models.vqvae.VQVAELitModel
-interval: step
-monitor: val/loss
-latent_loss_weight: 1.0
+latent_loss_weight: 0.25
diff --git a/training/conf/network/decoder/vae_decoder.yaml b/training/conf/network/decoder/vae_decoder.yaml
index b2090b3..0a36a54 100644
--- a/training/conf/network/decoder/vae_decoder.yaml
+++ b/training/conf/network/decoder/vae_decoder.yaml
@@ -1,5 +1,5 @@
 _target_: text_recognizer.networks.vqvae.decoder.Decoder
 out_channels: 1 
 hidden_dim: 32
-channels_multipliers: [8, 6, 2, 1]
+channels_multipliers: [4, 4, 2, 1]
 dropout_rate: 0.25
diff --git a/training/conf/network/encoder/vae_encoder.yaml b/training/conf/network/encoder/vae_encoder.yaml
index 5dc6814..dacd389 100644
--- a/training/conf/network/encoder/vae_encoder.yaml
+++ b/training/conf/network/encoder/vae_encoder.yaml
@@ -1,5 +1,5 @@
 _target_: text_recognizer.networks.vqvae.encoder.Encoder
 in_channels: 1 
 hidden_dim: 32
-channels_multipliers: [1, 2, 6, 8]
+channels_multipliers: [1, 2, 4, 4]
 dropout_rate: 0.25
diff --git a/training/conf/network/vqvae.yaml b/training/conf/network/vqvae.yaml
index 835d0b7..d97e9b6 100644
--- a/training/conf/network/vqvae.yaml
+++ b/training/conf/network/vqvae.yaml
@@ -3,7 +3,7 @@ defaults:
   - decoder: vae_decoder
 
 _target_: text_recognizer.networks.vqvae.vqvae.VQVAE
-hidden_dim: 256
+hidden_dim: 128
 embedding_dim: 32
 num_embeddings: 1024
 decay: 0.99
diff --git a/training/conf/optimizer/madgrad.yaml b/training/conf/optimizer/madgrad.yaml
index 46b2fff..458b116 100644
--- a/training/conf/optimizer/madgrad.yaml
+++ b/training/conf/optimizer/madgrad.yaml
@@ -1,5 +1,5 @@
 _target_: madgrad.MADGRAD
-lr: 2.0e-4
+lr: 3.0e-4
 momentum: 0.9
 weight_decay: 0
 eps: 1.0e-6