summaryrefslogtreecommitdiff
path: root/training/conf/experiment/conv_transformer_paragraphs.yaml
blob: e958367303c75f796279152cb35f23f215273d6a (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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
# @package _global_

defaults:
  - override /mapping: null
  - override /criterion: null
  - override /callbacks: htr
  - override /datamodule: iam_extended_paragraphs
  - override /network: null
  - override /model: null
  - override /lr_schedulers: null
  - override /optimizers: null


epochs: &epochs 720
ignore_index: &ignore_index 3
num_classes: &num_classes 58
max_output_len: &max_output_len 682
summary: [[1, 1, 576, 640], [1, 682]]

criterion:
  _target_: torch.nn.CrossEntropyLoss
  ignore_index: *ignore_index
    
mapping: &mapping
  mapping:
    _target_: text_recognizer.data.mappings.emnist.EmnistMapping
    extra_symbols: [ "\n" ]

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: 1.0e-4
    momentum: 0.9
    weight_decay: 5.0e-6
    eps: 1.0e-6
    parameters: network

lr_schedulers:
  network:
    _target_: torch.optim.lr_scheduler.OneCycleLR
    max_lr: 1.0e-4
    total_steps: null
    epochs: *epochs
    steps_per_epoch: 722
    pct_start: 0.01
    anneal_strategy: cos
    cycle_momentum: true
    base_momentum: 0.85
    max_momentum: 0.95
    div_factor: 25
    final_div_factor: 1.0e2
    three_phase: false
    last_epoch: -1
    verbose: false
    interval: step
    monitor: val/loss

datamodule:
  _target_: text_recognizer.data.iam_extended_paragraphs.IAMExtendedParagraphs
  batch_size: 6
  num_workers: 12
  train_fraction: 0.8
  pin_memory: true
  << : *mapping

rotary_embedding: &rotary_embedding
  rotary_embedding: 
    _target_: text_recognizer.networks.transformer.embeddings.rotary.RotaryEmbedding
    dim: 64

attn: &attn
  dim: 192
  num_heads: 4
  dim_head: 64
  dropout_rate: 0.05

network:
  _target_: text_recognizer.networks.conv_transformer.ConvTransformer
  input_dims: [1, 576, 640]
  hidden_dim: &hidden_dim 192
  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:
    depth: 3
    local_depth: 2
    _target_: text_recognizer.networks.transformer.layers.Decoder
    self_attn:
      _target_: text_recognizer.networks.transformer.attention.Attention
      << : *attn
      causal: true
      << : *rotary_embedding
    cross_attn:
      _target_: text_recognizer.networks.transformer.attention.Attention
      << : *attn
      causal: false
    local_self_attn:
      _target_: text_recognizer.networks.transformer.local_attention.LocalAttention
      << : *attn
      window_size: 31
      look_back: 1
      autopad: true
      << : *rotary_embedding
    norm:
      _target_: text_recognizer.networks.transformer.norm.ScaleNorm
      normalized_shape: *hidden_dim
    ff: 
      _target_: text_recognizer.networks.transformer.mlp.FeedForward
      dim: *hidden_dim
      dim_out: null
      expansion_factor: 4
      glu: true
      dropout_rate: 0.05
    pre_norm: true
  pixel_pos_embedding:
    _target_: text_recognizer.networks.transformer.embeddings.axial.AxialPositionalEmbedding
    dim: *hidden_dim
    shape: [18, 20]

model:
  _target_: text_recognizer.models.transformer.TransformerLitModel
  << : *mapping
  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.5
  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: 7
  overfit_batches: 0