summaryrefslogtreecommitdiff
path: root/notebooks/04-conv-transformer.ipynb
blob: 3303b63816033080d20afd0b4a2469a6b5ff7b3c (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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c02ae76-b540-4b16-9492-e9210b3b9249",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICE'] = ''\n",
    "import random\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "from omegaconf import OmegaConf\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from importlib.util import find_spec\n",
    "if find_spec(\"text_recognizer\") is None:\n",
    "    import sys\n",
    "    sys.path.append('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ccdb6dde-47e5-429a-88f2-0764fb7e259a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hydra import compose, initialize\n",
    "from omegaconf import OmegaConf\n",
    "from hydra.utils import instantiate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3cf50475-39f2-4642-a7d1-5bcbc0a036f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"../training/conf/network/conv_transformer.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e52ecb01-c975-4e55-925d-1182c7aea473",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(path, \"rb\") as f:\n",
    "    cfg = OmegaConf.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f939aa37-7b1d-45cc-885c-323c4540bda1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_target_': 'text_recognizer.networks.ConvTransformer', 'input_dims': [1, 1, 576, 640], 'hidden_dim': 128, 'num_classes': 58, 'pad_index': 3, 'encoder': {'_target_': 'text_recognizer.networks.EfficientNet', 'arch': 'b0', 'stochastic_dropout_rate': 0.2, 'bn_momentum': 0.99, 'bn_eps': 0.001, 'depth': 5, 'out_channels': 128, 'stride': [2, 1]}, 'decoder': {'_target_': 'text_recognizer.networks.transformer.Decoder', 'depth': 6, 'block': {'_target_': 'text_recognizer.networks.transformer.DecoderBlock', 'self_attn': {'_target_': 'text_recognizer.networks.transformer.Attention', 'dim': 128, 'num_heads': 8, 'dim_head': 64, 'dropout_rate': 0.4, 'causal': True, 'rotary_embedding': {'_target_': 'text_recognizer.networks.transformer.RotaryEmbedding', 'dim': 64}}, 'cross_attn': {'_target_': 'text_recognizer.networks.transformer.Attention', 'dim': 128, 'num_heads': 8, 'dim_head': 64, 'dropout_rate': 0.4, 'causal': False}, 'norm': {'_target_': 'text_recognizer.networks.transformer.RMSNorm', 'dim': 128}, 'ff': {'_target_': 'text_recognizer.networks.transformer.FeedForward', 'dim': 128, 'dim_out': None, 'expansion_factor': 2, 'glu': True, 'dropout_rate': 0.4}}}, 'pixel_embedding': {'_target_': 'text_recognizer.networks.transformer.AxialPositionalEmbedding', 'dim': 128, 'shape': [18, 80]}}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "aaeab329-aeb0-4a1b-aa35-5a2aab81b1d0",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "net = instantiate(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "618b997c-e6a6-4487-b70c-9d260cb556d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchinfo import summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "25759b7b-8deb-4163-b75d-a1357c9fe88f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==============================================================================================================\n",
       "Layer (type:depth-idx)                                       Output Shape              Param #\n",
       "==============================================================================================================\n",
       "ConvTransformer                                              [1, 58, 682]              --\n",
       "├─EfficientNet: 1-1                                          [1, 128, 18, 80]          850,044\n",
       "│    └─Sequential: 2-1                                       [1, 16, 288, 320]         --\n",
       "│    │    └─ZeroPad2d: 3-1                                   [1, 1, 577, 641]          --\n",
       "│    │    └─Conv2d: 3-2                                      [1, 16, 288, 320]         144\n",
       "│    │    └─BatchNorm2d: 3-3                                 [1, 16, 288, 320]         32\n",
       "│    │    └─Mish: 3-4                                        [1, 16, 288, 320]         --\n",
       "│    └─ModuleList: 2                                         --                        --\n",
       "│    │    └─MBConvBlock: 3-5                                 [1, 16, 288, 320]         --\n",
       "│    │    │    └─Depthwise: 4-1                              [1, 16, 288, 320]         176\n",
       "│    │    │    └─SqueezeAndExcite: 4-2                       [1, 16, 288, 320]         148\n",
       "│    │    │    └─Pointwise: 4-3                              [1, 16, 288, 320]         288\n",
       "│    │    └─MBConvBlock: 3-6                                 [1, 24, 144, 160]         --\n",
       "│    │    │    └─InvertedBottleneck: 4-4                     [1, 96, 288, 320]         1,728\n",
       "│    │    │    └─Depthwise: 4-5                              [1, 96, 144, 160]         1,056\n",
       "│    │    │    └─SqueezeAndExcite: 4-6                       [1, 96, 144, 160]         868\n",
       "│    │    │    └─Pointwise: 4-7                              [1, 24, 144, 160]         2,352\n",
       "│    │    └─MBConvBlock: 3-7                                 [1, 24, 144, 160]         --\n",
       "│    │    │    └─InvertedBottleneck: 4-8                     [1, 144, 144, 160]        3,744\n",
       "│    │    │    └─Depthwise: 4-9                              [1, 144, 144, 160]        1,584\n",
       "│    │    │    └─SqueezeAndExcite: 4-10                      [1, 144, 144, 160]        1,878\n",
       "│    │    │    └─Pointwise: 4-11                             [1, 24, 144, 160]         3,504\n",
       "│    │    └─MBConvBlock: 3-8                                 [1, 40, 72, 80]           --\n",
       "│    │    │    └─InvertedBottleneck: 4-12                    [1, 144, 144, 160]        3,744\n",
       "│    │    │    └─Depthwise: 4-13                             [1, 144, 72, 80]          3,888\n",
       "│    │    │    └─SqueezeAndExcite: 4-14                      [1, 144, 72, 80]          1,878\n",
       "│    │    │    └─Pointwise: 4-15                             [1, 40, 72, 80]           5,840\n",
       "│    │    └─MBConvBlock: 3-9                                 [1, 40, 72, 80]           --\n",
       "│    │    │    └─InvertedBottleneck: 4-16                    [1, 240, 72, 80]          10,080\n",
       "│    │    │    └─Depthwise: 4-17                             [1, 240, 72, 80]          6,480\n",
       "│    │    │    └─SqueezeAndExcite: 4-18                      [1, 240, 72, 80]          5,050\n",
       "│    │    │    └─Pointwise: 4-19                             [1, 40, 72, 80]           9,680\n",
       "│    │    └─MBConvBlock: 3-10                                [1, 80, 36, 80]           --\n",
       "│    │    │    └─InvertedBottleneck: 4-20                    [1, 240, 72, 80]          10,080\n",
       "│    │    │    └─Depthwise: 4-21                             [1, 240, 36, 80]          2,640\n",
       "│    │    │    └─SqueezeAndExcite: 4-22                      [1, 240, 36, 80]          5,050\n",
       "│    │    │    └─Pointwise: 4-23                             [1, 80, 36, 80]           19,360\n",
       "│    │    └─MBConvBlock: 3-11                                [1, 80, 36, 80]           --\n",
       "│    │    │    └─InvertedBottleneck: 4-24                    [1, 480, 36, 80]          39,360\n",
       "│    │    │    └─Depthwise: 4-25                             [1, 480, 36, 80]          5,280\n",
       "│    │    │    └─SqueezeAndExcite: 4-26                      [1, 480, 36, 80]          19,700\n",
       "│    │    │    └─Pointwise: 4-27                             [1, 80, 36, 80]           38,560\n",
       "│    │    └─MBConvBlock: 3-12                                [1, 80, 36, 80]           --\n",
       "│    │    │    └─InvertedBottleneck: 4-28                    [1, 480, 36, 80]          39,360\n",
       "│    │    │    └─Depthwise: 4-29                             [1, 480, 36, 80]          5,280\n",
       "│    │    │    └─SqueezeAndExcite: 4-30                      [1, 480, 36, 80]          19,700\n",
       "│    │    │    └─Pointwise: 4-31                             [1, 80, 36, 80]           38,560\n",
       "│    │    └─MBConvBlock: 3-13                                [1, 112, 18, 80]          --\n",
       "│    │    │    └─InvertedBottleneck: 4-32                    [1, 480, 36, 80]          39,360\n",
       "│    │    │    └─Depthwise: 4-33                             [1, 480, 18, 80]          12,960\n",
       "│    │    │    └─SqueezeAndExcite: 4-34                      [1, 480, 18, 80]          19,700\n",
       "│    │    │    └─Pointwise: 4-35                             [1, 112, 18, 80]          53,984\n",
       "│    │    └─MBConvBlock: 3-14                                [1, 112, 18, 80]          --\n",
       "│    │    │    └─InvertedBottleneck: 4-36                    [1, 672, 18, 80]          76,608\n",
       "│    │    │    └─Depthwise: 4-37                             [1, 672, 18, 80]          18,144\n",
       "│    │    │    └─SqueezeAndExcite: 4-38                      [1, 672, 18, 80]          38,332\n",
       "│    │    │    └─Pointwise: 4-39                             [1, 112, 18, 80]          75,488\n",
       "│    │    └─MBConvBlock: 3-15                                [1, 112, 18, 80]          --\n",
       "│    │    │    └─InvertedBottleneck: 4-40                    [1, 672, 18, 80]          76,608\n",
       "│    │    │    └─Depthwise: 4-41                             [1, 672, 18, 80]          18,144\n",
       "│    │    │    └─SqueezeAndExcite: 4-42                      [1, 672, 18, 80]          38,332\n",
       "│    │    │    └─Pointwise: 4-43                             [1, 112, 18, 80]          75,488\n",
       "│    └─Sequential: 2-2                                       [1, 128, 18, 80]          --\n",
       "│    │    └─Conv2d: 3-16                                     [1, 128, 18, 80]          14,336\n",
       "│    │    └─BatchNorm2d: 3-17                                [1, 128, 18, 80]          256\n",
       "│    │    └─Dropout: 3-18                                    [1, 128, 18, 80]          --\n",
       "├─Conv2d: 1-2                                                [1, 128, 18, 80]          16,512\n",
       "├─AxialPositionalEmbedding: 1-3                              [1, 128, 18, 80]          12,544\n",
       "├─Embedding: 1-4                                             [1, 682, 128]             7,424\n",
       "├─Decoder: 1-5                                               [1, 682, 128]             --\n",
       "│    └─ModuleList: 2                                         --                        --\n",
       "│    │    └─DecoderBlock: 3-19                               [1, 682, 128]             --\n",
       "│    │    └─DecoderBlock: 3-20                               [1, 682, 128]             --\n",
       "│    │    └─DecoderBlock: 3-21                               [1, 682, 128]             --\n",
       "│    │    └─DecoderBlock: 3-22                               [1, 682, 128]             --\n",
       "│    │    └─DecoderBlock: 3-23                               [1, 682, 128]             --\n",
       "│    │    └─DecoderBlock: 3-24                               [1, 682, 128]             --\n",
       "├─Linear: 1-6                                                [1, 682, 58]              7,482\n",
       "==============================================================================================================\n",
       "Total params: 4,652,006\n",
       "Trainable params: 4,652,006\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 2.44\n",
       "==============================================================================================================\n",
       "Input size (MB): 1.48\n",
       "Forward/backward pass size (MB): 1041.70\n",
       "Params size (MB): 18.61\n",
       "Estimated Total Size (MB): 1061.78\n",
       "=============================================================================================================="
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(net, ((1, 1, 576, 640), (1, 682)), device=\"cpu\", depth=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "506f01a3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}