Update conv transformer model

1 files changed, 70 insertions, 66 deletions

diff --git a/training/conf/experiment/conv_transformer_lines.yaml b/training/conf/experiment/conv_transformer_lines.yaml
index e0e426c..d32c7d6 100644
--- a/training/conf/experiment/conv_transformer_lines.yaml
+++ b/training/conf/experiment/conv_transformer_lines.yaml
@@ -12,7 +12,7 @@ defaults:
 tags: [lines]
 epochs: &epochs 260
 ignore_index: &ignore_index 3
-num_classes: &num_classes 57
+num_classes: &num_classes 58
 max_output_len: &max_output_len 89
 # summary: [[1, 1, 56, 1024], [1, 89]]
 
@@ -35,7 +35,7 @@ callbacks:
 
 optimizer:
   _target_: adan_pytorch.Adan
-  lr: 1.0e-3
+  lr: 3.0e-4
   betas: [0.02, 0.08, 0.01]
   weight_decay: 0.02
 
@@ -59,73 +59,77 @@ datamodule:
 
 network:
   _target_: text_recognizer.networks.ConvTransformer
-  input_dims: [1, 1, 56, 1024]
-  hidden_dim: &hidden_dim 384
-  num_classes: 58
-  pad_index: 3
   encoder:
-    _target_: text_recognizer.networks.convnext.ConvNext
-    dim: 16
-    dim_mults: [2, 4, 24]
-    depths: [3, 3, 6]
-    downsampling_factors: [[2, 2], [2, 2], [2, 2]]
-    attn:
-      _target_: text_recognizer.networks.convnext.TransformerBlock
+    _target_: text_recognizer.networks.image_encoder.ImageEncoder
+    encoder:
+      _target_: text_recognizer.networks.convnext.ConvNext
+      dim: 16
+      dim_mults: [2, 4, 24]
+      depths: [3, 3, 6]
+      downsampling_factors: [[2, 2], [2, 2], [2, 2]]
       attn:
-        _target_: text_recognizer.networks.convnext.Attention
-        dim: *hidden_dim
-        heads: 4
-        dim_head: 64
-        scale: 8
-      ff:
-        _target_: text_recognizer.networks.convnext.FeedForward
-        dim: *hidden_dim
-        mult: 2
+        _target_: text_recognizer.networks.convnext.TransformerBlock
+        attn:
+          _target_: text_recognizer.networks.convnext.Attention
+          dim: *hidden_dim
+          heads: 4
+          dim_head: 64
+          scale: 8
+        ff:
+          _target_: text_recognizer.networks.convnext.FeedForward
+          dim: *hidden_dim
+          mult: 2
+    pixel_embedding:
+      _target_: "text_recognizer.networks.transformer.embeddings.axial.\
+        AxialPositionalEmbeddingImage"
+      dim: &hidden_dim 384
+      axial_shape: [7, 128]
+      axial_dims: [192, 192]
   decoder:
-    _target_: text_recognizer.networks.transformer.Decoder
-    depth: 6
-    dim: *hidden_dim
-    block:
-      _target_: text_recognizer.networks.transformer.decoder_block.DecoderBlock
-      self_attn:
-        _target_: text_recognizer.networks.transformer.Attention
-        dim: *hidden_dim
-        num_heads: 8
-        dim_head: 64
-        dropout_rate: &dropout_rate 0.2
-        causal: true
-        rotary_embedding:
-          _target_: text_recognizer.networks.transformer.RotaryEmbedding
-          dim: 64
-      cross_attn:
-        _target_: text_recognizer.networks.transformer.Attention
-        dim: *hidden_dim
-        num_heads: 8
-        dim_head: 64
-        dropout_rate: *dropout_rate
-        causal: false
-      norm:
-        _target_: text_recognizer.networks.transformer.RMSNorm
-        dim: *hidden_dim
-      ff:
-        _target_: text_recognizer.networks.transformer.FeedForward
-        dim: *hidden_dim
-        dim_out: null
-        expansion_factor: 2
-        glu: true
-        dropout_rate: *dropout_rate
-  pixel_embedding:
-    _target_: "text_recognizer.networks.transformer.embeddings.axial.\
-      AxialPositionalEmbeddingImage"
-    dim: *hidden_dim
-    axial_shape: [7, 128]
-    axial_dims: [192, 192]
-  token_pos_embedding:
-    _target_: "text_recognizer.networks.transformer.embeddings.fourier.\
-      PositionalEncoding"
-    dim: *hidden_dim
-    dropout_rate: 0.1
-    max_len: 89
+    _target_: text_recognizer.networks.text_decoder.TextDecoder
+    hidden_dim: *hidden_dim
+    num_classes: *num_classes
+    pad_index: *ignore_index
+    decoder:
+      _target_: text_recognizer.networks.transformer.Decoder
+      dim: *hidden_dim
+      depth: 6
+      block:
+        _target_: text_recognizer.networks.transformer.decoder_block.\
+          DecoderBlock
+        self_attn:
+          _target_: text_recognizer.networks.transformer.Attention
+          dim: *hidden_dim
+          num_heads: 10
+          dim_head: 64
+          dropout_rate: &dropout_rate 0.2
+          causal: true
+        cross_attn:
+          _target_: text_recognizer.networks.transformer.Attention
+          dim: *hidden_dim
+          num_heads: 10
+          dim_head: 64
+          dropout_rate: *dropout_rate
+          causal: false
+        norm:
+          _target_: text_recognizer.networks.transformer.RMSNorm
+          dim: *hidden_dim
+        ff:
+          _target_: text_recognizer.networks.transformer.FeedForward
+          dim: *hidden_dim
+          dim_out: null
+          expansion_factor: 2
+          glu: true
+          dropout_rate: *dropout_rate
+      rotary_embedding:
+        _target_: text_recognizer.networks.transformer.RotaryEmbedding
+        dim: 64
+    token_pos_embedding:
+      _target_: "text_recognizer.networks.transformer.embeddings.fourier.\
+        PositionalEncoding"
+      dim: *hidden_dim
+      dropout_rate: 0.1
+      max_len: *max_output_len
 
 model:
   max_output_len: *max_output_len