diff options
| -rw-r--r-- | notebooks/05c-test-model-end-to-end.ipynb | 308 | ||||
| -rw-r--r-- | text_recognizer/criterions/vqgan_loss.py | 23 | ||||
| -rw-r--r-- | text_recognizer/models/base.py | 8 | ||||
| -rw-r--r-- | text_recognizer/models/vqgan.py | 24 | ||||
| -rw-r--r-- | training/conf/experiment/vqgan.yaml | 37 | ||||
| -rw-r--r-- | training/conf/model/lit_vqgan.yaml | 1 | ||||
| -rw-r--r-- | training/conf/network/encoder/vae_encoder.yaml | 2 |
7 files changed, 205 insertions, 198 deletions
diff --git a/notebooks/05c-test-model-end-to-end.ipynb b/notebooks/05c-test-model-end-to-end.ipynb index 42621da..23361b6 100644 --- a/notebooks/05c-test-model-end-to-end.ipynb +++ b/notebooks/05c-test-model-end-to-end.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": 32, "id": "1e40a88b", "metadata": {}, "outputs": [], @@ -25,53 +25,7 @@ }, { "cell_type": "code", - "execution_count": 2, - "id": "f40fc669-829c-4de8-83ed-475fc6a0b8c1", - "metadata": {}, - "outputs": [], - "source": [ - "class T:\n", - " def __init__(self):\n", - " self.network = nn.Linear(1, 1)\n", - " \n", - " def get(self):\n", - " return getattr(self, \"network\")" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "d2bedf96-5388-4c7a-a048-1b97041cbedc", - "metadata": {}, - "outputs": [], - "source": [ - "t = T()" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "a6fbe3be-2a9f-4050-a397-7ad982d6cd05", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "<generator object Module.parameters at 0x7f29ad6d6120>" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "t.get().parameters()" - ] - }, - { - "cell_type": "code", - "execution_count": 5, + "execution_count": 33, "id": "38fb3d9d-a163-4b72-981f-f31b51be39f2", "metadata": {}, "outputs": [], @@ -83,7 +37,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 47, "id": "74780b21-3313-452b-b580-703cac878416", "metadata": {}, "outputs": [ @@ -91,49 +45,181 @@ "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", + "callbacks:\n", + " model_checkpoint:\n", + " _target_: pytorch_lightning.callbacks.ModelCheckpoint\n", + " monitor: val/loss\n", + " save_top_k: 1\n", + " save_last: true\n", + " mode: min\n", + " verbose: false\n", + " dirpath: checkpoints/\n", + " filename: '{epoch:02d}'\n", + " learning_rate_monitor:\n", + " _target_: pytorch_lightning.callbacks.LearningRateMonitor\n", + " logging_interval: step\n", + " log_momentum: false\n", + " watch_model:\n", + " _target_: callbacks.wandb_callbacks.WatchModel\n", + " log: all\n", + " log_freq: 100\n", + " upload_ckpts_as_artifact:\n", + " _target_: callbacks.wandb_callbacks.UploadCheckpointsAsArtifact\n", + " ckpt_dir: checkpoints/\n", + " upload_best_only: true\n", + " log_image_reconstruction:\n", + " _target_: callbacks.wandb_callbacks.LogReconstuctedImages\n", + " num_samples: 8\n", + "criterion:\n", + " _target_: text_recognizer.criterions.vqgan_loss.VQGANLoss\n", + " reconstruction_loss:\n", + " _target_: torch.nn.L1Loss\n", + " reduction: mean\n", + " discriminator:\n", + " _target_: text_recognizer.criterions.n_layer_discriminator.NLayerDiscriminator\n", + " in_channels: 1\n", + " num_channels: 32\n", + " num_layers: 3\n", + " vq_loss_weight: 1.0\n", + " discriminator_weight: 1.0\n", + "datamodule:\n", + " _target_: text_recognizer.data.iam_extended_paragraphs.IAMExtendedParagraphs\n", + " batch_size: 8\n", + " num_workers: 12\n", + " train_fraction: 0.8\n", + " augment: true\n", + " pin_memory: false\n", + " word_pieces: true\n", + "logger:\n", + " wandb:\n", + " _target_: pytorch_lightning.loggers.wandb.WandbLogger\n", + " project: text-recognizer\n", + " name: null\n", + " save_dir: .\n", + " offline: false\n", + " id: null\n", + " log_model: false\n", + " prefix: ''\n", + " job_type: train\n", + " group: ''\n", + " tags: []\n", + "mapping:\n", + " _target_: text_recognizer.data.word_piece_mapping.WordPieceMapping\n", + " num_features: 1000\n", + " tokens: iamdb_1kwp_tokens_1000.txt\n", + " lexicon: iamdb_1kwp_lex_1000.txt\n", + " data_dir: null\n", + " use_words: false\n", + " prepend_wordsep: false\n", + " special_tokens:\n", + " - <s>\n", + " - <e>\n", + " - <p>\n", + " extra_symbols:\n", + " - '\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" + " '\n", + "model:\n", + " _target_: text_recognizer.models.vqgan.VQGANLitModel\n", + "network:\n", + " 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", + " - 8\n", + " - 8\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", + " - 8\n", + " - 8\n", + " - 4\n", + " - 1\n", + " dropout_rate: 0.25\n", + " _target_: text_recognizer.networks.vqvae.vqvae.VQVAE\n", + " hidden_dim: 256\n", + " embedding_dim: 32\n", + " num_embeddings: 1024\n", + " decay: 0.99\n", + "trainer:\n", + " _target_: pytorch_lightning.Trainer\n", + " stochastic_weight_avg: false\n", + " auto_scale_batch_size: binsearch\n", + " auto_lr_find: false\n", + " gradient_clip_val: 0\n", + " fast_dev_run: false\n", + " gpus: 1\n", + " precision: 16\n", + " max_epochs: 256\n", + " terminate_on_nan: true\n", + " weights_summary: top\n", + " limit_train_batches: 1.0\n", + " limit_val_batches: 1.0\n", + " limit_test_batches: 1.0\n", + " resume_from_checkpoint: null\n", + "seed: 4711\n", + "tune: false\n", + "train: true\n", + "test: true\n", + "logging: INFO\n", + "work_dir: ${hydra:runtime.cwd}\n", + "debug: false\n", + "print_config: false\n", + "ignore_warnings: true\n", + "summary: null\n", + "lr_schedulers:\n", + " generator:\n", + " _target_: torch.optim.lr_scheduler.CosineAnnealingLR\n", + " T_max: 256\n", + " eta_min: 0.0\n", + " last_epoch: -1\n", + " interval: epoch\n", + " monitor: val/loss\n", + " discriminator:\n", + " _target_: torch.optim.lr_scheduler.CosineAnnealingLR\n", + " T_max: 256\n", + " eta_min: 0.0\n", + " last_epoch: -1\n", + " interval: epoch\n", + " monitor: val/loss\n", + "optimizers:\n", + " generator:\n", + " _target_: madgrad.MADGRAD\n", + " lr: 0.001\n", + " momentum: 0.5\n", + " weight_decay: 0\n", + " eps: 1.0e-06\n", + " parameters: network\n", + " discriminator:\n", + " _target_: madgrad.MADGRAD\n", + " lr: 0.001\n", + " momentum: 0.5\n", + " weight_decay: 0\n", + " eps: 1.0e-06\n", + " parameters: loss_fn.discriminator\n", + "\n", + "{'callbacks': {'model_checkpoint': {'_target_': 'pytorch_lightning.callbacks.ModelCheckpoint', 'monitor': 'val/loss', 'save_top_k': 1, 'save_last': True, 'mode': 'min', 'verbose': False, 'dirpath': 'checkpoints/', 'filename': '{epoch:02d}'}, 'learning_rate_monitor': {'_target_': 'pytorch_lightning.callbacks.LearningRateMonitor', 'logging_interval': 'step', 'log_momentum': False}, 'watch_model': {'_target_': 'callbacks.wandb_callbacks.WatchModel', 'log': 'all', 'log_freq': 100}, 'upload_ckpts_as_artifact': {'_target_': 'callbacks.wandb_callbacks.UploadCheckpointsAsArtifact', 'ckpt_dir': 'checkpoints/', 'upload_best_only': True}, 'log_image_reconstruction': {'_target_': 'callbacks.wandb_callbacks.LogReconstuctedImages', 'num_samples': 8}}, 'criterion': {'_target_': 'text_recognizer.criterions.vqgan_loss.VQGANLoss', 'reconstruction_loss': {'_target_': 'torch.nn.L1Loss', 'reduction': 'mean'}, 'discriminator': {'_target_': 'text_recognizer.criterions.n_layer_discriminator.NLayerDiscriminator', 'in_channels': 1, 'num_channels': 32, 'num_layers': 3}, 'vq_loss_weight': 1.0, 'discriminator_weight': 1.0}, 'datamodule': {'_target_': 'text_recognizer.data.iam_extended_paragraphs.IAMExtendedParagraphs', 'batch_size': 8, 'num_workers': 12, 'train_fraction': 0.8, 'augment': True, 'pin_memory': False, 'word_pieces': True}, 'logger': {'wandb': {'_target_': 'pytorch_lightning.loggers.wandb.WandbLogger', 'project': 'text-recognizer', 'name': None, 'save_dir': '.', 'offline': False, 'id': None, 'log_model': False, 'prefix': '', 'job_type': 'train', 'group': '', 'tags': []}}, 'mapping': {'_target_': 'text_recognizer.data.word_piece_mapping.WordPieceMapping', 'num_features': 1000, 'tokens': 'iamdb_1kwp_tokens_1000.txt', 'lexicon': 'iamdb_1kwp_lex_1000.txt', 'data_dir': None, 'use_words': False, 'prepend_wordsep': False, 'special_tokens': ['<s>', '<e>', '<p>'], 'extra_symbols': ['\\n']}, 'model': {'_target_': 'text_recognizer.models.vqgan.VQGANLitModel'}, 'network': {'encoder': {'_target_': 'text_recognizer.networks.vqvae.encoder.Encoder', 'in_channels': 1, 'hidden_dim': 32, 'channels_multipliers': [1, 2, 4, 8, 8], 'dropout_rate': 0.25}, 'decoder': {'_target_': 'text_recognizer.networks.vqvae.decoder.Decoder', 'out_channels': 1, 'hidden_dim': 32, 'channels_multipliers': [8, 8, 4, 1], 'dropout_rate': 0.25}, '_target_': 'text_recognizer.networks.vqvae.vqvae.VQVAE', 'hidden_dim': 256, 'embedding_dim': 32, 'num_embeddings': 1024, 'decay': 0.99}, 'trainer': {'_target_': 'pytorch_lightning.Trainer', 'stochastic_weight_avg': False, 'auto_scale_batch_size': 'binsearch', 'auto_lr_find': False, 'gradient_clip_val': 0, 'fast_dev_run': False, 'gpus': 1, 'precision': 16, 'max_epochs': 256, 'terminate_on_nan': True, 'weights_summary': 'top', 'limit_train_batches': 1.0, 'limit_val_batches': 1.0, 'limit_test_batches': 1.0, 'resume_from_checkpoint': None}, 'seed': 4711, 'tune': False, 'train': True, 'test': True, 'logging': 'INFO', 'work_dir': '${hydra:runtime.cwd}', 'debug': False, 'print_config': False, 'ignore_warnings': True, 'summary': None, 'lr_schedulers': {'generator': {'_target_': 'torch.optim.lr_scheduler.CosineAnnealingLR', 'T_max': 256, 'eta_min': 0.0, 'last_epoch': -1, 'interval': 'epoch', 'monitor': 'val/loss'}, 'discriminator': {'_target_': 'torch.optim.lr_scheduler.CosineAnnealingLR', 'T_max': 256, 'eta_min': 0.0, 'last_epoch': -1, 'interval': 'epoch', 'monitor': 'val/loss'}}, 'optimizers': {'generator': {'_target_': 'madgrad.MADGRAD', 'lr': 0.001, 'momentum': 0.5, 'weight_decay': 0, 'eps': 1e-06, 'parameters': 'network'}, 'discriminator': {'_target_': 'madgrad.MADGRAD', 'lr': 0.001, 'momentum': 0.5, 'weight_decay': 0, 'eps': 1e-06, 'parameters': 'loss_fn.discriminator'}}}\n" ] } ], "source": [ "# context initialization\n", - "with initialize(config_path=\"../training/conf/network/\", job_name=\"test_app\"):\n", - " cfg = compose(config_name=\"vqvae\")\n", + "with initialize(config_path=\"../training/conf/\", job_name=\"test_app\"):\n", + " cfg = compose(config_name=\"config\", overrides=[\"+experiment=vqgan\"])\n", " print(OmegaConf.to_yaml(cfg))\n", " print(cfg)" ] }, { "cell_type": "code", - "execution_count": 11, + "execution_count": null, "id": "205a03e8-7aa1-407f-afa5-92693715b677", "metadata": {}, "outputs": [], @@ -143,7 +229,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": null, "id": "c74384f0-754e-4c29-8f06-339372d6e4c1", "metadata": {}, "outputs": [], @@ -153,66 +239,10 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": null, "id": "5ebab599-2497-42f8-b54b-1663ee66fde9", "metadata": {}, - "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" - ] - } - ], + "outputs": [], "source": [ "summary(net, (1, 576, 640), device=\"cpu\");" ] diff --git a/text_recognizer/criterions/vqgan_loss.py b/text_recognizer/criterions/vqgan_loss.py index 8bb568f..87f0f1c 100644 --- a/text_recognizer/criterions/vqgan_loss.py +++ b/text_recognizer/criterions/vqgan_loss.py @@ -1,5 +1,5 @@ """VQGAN loss for PyTorch Lightning.""" -from typing import Dict +from typing import Dict, Optional from click.types import Tuple import torch @@ -40,9 +40,9 @@ class VQGANLoss(nn.Module): vq_loss: Tensor, optimizer_idx: int, stage: str, - ) -> Tuple[Tensor, Dict[str, Tensor]]: + ) -> Optional[Tuple]: """Calculates the VQGAN loss.""" - rec_loss = self.reconstruction_loss( + rec_loss: Tensor = self.reconstruction_loss( data.contiguous(), reconstructions.contiguous() ) @@ -51,13 +51,13 @@ class VQGANLoss(nn.Module): logits_fake = self.discriminator(reconstructions.contiguous()) g_loss = -torch.mean(logits_fake) - loss = ( + loss: Tensor = ( rec_loss + self.discriminator_weight * g_loss + self.vq_loss_weight * vq_loss ) log = { - f"{stage}/loss": loss, + f"{stage}/total_loss": loss, f"{stage}/vq_loss": vq_loss, f"{stage}/rec_loss": rec_loss, f"{stage}/g_loss": g_loss, @@ -68,18 +68,11 @@ class VQGANLoss(nn.Module): logits_fake = self.discriminator(reconstructions.contiguous().detach()) logits_real = self.discriminator(data.contiguous().detach()) - d_loss = self.adversarial_loss( + d_loss = self.discriminator_weight * self.adversarial_loss( logits_real=logits_real, logits_fake=logits_fake ) - loss = ( - rec_loss - + self.discriminator_weight * d_loss - + self.vq_loss_weight * vq_loss - ) + log = { - f"{stage}/loss": loss, - f"{stage}/vq_loss": vq_loss, - f"{stage}/rec_loss": rec_loss, f"{stage}/d_loss": d_loss, } - return loss, log + return d_loss, log diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py index 8b68ed9..94dbde5 100644 --- a/text_recognizer/models/base.py +++ b/text_recognizer/models/base.py @@ -49,12 +49,14 @@ class BaseLitModel(LightningModule): """Configures the optimizer.""" optimizers = [] for optimizer_config in self.optimizer_configs.values(): - network = getattr(self, optimizer_config.parameters) + module = self + for m in str(optimizer_config.parameters).split("."): + module = getattr(module, m) del optimizer_config.parameters log.info(f"Instantiating optimizer <{optimizer_config._target_}>") optimizers.append( hydra.utils.instantiate( - self.optimizer_config, params=network.parameters() + optimizer_config, params=module.parameters() ) ) return optimizers @@ -92,7 +94,7 @@ class BaseLitModel(LightningModule): ) -> Tuple[List[Type[torch.optim.Optimizer]], List[Dict[str, Any]]]: """Configures optimizer and lr scheduler.""" optimizers = self._configure_optimizer() - schedulers = self._configure_lr_scheduler(optimizers) + schedulers = self._configure_lr_schedulers(optimizers) return optimizers, schedulers def forward(self, data: Tensor) -> Tensor: diff --git a/text_recognizer/models/vqgan.py b/text_recognizer/models/vqgan.py index 8ff65cc..80653b6 100644 --- a/text_recognizer/models/vqgan.py +++ b/text_recognizer/models/vqgan.py @@ -9,7 +9,7 @@ from text_recognizer.criterions.vqgan_loss import VQGANLoss @attr.s(auto_attribs=True, eq=False) -class VQVAELitModel(BaseLitModel): +class VQGANLitModel(BaseLitModel): """A PyTorch Lightning model for transformer networks.""" loss_fn: VQGANLoss = attr.ib() @@ -26,7 +26,6 @@ class VQVAELitModel(BaseLitModel): data, _ = batch reconstructions, vq_loss = self(data) - loss = self.loss_fn(reconstructions, data) if optimizer_idx == 0: loss, log = self.loss_fn( @@ -81,14 +80,6 @@ class VQVAELitModel(BaseLitModel): self.log( "val/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True ) - self.log( - "val/rec_loss", - log["val/rec_loss"], - prog_bar=True, - logger=True, - on_step=True, - on_epoch=True, - ) self.log_dict(log) _, log = self.loss_fn( @@ -105,24 +96,13 @@ class VQVAELitModel(BaseLitModel): data, _ = batch reconstructions, vq_loss = self(data) - loss, log = self.loss_fn( + _, log = self.loss_fn( data=data, reconstructions=reconstructions, vq_loss=vq_loss, optimizer_idx=0, stage="test", ) - self.log( - "test/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True - ) - self.log( - "test/rec_loss", - log["test/rec_loss"], - prog_bar=True, - logger=True, - on_step=True, - on_epoch=True, - ) self.log_dict(log) _, log = self.loss_fn( diff --git a/training/conf/experiment/vqgan.yaml b/training/conf/experiment/vqgan.yaml index 3d97892..570e7f9 100644 --- a/training/conf/experiment/vqgan.yaml +++ b/training/conf/experiment/vqgan.yaml @@ -5,13 +5,15 @@ defaults: - override /criterion: vqgan_loss - override /model: lit_vqgan - override /callbacks: wandb_vae + - override /optimizers: null - override /lr_schedulers: null datamodule: batch_size: 8 lr_schedulers: - - generator: + generator: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR T_max: 256 eta_min: 0.0 last_epoch: -1 @@ -19,7 +21,8 @@ lr_schedulers: interval: epoch monitor: val/loss - - discriminator: + discriminator: + _target_: torch.optim.lr_scheduler.CosineAnnealingLR T_max: 256 eta_min: 0.0 last_epoch: -1 @@ -27,26 +30,24 @@ lr_schedulers: interval: epoch monitor: val/loss -optimizer: - - generator: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - T_max: 256 - eta_min: 0.0 - last_epoch: -1 +optimizers: + generator: + _target_: madgrad.MADGRAD + lr: 2.0e-5 + momentum: 0.5 + weight_decay: 0 + eps: 1.0e-6 - interval: epoch - monitor: val/loss parameters: network - - discriminator: - _target_: torch.optim.lr_scheduler.CosineAnnealingLR - T_max: 256 - eta_min: 0.0 - last_epoch: -1 + discriminator: + _target_: madgrad.MADGRAD + lr: 2.0e-5 + momentum: 0.5 + weight_decay: 0 + eps: 1.0e-6 - interval: epoch - monitor: val/loss - parameters: loss_fn + parameters: loss_fn.discriminator trainer: max_epochs: 256 diff --git a/training/conf/model/lit_vqgan.yaml b/training/conf/model/lit_vqgan.yaml new file mode 100644 index 0000000..9ee1046 --- /dev/null +++ b/training/conf/model/lit_vqgan.yaml @@ -0,0 +1 @@ +_target_: text_recognizer.models.vqgan.VQGANLitModel diff --git a/training/conf/network/encoder/vae_encoder.yaml b/training/conf/network/encoder/vae_encoder.yaml index 58e905d..099c36a 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, 4, 8, 8] +channels_multipliers: [1, 4, 8, 8] dropout_rate: 0.25 |