summaryrefslogtreecommitdiff
path: root/training/configs/vqvae.yaml
blob: 13d7c97a03366fae3e30aebc41835934cde70371 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
seed: 4711

network:
        desc: Configuration of the PyTorch neural network.
        type: VQVAE
        args:
            in_channels: 1
            channels: [32, 64, 64, 96, 96]
            kernel_sizes: [4, 4, 4, 4, 4]
            strides: [2, 2, 2, 2, 2]
            num_residual_layers: 2 
            embedding_dim: 512
            num_embeddings: 1024
            upsampling: null
            beta: 0.25
            activation: leaky_relu
            dropout_rate: 0.2

model:
        desc: Configuration of the PyTorch Lightning model.
        type: LitVQVAEModel
        args:
                optimizer:
                        type: MADGRAD
                        args:
                                lr: 1.0e-3
                                momentum: 0.9
                                weight_decay: 0
                                eps: 1.0e-6
                lr_scheduler:
                        type: OneCycleLR
                        args:
                                interval: &interval step
                                max_lr: 1.0e-3
                                three_phase: true
                                epochs: 64
                                steps_per_epoch: 633 # num_samples / batch_size
                criterion:
                        type: MSELoss
                        args:
                                reduction: mean
                monitor: val_loss
                mapping: sentence_piece

data:
        desc: Configuration of the training/test data.
        type: IAMExtendedParagraphs
        args:
                batch_size: 32
                num_workers: 12
                train_fraction: 0.8
                augment: true

callbacks:
        - type: ModelCheckpoint
          args:
                  monitor: val_loss
                  mode: min
                  save_last: true
        - type: StochasticWeightAveraging
          args:
                  swa_epoch_start: 0.8
                  swa_lrs: 0.05
                  annealing_epochs: 10
                  annealing_strategy: cos
                  device: null
        - type: LearningRateMonitor
          args:
                  logging_interval: *interval
        # - type: EarlyStopping
        #   args:
        #           monitor: val_loss
        #           mode: min
        #           patience: 10

trainer:
        desc: Configuration of the PyTorch Lightning Trainer.
        args:
                stochastic_weight_avg: true
                auto_scale_batch_size: binsearch
                gradient_clip_val: 0
                fast_dev_run: false
                gpus: 1
                precision: 16
                max_epochs: 64
                terminate_on_nan: true
                weights_summary: top

load_checkpoint: null