From 263f2b7158d76bc0adad45309625910c0fa7b1fe Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Fri, 6 Aug 2021 14:19:37 +0200 Subject: Remove lr args from model, add Cosine lr, fix to vqvae stack --- notebooks/05c-test-model-end-to-end.ipynb | 119 +++++++++++++++++++++-- text_recognizer/models/base.py | 12 ++- text_recognizer/models/vqvae.py | 3 + training/conf/experiment/vqvae.yaml | 11 ++- training/conf/lr_scheduler/cosine_annealing.yaml | 7 ++ training/conf/lr_scheduler/one_cycle.yaml | 4 + training/conf/model/lit_vqvae.yaml | 4 +- training/conf/network/decoder/vae_decoder.yaml | 2 +- training/conf/network/encoder/vae_encoder.yaml | 2 +- training/conf/network/vqvae.yaml | 2 +- training/conf/optimizer/madgrad.yaml | 2 +- 11 files changed, 143 insertions(+), 25 deletions(-) create mode 100644 training/conf/lr_scheduler/cosine_annealing.yaml diff --git a/notebooks/05c-test-model-end-to-end.ipynb b/notebooks/05c-test-model-end-to-end.ipynb index 7996257..b26a1fe 100644 --- a/notebooks/05c-test-model-end-to-end.ipynb +++ b/notebooks/05c-test-model-end-to-end.ipynb @@ -2,10 +2,19 @@ "cells": [ { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "1e40a88b", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The autoreload extension is already loaded. To reload it, use:\n", + " %reload_ext autoreload\n" + ] + } + ], "source": [ "%load_ext autoreload\n", "%autoreload 2\n", @@ -25,7 +34,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "38fb3d9d-a163-4b72-981f-f31b51be39f2", "metadata": {}, "outputs": [], @@ -37,10 +46,46 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 10, "id": "74780b21-3313-452b-b580-703cac878416", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "encoder:\n", + " _target_: text_recognizer.networks.vqvae.encoder.Encoder\n", + " in_channels: 1\n", + " hidden_dim: 32\n", + " channels_multipliers:\n", + " - 1\n", + " - 2\n", + " - 4\n", + " - 4\n", + " - 4\n", + " dropout_rate: 0.25\n", + "decoder:\n", + " _target_: text_recognizer.networks.vqvae.decoder.Decoder\n", + " out_channels: 1\n", + " hidden_dim: 32\n", + " channels_multipliers:\n", + " - 4\n", + " - 4\n", + " - 4\n", + " - 2\n", + " - 1\n", + " dropout_rate: 0.25\n", + "_target_: text_recognizer.networks.vqvae.vqvae.VQVAE\n", + "hidden_dim: 128\n", + "embedding_dim: 32\n", + "num_embeddings: 1024\n", + "decay: 0.99\n", + "\n", + "{'encoder': {'_target_': 'text_recognizer.networks.vqvae.encoder.Encoder', 'in_channels': 1, 'hidden_dim': 32, 'channels_multipliers': [1, 2, 4, 4, 4], 'dropout_rate': 0.25}, 'decoder': {'_target_': 'text_recognizer.networks.vqvae.decoder.Decoder', 'out_channels': 1, 'hidden_dim': 32, 'channels_multipliers': [4, 4, 4, 2, 1], 'dropout_rate': 0.25}, '_target_': 'text_recognizer.networks.vqvae.vqvae.VQVAE', 'hidden_dim': 128, 'embedding_dim': 32, 'num_embeddings': 1024, 'decay': 0.99}\n" + ] + } + ], "source": [ "# context initialization\n", "with initialize(config_path=\"../training/conf/network/\", job_name=\"test_app\"):\n", @@ -51,7 +96,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 11, "id": "205a03e8-7aa1-407f-afa5-92693715b677", "metadata": {}, "outputs": [], @@ -61,7 +106,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 12, "id": "c74384f0-754e-4c29-8f06-339372d6e4c1", "metadata": {}, "outputs": [], @@ -71,10 +116,66 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 13, "id": "5ebab599-2497-42f8-b54b-1663ee66fde9", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "==========================================================================================\n", + "Layer (type:depth-idx) Output Shape Param #\n", + "==========================================================================================\n", + "├─Encoder: 1-1 [-1, 128, 18, 20] --\n", + "| └─Sequential: 2-1 [-1, 128, 18, 20] --\n", + "| | └─Conv2d: 3-1 [-1, 32, 576, 640] 320\n", + "| | └─Conv2d: 3-2 [-1, 32, 288, 320] 16,416\n", + "| | └─Mish: 3-3 [-1, 32, 288, 320] --\n", + "| | └─Conv2d: 3-4 [-1, 64, 144, 160] 32,832\n", + "| | └─Mish: 3-5 [-1, 64, 144, 160] --\n", + "| | └─Conv2d: 3-6 [-1, 128, 72, 80] 131,200\n", + "| | └─Mish: 3-7 [-1, 128, 72, 80] --\n", + "| | └─Conv2d: 3-8 [-1, 128, 36, 40] 262,272\n", + "| | └─Mish: 3-9 [-1, 128, 36, 40] --\n", + "| | └─Conv2d: 3-10 [-1, 128, 18, 20] 262,272\n", + "| | └─Mish: 3-11 [-1, 128, 18, 20] --\n", + "| | └─Residual: 3-12 [-1, 128, 18, 20] 164,352\n", + "| | └─Residual: 3-13 [-1, 128, 18, 20] 164,352\n", + "├─Conv2d: 1-2 [-1, 32, 18, 20] 4,128\n", + "├─VectorQuantizer: 1-3 [-1, 32, 18, 20] --\n", + "├─Conv2d: 1-4 [-1, 128, 18, 20] 4,224\n", + "├─Decoder: 1-5 [-1, 1, 576, 640] --\n", + "| └─Sequential: 2-2 [-1, 1, 576, 640] --\n", + "| | └─Residual: 3-14 [-1, 128, 18, 20] 164,352\n", + "| | └─Residual: 3-15 [-1, 128, 18, 20] 164,352\n", + "| | └─ConvTranspose2d: 3-16 [-1, 128, 36, 40] 262,272\n", + "| | └─Mish: 3-17 [-1, 128, 36, 40] --\n", + "| | └─ConvTranspose2d: 3-18 [-1, 128, 72, 80] 262,272\n", + "| | └─Mish: 3-19 [-1, 128, 72, 80] --\n", + "| | └─ConvTranspose2d: 3-20 [-1, 64, 144, 160] 131,136\n", + "| | └─Mish: 3-21 [-1, 64, 144, 160] --\n", + "| | └─ConvTranspose2d: 3-22 [-1, 32, 288, 320] 32,800\n", + "| | └─Mish: 3-23 [-1, 32, 288, 320] --\n", + "| | └─ConvTranspose2d: 3-24 [-1, 32, 576, 640] 16,416\n", + "| | └─Mish: 3-25 [-1, 32, 576, 640] --\n", + "| | └─Normalize: 3-26 [-1, 32, 576, 640] 64\n", + "| | └─Mish: 3-27 [-1, 32, 576, 640] --\n", + "| | └─Conv2d: 3-28 [-1, 1, 576, 640] 289\n", + "==========================================================================================\n", + "Total params: 2,076,321\n", + "Trainable params: 2,076,321\n", + "Non-trainable params: 0\n", + "Total mult-adds (G): 17.68\n", + "==========================================================================================\n", + "Input size (MB): 1.41\n", + "Forward/backward pass size (MB): 355.17\n", + "Params size (MB): 7.92\n", + "Estimated Total Size (MB): 364.49\n", + "==========================================================================================\n" + ] + } + ], "source": [ "summary(net, (1, 576, 640), device=\"cpu\");" ] diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py index 57c5964..ab3fa35 100644 --- a/text_recognizer/models/base.py +++ b/text_recognizer/models/base.py @@ -26,8 +26,6 @@ class BaseLitModel(LightningModule): loss_fn: Type[nn.Module] = attr.ib() optimizer_config: DictConfig = attr.ib() lr_scheduler_config: DictConfig = attr.ib() - interval: str = attr.ib() - monitor: str = attr.ib(default="val/loss") train_acc: torchmetrics.Accuracy = attr.ib( init=False, default=torchmetrics.Accuracy() ) @@ -58,12 +56,18 @@ class BaseLitModel(LightningModule): self, optimizer: Type[torch.optim.Optimizer] ) -> Dict[str, Any]: """Configures the lr scheduler.""" + # Extract non-class arguments. + monitor = self.lr_scheduler_config.monitor + interval = self.lr_scheduler_config.interval + del self.lr_scheduler_config.monitor + del self.lr_scheduler_config.interval + log.info( f"Instantiating learning rate scheduler <{self.lr_scheduler_config._target_}>" ) scheduler = { - "monitor": self.monitor, - "interval": self.interval, + "monitor": monitor, + "interval": interval, "scheduler": hydra.utils.instantiate( self.lr_scheduler_config, optimizer=optimizer ), diff --git a/text_recognizer/models/vqvae.py b/text_recognizer/models/vqvae.py index 7f79b78..76b7ba6 100644 --- a/text_recognizer/models/vqvae.py +++ b/text_recognizer/models/vqvae.py @@ -23,6 +23,7 @@ class VQVAELitModel(BaseLitModel): reconstructions, vq_loss = self(data) loss = self.loss_fn(reconstructions, data) loss = loss + self.latent_loss_weight * vq_loss + self.log("train/vq_loss", vq_loss) self.log("train/loss", loss) return loss @@ -32,6 +33,7 @@ class VQVAELitModel(BaseLitModel): reconstructions, vq_loss = self(data) loss = self.loss_fn(reconstructions, data) loss = loss + self.latent_loss_weight * vq_loss + self.log("val/vq_loss", vq_loss) self.log("val/loss", loss, prog_bar=True) def test_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None: @@ -40,4 +42,5 @@ class VQVAELitModel(BaseLitModel): reconstructions, vq_loss = self(data) loss = self.loss_fn(reconstructions, data) loss = loss + self.latent_loss_weight * vq_loss + self.log("test/vq_loss", vq_loss) self.log("test/loss", loss) 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 -- cgit v1.2.3-70-g09d2