Hugging Face's logo

Spaces:
sneedium
/
captcha_pixelplanet
Build error

App Files Community
2
nigger game commited on
Commit
166850f
·
1 Parent(s): b0fb16e

gfgf

Browse files
Files changed (36) hide show
  1. .config/.last_opt_in_prompt.yaml +1 -0
  2. .config/.last_survey_prompt.yaml +1 -0
  3. .config/.last_update_check.json +1 -0
  4. .config/active_config +1 -0
  5. .config/config_sentinel +0 -0
  6. .config/configurations/config_default +6 -0
  7. .config/gce +1 -0
  8. .config/logs/2022.09.14/13.43.12.857966.log +503 -0
  9. .config/logs/2022.09.14/13.43.34.754807.log +4 -0
  10. .config/logs/2022.09.14/13.43.54.433400.log +165 -0
  11. .config/logs/2022.09.14/13.44.02.480346.log +4 -0
  12. .config/logs/2022.09.14/13.44.22.887116.log +7 -0
  13. .config/logs/2022.09.14/13.44.23.609544.log +7 -0
  14. README.md +136 -12
  15. app.py +30 -0
  16. callbacks.py +360 -0
  17. .gitattributes → configs/.gitattributes +0 -0
  18. configs/train_iternet.yaml +65 -0
  19. dataset.py +278 -0
  20. demo.py +109 -0
  21. losses.py +72 -0
  22. main.py +246 -0
  23. modules/__init__.py +0 -0
  24. modules/attention.py +97 -0
  25. modules/backbone.py +42 -0
  26. modules/model.py +50 -0
  27. modules/model_alignment.py +34 -0
  28. modules/model_iternet.py +38 -0
  29. modules/model_language.py +67 -0
  30. modules/model_vision.py +151 -0
  31. modules/resnet.py +133 -0
  32. modules/transformer.py +901 -0
  33. packages.txt +1 -0
  34. requirements.txt +8 -0
  35. transforms.py +329 -0
  36. utils.py +304 -0
.config/.last_opt_in_prompt.yaml ADDED
@@ -0,0 +1 @@
 
 
1
+ {}
.config/.last_survey_prompt.yaml ADDED
@@ -0,0 +1 @@
 
 
1
+ last_prompt_time: 1663163033.713858
.config/.last_update_check.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"last_update_check_time": 1663163041.8911226, "last_update_check_revision": 20220912133630, "notifications": [], "last_nag_times": {}}
.config/active_config ADDED
@@ -0,0 +1 @@
 
 
1
+ default
.config/config_sentinel ADDED
File without changes
.config/configurations/config_default ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ [component_manager]
2
+ disable_update_check = true
3
+
4
+ [compute]
5
+ gce_metadata_read_timeout_sec = 0
6
+
.config/gce ADDED
@@ -0,0 +1 @@
 
 
1
+ False
.config/logs/2022.09.14/13.43.12.857966.log ADDED
@@ -0,0 +1,503 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ 2022-09-14 13:43:12,862 DEBUG root Loaded Command Group: ['gcloud', 'components']
2
+ 2022-09-14 13:43:12,866 DEBUG root Loaded Command Group: ['gcloud', 'components', 'update']
3
+ 2022-09-14 13:43:12,868 DEBUG root Running [gcloud.components.update] with arguments: [--allow-no-backup: "True", --quiet: "True", COMPONENT-IDS:7: "['core', 'gcloud-deps', 'bq', 'gcloud', 'gcloud-crc32c', 'gsutil', 'anthoscli']"]
4
+ 2022-09-14 13:43:12,875 INFO ___FILE_ONLY___ Beginning update. This process may take several minutes.
5
+
6
+ 2022-09-14 13:43:24,906 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
7
+ 2022-09-14 13:43:24,995 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components-2.json HTTP/1.1" 200 197751
8
+ 2022-09-14 13:43:25,045 INFO ___FILE_ONLY___
9
+
10
+ 2022-09-14 13:43:25,045 INFO ___FILE_ONLY___
11
+ Your current Google Cloud CLI version is: 402.0.0
12
+
13
+ 2022-09-14 13:43:25,045 INFO ___FILE_ONLY___ Installing components from version: 402.0.0
14
+
15
+ 2022-09-14 13:43:25,045 INFO ___FILE_ONLY___
16
+
17
+ 2022-09-14 13:43:25,052 INFO ___FILE_ONLY___ ┌─────────────────────────────────────────────────────────────────────────────┐
18
+ 2022-09-14 13:43:25,052 INFO ___FILE_ONLY___
19
+
20
+ 2022-09-14 13:43:25,052 INFO ___FILE_ONLY___ │ These components will be installed. │
21
+ 2022-09-14 13:43:25,052 INFO ___FILE_ONLY___
22
+
23
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ ├─────────────────────────────────────────────────────┬────────────┬──────────┤
24
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___
25
+
26
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ │ Name │ Version │ Size │
27
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___
28
+
29
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ ├─────────────────────────────────────────────────────┼────────────┼──────────┤
30
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___
31
+
32
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ │
33
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ BigQuery Command Line Tool
34
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___
35
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ │
36
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ 2.0.75
37
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___
38
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ │
39
+ 2022-09-14 13:43:25,053 INFO ___FILE_ONLY___ 1.6 MiB
40
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___
41
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___ │
42
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___
43
+
44
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___ │
45
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___ BigQuery Command Line Tool (Platform Specific)
46
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___
47
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___ │
48
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___ 2.0.75
49
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___
50
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___ │
51
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___ < 1 MiB
52
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___
53
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___ │
54
+ 2022-09-14 13:43:25,054 INFO ___FILE_ONLY___
55
+
56
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ │
57
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ Cloud Storage Command Line Tool
58
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___
59
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ │
60
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ 5.13
61
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___
62
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ │
63
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ 15.5 MiB
64
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___
65
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ │
66
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___
67
+
68
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ │
69
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___ Cloud Storage Command Line Tool (Platform Specific)
70
+ 2022-09-14 13:43:25,055 INFO ___FILE_ONLY___
71
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ │
72
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ 5.11
73
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___
74
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ │
75
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ < 1 MiB
76
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___
77
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ │
78
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___
79
+
80
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ │
81
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ Google Cloud CLI Core Libraries (Platform Specific)
82
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___
83
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ │
84
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___ 2022.08.05
85
+ 2022-09-14 13:43:25,056 INFO ___FILE_ONLY___
86
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ │
87
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ < 1 MiB
88
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___
89
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ │
90
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___
91
+
92
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ │
93
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ Google Cloud CRC32C Hash Tool
94
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___
95
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ │
96
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ 1.0.0
97
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___
98
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ │
99
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___ 1.2 MiB
100
+ 2022-09-14 13:43:25,057 INFO ___FILE_ONLY___
101
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ │
102
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___
103
+
104
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ │
105
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ anthoscli
106
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___
107
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ │
108
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ 0.2.29
109
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___
110
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ │
111
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ 37.6 MiB
112
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___
113
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ │
114
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___
115
+
116
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ │
117
+ 2022-09-14 13:43:25,058 INFO ___FILE_ONLY___ gcloud cli dependencies
118
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___
119
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___ │
120
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___ 2021.04.16
121
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___
122
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___ │
123
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___ < 1 MiB
124
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___
125
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___ │
126
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___
127
+
128
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___ └─────────────────────────────────────────────────────┴────────────┴──────────┘
129
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___
130
+
131
+ 2022-09-14 13:43:25,059 INFO ___FILE_ONLY___
132
+
133
+ 2022-09-14 13:43:25,062 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
134
+ 2022-09-14 13:43:25,093 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/RELEASE_NOTES HTTP/1.1" 200 867084
135
+ 2022-09-14 13:43:25,185 INFO ___FILE_ONLY___ For the latest full release notes, please visit:
136
+ https://cloud.google.com/sdk/release_notes
137
+
138
+
139
+ 2022-09-14 13:43:25,191 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
140
+
141
+ 2022-09-14 13:43:25,191 INFO ___FILE_ONLY___ ╠═ Creating update staging area ═╣
142
+
143
+ 2022-09-14 13:43:25,191 INFO ___FILE_ONLY___ ╚
144
+ 2022-09-14 13:43:25,192 INFO ___FILE_ONLY___ ══════
145
+ 2022-09-14 13:43:25,192 INFO ___FILE_ONLY___ ══════
146
+ 2022-09-14 13:43:25,192 INFO ___FILE_ONLY___ ══════
147
+ 2022-09-14 13:43:25,441 INFO ___FILE_ONLY___ ═
148
+ 2022-09-14 13:43:25,497 INFO ___FILE_ONLY___ ═
149
+ 2022-09-14 13:43:25,543 INFO ___FILE_ONLY___ ═
150
+ 2022-09-14 13:43:25,605 INFO ___FILE_ONLY___ ═
151
+ 2022-09-14 13:43:25,656 INFO ___FILE_ONLY___ ═
152
+ 2022-09-14 13:43:25,707 INFO ___FILE_ONLY___ ═
153
+ 2022-09-14 13:43:25,761 INFO ___FILE_ONLY___ ═
154
+ 2022-09-14 13:43:25,817 INFO ___FILE_ONLY___ ═
155
+ 2022-09-14 13:43:25,892 INFO ___FILE_ONLY___ ═
156
+ 2022-09-14 13:43:25,953 INFO ___FILE_ONLY___ ═
157
+ 2022-09-14 13:43:26,013 INFO ___FILE_ONLY___ ═
158
+ 2022-09-14 13:43:26,068 INFO ___FILE_ONLY___ ═
159
+ 2022-09-14 13:43:26,121 INFO ___FILE_ONLY___ ═
160
+ 2022-09-14 13:43:26,177 INFO ___FILE_ONLY___ ═
161
+ 2022-09-14 13:43:26,230 INFO ___FILE_ONLY___ ═
162
+ 2022-09-14 13:43:26,358 INFO ___FILE_ONLY___ ═
163
+ 2022-09-14 13:43:26,415 INFO ___FILE_ONLY___ ═
164
+ 2022-09-14 13:43:26,594 INFO ___FILE_ONLY___ ═
165
+ 2022-09-14 13:43:26,669 INFO ___FILE_ONLY___ ═
166
+ 2022-09-14 13:43:26,746 INFO ___FILE_ONLY___ ═
167
+ 2022-09-14 13:43:26,786 INFO ___FILE_ONLY___ ═
168
+ 2022-09-14 13:43:26,822 INFO ___FILE_ONLY___ ═
169
+ 2022-09-14 13:43:26,852 INFO ___FILE_ONLY___ ═
170
+ 2022-09-14 13:43:26,882 INFO ___FILE_ONLY___ ═
171
+ 2022-09-14 13:43:26,911 INFO ___FILE_ONLY___ ═
172
+ 2022-09-14 13:43:26,942 INFO ___FILE_ONLY___ ═
173
+ 2022-09-14 13:43:27,036 INFO ___FILE_ONLY___ ═
174
+ 2022-09-14 13:43:27,161 INFO ___FILE_ONLY___ ═
175
+ 2022-09-14 13:43:27,224 INFO ___FILE_ONLY___ ═
176
+ 2022-09-14 13:43:27,270 INFO ___FILE_ONLY___ ═
177
+ 2022-09-14 13:43:27,350 INFO ___FILE_ONLY___ ═
178
+ 2022-09-14 13:43:27,432 INFO ___FILE_ONLY___ ═
179
+ 2022-09-14 13:43:27,489 INFO ___FILE_ONLY___ ═
180
+ 2022-09-14 13:43:27,550 INFO ___FILE_ONLY___ ═
181
+ 2022-09-14 13:43:27,627 INFO ___FILE_ONLY___ ═
182
+ 2022-09-14 13:43:27,687 INFO ___FILE_ONLY___ ═
183
+ 2022-09-14 13:43:27,767 INFO ___FILE_ONLY___ ═
184
+ 2022-09-14 13:43:27,828 INFO ___FILE_ONLY___ ═
185
+ 2022-09-14 13:43:27,865 INFO ___FILE_ONLY___ ═
186
+ 2022-09-14 13:43:27,918 INFO ___FILE_ONLY___ ═
187
+ 2022-09-14 13:43:27,974 INFO ___FILE_ONLY___ ═
188
+ 2022-09-14 13:43:28,015 INFO ___FILE_ONLY___ ═
189
+ 2022-09-14 13:43:28,015 INFO ___FILE_ONLY___ ╝
190
+
191
+ 2022-09-14 13:43:28,116 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
192
+
193
+ 2022-09-14 13:43:28,116 INFO ___FILE_ONLY___ ╠═ Installing: BigQuery Command Line Tool ═╣
194
+
195
+ 2022-09-14 13:43:28,117 INFO ___FILE_ONLY___ ╚
196
+ 2022-09-14 13:43:28,121 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
197
+ 2022-09-14 13:43:28,216 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-bq-20220826172526.tar.gz HTTP/1.1" 200 1649904
198
+ 2022-09-14 13:43:28,301 INFO ___FILE_ONLY___ ═
199
+ 2022-09-14 13:43:28,301 INFO ___FILE_ONLY___ ═
200
+ 2022-09-14 13:43:28,302 INFO ___FILE_ONLY___ ═
201
+ 2022-09-14 13:43:28,302 INFO ___FILE_ONLY___ ═
202
+ 2022-09-14 13:43:28,302 INFO ___FILE_ONLY___ ═
203
+ 2022-09-14 13:43:28,302 INFO ___FILE_ONLY___ ═
204
+ 2022-09-14 13:43:28,302 INFO ___FILE_ONLY___ ═
205
+ 2022-09-14 13:43:28,302 INFO ___FILE_ONLY___ ═
206
+ 2022-09-14 13:43:28,303 INFO ___FILE_ONLY___ ═
207
+ 2022-09-14 13:43:28,303 INFO ___FILE_ONLY___ ═
208
+ 2022-09-14 13:43:28,303 INFO ___FILE_ONLY___ ═
209
+ 2022-09-14 13:43:28,303 INFO ___FILE_ONLY___ ═
210
+ 2022-09-14 13:43:28,303 INFO ___FILE_ONLY___ ═
211
+ 2022-09-14 13:43:28,303 INFO ___FILE_ONLY___ ═
212
+ 2022-09-14 13:43:28,303 INFO ___FILE_ONLY___ ═
213
+ 2022-09-14 13:43:28,304 INFO ___FILE_ONLY___ ═
214
+ 2022-09-14 13:43:28,304 INFO ___FILE_ONLY___ ═
215
+ 2022-09-14 13:43:28,304 INFO ___FILE_ONLY___ ═
216
+ 2022-09-14 13:43:28,304 INFO ___FILE_ONLY___ ═
217
+ 2022-09-14 13:43:28,304 INFO ___FILE_ONLY___ ═
218
+ 2022-09-14 13:43:28,304 INFO ___FILE_ONLY___ ═
219
+ 2022-09-14 13:43:28,305 INFO ___FILE_ONLY___ ═
220
+ 2022-09-14 13:43:28,305 INFO ___FILE_ONLY___ ═
221
+ 2022-09-14 13:43:28,305 INFO ___FILE_ONLY___ ═
222
+ 2022-09-14 13:43:28,305 INFO ___FILE_ONLY___ ═
223
+ 2022-09-14 13:43:28,305 INFO ___FILE_ONLY___ ═
224
+ 2022-09-14 13:43:28,305 INFO ___FILE_ONLY___ ═
225
+ 2022-09-14 13:43:28,305 INFO ___FILE_ONLY___ ═
226
+ 2022-09-14 13:43:28,306 INFO ___FILE_ONLY___ ═
227
+ 2022-09-14 13:43:28,306 INFO ___FILE_ONLY___ ═
228
+ 2022-09-14 13:43:28,430 INFO ___FILE_ONLY___ ═
229
+ 2022-09-14 13:43:28,434 INFO ___FILE_ONLY___ ═
230
+ 2022-09-14 13:43:28,439 INFO ___FILE_ONLY___ ═
231
+ 2022-09-14 13:43:28,446 INFO ___FILE_ONLY___ ═
232
+ 2022-09-14 13:43:28,450 INFO ___FILE_ONLY___ ═
233
+ 2022-09-14 13:43:28,454 INFO ___FILE_ONLY___ ═
234
+ 2022-09-14 13:43:28,458 INFO ___FILE_ONLY___ ═
235
+ 2022-09-14 13:43:28,463 INFO ___FILE_ONLY___ ═
236
+ 2022-09-14 13:43:28,467 INFO ___FILE_ONLY___ ═
237
+ 2022-09-14 13:43:28,472 INFO ___FILE_ONLY___ ═
238
+ 2022-09-14 13:43:28,476 INFO ___FILE_ONLY___ ═
239
+ 2022-09-14 13:43:28,481 INFO ___FILE_ONLY___ ═
240
+ 2022-09-14 13:43:28,487 INFO ___FILE_ONLY___ ═
241
+ 2022-09-14 13:43:28,493 INFO ___FILE_ONLY___ ═
242
+ 2022-09-14 13:43:28,498 INFO ___FILE_ONLY___ ═
243
+ 2022-09-14 13:43:28,502 INFO ___FILE_ONLY___ ═
244
+ 2022-09-14 13:43:28,506 INFO ___FILE_ONLY___ ═
245
+ 2022-09-14 13:43:28,510 INFO ___FILE_ONLY___ ═
246
+ 2022-09-14 13:43:28,514 INFO ___FILE_ONLY___ ═
247
+ 2022-09-14 13:43:28,518 INFO ___FILE_ONLY___ ═
248
+ 2022-09-14 13:43:28,524 INFO ___FILE_ONLY___ ═
249
+ 2022-09-14 13:43:28,528 INFO ___FILE_ONLY___ ═
250
+ 2022-09-14 13:43:28,533 INFO ___FILE_ONLY___ ═
251
+ 2022-09-14 13:43:28,537 INFO ___FILE_ONLY___ ═
252
+ 2022-09-14 13:43:28,542 INFO ___FILE_ONLY___ ═
253
+ 2022-09-14 13:43:28,546 INFO ___FILE_ONLY___ ═
254
+ 2022-09-14 13:43:28,551 INFO ___FILE_ONLY___ ═
255
+ 2022-09-14 13:43:28,555 INFO ___FILE_ONLY___ ═
256
+ 2022-09-14 13:43:28,560 INFO ___FILE_ONLY___ ═
257
+ 2022-09-14 13:43:28,564 INFO ___FILE_ONLY___ ═
258
+ 2022-09-14 13:43:28,564 INFO ___FILE_ONLY___ ╝
259
+
260
+ 2022-09-14 13:43:28,581 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
261
+
262
+ 2022-09-14 13:43:28,581 INFO ___FILE_ONLY___ ╠═ Installing: BigQuery Command Line Tool (Platform Spec... ═╣
263
+
264
+ 2022-09-14 13:43:28,581 INFO ___FILE_ONLY___ ╚
265
+ 2022-09-14 13:43:28,586 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
266
+ 2022-09-14 13:43:28,671 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-bq-nix-20220805142045.tar.gz HTTP/1.1" 200 1816
267
+ 2022-09-14 13:43:28,672 INFO ___FILE_ONLY___ ══════════════════════════════
268
+ 2022-09-14 13:43:28,673 INFO ___FILE_ONLY___ ══════════════════════════════
269
+ 2022-09-14 13:43:28,673 INFO ___FILE_ONLY___ ╝
270
+
271
+ 2022-09-14 13:43:28,685 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
272
+
273
+ 2022-09-14 13:43:28,685 INFO ___FILE_ONLY___ ╠═ Installing: Cloud Storage Command Line Tool ═╣
274
+
275
+ 2022-09-14 13:43:28,685 INFO ___FILE_ONLY___ ╚
276
+ 2022-09-14 13:43:28,689 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
277
+ 2022-09-14 13:43:28,785 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-gsutil-20220912133630.tar.gz HTTP/1.1" 200 16234713
278
+ 2022-09-14 13:43:28,948 INFO ___FILE_ONLY___ ═
279
+ 2022-09-14 13:43:28,949 INFO ___FILE_ONLY___ ═
280
+ 2022-09-14 13:43:28,949 INFO ___FILE_ONLY___ ═
281
+ 2022-09-14 13:43:28,950 INFO ___FILE_ONLY___ ═
282
+ 2022-09-14 13:43:28,951 INFO ___FILE_ONLY___ ═
283
+ 2022-09-14 13:43:28,952 INFO ___FILE_ONLY___ ═
284
+ 2022-09-14 13:43:28,953 INFO ___FILE_ONLY___ ═
285
+ 2022-09-14 13:43:28,954 INFO ___FILE_ONLY___ ═
286
+ 2022-09-14 13:43:28,954 INFO ___FILE_ONLY___ ═
287
+ 2022-09-14 13:43:28,955 INFO ___FILE_ONLY___ ═
288
+ 2022-09-14 13:43:28,956 INFO ___FILE_ONLY___ ═
289
+ 2022-09-14 13:43:28,957 INFO ___FILE_ONLY___ ═
290
+ 2022-09-14 13:43:28,957 INFO ___FILE_ONLY___ ═
291
+ 2022-09-14 13:43:28,958 INFO ___FILE_ONLY___ ═
292
+ 2022-09-14 13:43:28,959 INFO ___FILE_ONLY___ ═
293
+ 2022-09-14 13:43:28,960 INFO ___FILE_ONLY___ ═
294
+ 2022-09-14 13:43:28,961 INFO ___FILE_ONLY___ ═
295
+ 2022-09-14 13:43:28,961 INFO ___FILE_ONLY___ ═
296
+ 2022-09-14 13:43:28,962 INFO ___FILE_ONLY___ ═
297
+ 2022-09-14 13:43:28,963 INFO ___FILE_ONLY___ ═
298
+ 2022-09-14 13:43:28,964 INFO ___FILE_ONLY___ ═
299
+ 2022-09-14 13:43:28,965 INFO ___FILE_ONLY___ ═
300
+ 2022-09-14 13:43:28,965 INFO ___FILE_ONLY___ ═
301
+ 2022-09-14 13:43:28,966 INFO ___FILE_ONLY___ ═
302
+ 2022-09-14 13:43:28,967 INFO ___FILE_ONLY___ ═
303
+ 2022-09-14 13:43:28,968 INFO ___FILE_ONLY___ ═
304
+ 2022-09-14 13:43:28,968 INFO ___FILE_ONLY___ ═
305
+ 2022-09-14 13:43:28,969 INFO ___FILE_ONLY___ ═
306
+ 2022-09-14 13:43:28,970 INFO ___FILE_ONLY___ ═
307
+ 2022-09-14 13:43:28,971 INFO ___FILE_ONLY___ ═
308
+ 2022-09-14 13:43:30,254 INFO ___FILE_ONLY___ ═
309
+ 2022-09-14 13:43:30,308 INFO ___FILE_ONLY___ ═
310
+ 2022-09-14 13:43:30,353 INFO ___FILE_ONLY___ ═
311
+ 2022-09-14 13:43:30,396 INFO ___FILE_ONLY___ ═
312
+ 2022-09-14 13:43:30,450 INFO ___FILE_ONLY___ ═
313
+ 2022-09-14 13:43:30,509 INFO ___FILE_ONLY___ ═
314
+ 2022-09-14 13:43:30,563 INFO ___FILE_ONLY___ ═
315
+ 2022-09-14 13:43:30,628 INFO ___FILE_ONLY___ ═
316
+ 2022-09-14 13:43:30,694 INFO ___FILE_ONLY___ ═
317
+ 2022-09-14 13:43:30,765 INFO ___FILE_ONLY___ ═
318
+ 2022-09-14 13:43:30,812 INFO ___FILE_ONLY___ ═
319
+ 2022-09-14 13:43:30,858 INFO ___FILE_ONLY___ ═
320
+ 2022-09-14 13:43:30,915 INFO ___FILE_ONLY___ ═
321
+ 2022-09-14 13:43:30,981 INFO ___FILE_ONLY___ ═
322
+ 2022-09-14 13:43:31,038 INFO ___FILE_ONLY___ ═
323
+ 2022-09-14 13:43:31,082 INFO ___FILE_ONLY___ ═
324
+ 2022-09-14 13:43:31,155 INFO ___FILE_ONLY___ ═
325
+ 2022-09-14 13:43:31,206 INFO ___FILE_ONLY___ ═
326
+ 2022-09-14 13:43:31,249 INFO ___FILE_ONLY___ ═
327
+ 2022-09-14 13:43:31,297 INFO ___FILE_ONLY___ ═
328
+ 2022-09-14 13:43:31,347 INFO ___FILE_ONLY___ ═
329
+ 2022-09-14 13:43:31,398 INFO ___FILE_ONLY___ ═
330
+ 2022-09-14 13:43:31,456 INFO ___FILE_ONLY___ ═
331
+ 2022-09-14 13:43:31,506 INFO ___FILE_ONLY___ ═
332
+ 2022-09-14 13:43:31,562 INFO ___FILE_ONLY___ ═
333
+ 2022-09-14 13:43:31,619 INFO ___FILE_ONLY___ ═
334
+ 2022-09-14 13:43:31,664 INFO ___FILE_ONLY___ ═
335
+ 2022-09-14 13:43:31,707 INFO ___FILE_ONLY___ ═
336
+ 2022-09-14 13:43:31,750 INFO ___FILE_ONLY___ ═
337
+ 2022-09-14 13:43:31,810 INFO ___FILE_ONLY___ ═
338
+ 2022-09-14 13:43:31,810 INFO ___FILE_ONLY___ ╝
339
+
340
+ 2022-09-14 13:43:31,908 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
341
+
342
+ 2022-09-14 13:43:31,909 INFO ___FILE_ONLY___ ╠═ Installing: Cloud Storage Command Line Tool (Platform... ═╣
343
+
344
+ 2022-09-14 13:43:31,909 INFO ___FILE_ONLY___ ╚
345
+ 2022-09-14 13:43:31,913 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
346
+ 2022-09-14 13:43:32,000 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-gsutil-nix-20220805142045.tar.gz HTTP/1.1" 200 1831
347
+ 2022-09-14 13:43:32,001 INFO ___FILE_ONLY___ ══════════════════════════════
348
+ 2022-09-14 13:43:32,002 INFO ___FILE_ONLY___ ══════════════════════════════
349
+ 2022-09-14 13:43:32,002 INFO ___FILE_ONLY___ ╝
350
+
351
+ 2022-09-14 13:43:32,014 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
352
+
353
+ 2022-09-14 13:43:32,015 INFO ___FILE_ONLY___ ╠═ Installing: Default set of gcloud commands ═╣
354
+
355
+ 2022-09-14 13:43:32,015 INFO ___FILE_ONLY___ ╚
356
+ 2022-09-14 13:43:32,020 INFO ___FILE_ONLY___ ════════════════════════════════════════════════════════════
357
+ 2022-09-14 13:43:32,020 INFO ___FILE_ONLY___ ╝
358
+
359
+ 2022-09-14 13:43:32,025 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
360
+
361
+ 2022-09-14 13:43:32,025 INFO ___FILE_ONLY___ ╠═ Installing: Google Cloud CLI Core Libraries (Platform... ═╣
362
+
363
+ 2022-09-14 13:43:32,025 INFO ___FILE_ONLY___ ╚
364
+ 2022-09-14 13:43:32,029 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
365
+ 2022-09-14 13:43:32,118 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-core-nix-20220805142045.tar.gz HTTP/1.1" 200 2202
366
+ 2022-09-14 13:43:32,119 INFO ___FILE_ONLY___ ══════════════════════════════
367
+ 2022-09-14 13:43:32,120 INFO ___FILE_ONLY___ ═══════════════
368
+ 2022-09-14 13:43:32,121 INFO ___FILE_ONLY___ ═══════════════
369
+ 2022-09-14 13:43:32,121 INFO ___FILE_ONLY___ ╝
370
+
371
+ 2022-09-14 13:43:32,132 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
372
+
373
+ 2022-09-14 13:43:32,133 INFO ___FILE_ONLY___ ╠═ Installing: Google Cloud CRC32C Hash Tool ═╣
374
+
375
+ 2022-09-14 13:43:32,133 INFO ___FILE_ONLY___ ╚
376
+ 2022-09-14 13:43:32,137 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
377
+ 2022-09-14 13:43:32,170 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-gcloud-crc32c-linux-x86_64-20220729144039.tar.gz HTTP/1.1" 200 1243294
378
+ 2022-09-14 13:43:32,248 INFO ___FILE_ONLY___ ═
379
+ 2022-09-14 13:43:32,248 INFO ___FILE_ONLY___ ═
380
+ 2022-09-14 13:43:32,248 INFO ___FILE_ONLY___ ═
381
+ 2022-09-14 13:43:32,248 INFO ___FILE_ONLY___ ═
382
+ 2022-09-14 13:43:32,248 INFO ___FILE_ONLY___ ═
383
+ 2022-09-14 13:43:32,248 INFO ___FILE_ONLY___ ═
384
+ 2022-09-14 13:43:32,249 INFO ___FILE_ONLY___ ═
385
+ 2022-09-14 13:43:32,249 INFO ___FILE_ONLY___ ═
386
+ 2022-09-14 13:43:32,249 INFO ___FILE_ONLY___ ═
387
+ 2022-09-14 13:43:32,249 INFO ___FILE_ONLY___ ═
388
+ 2022-09-14 13:43:32,249 INFO ___FILE_ONLY___ ═
389
+ 2022-09-14 13:43:32,249 INFO ___FILE_ONLY___ ═
390
+ 2022-09-14 13:43:32,249 INFO ___FILE_ONLY___ ═
391
+ 2022-09-14 13:43:32,250 INFO ___FILE_ONLY___ ═
392
+ 2022-09-14 13:43:32,250 INFO ___FILE_ONLY___ ═
393
+ 2022-09-14 13:43:32,250 INFO ___FILE_ONLY___ ═
394
+ 2022-09-14 13:43:32,250 INFO ___FILE_ONLY___ ═
395
+ 2022-09-14 13:43:32,250 INFO ___FILE_ONLY___ ═
396
+ 2022-09-14 13:43:32,250 INFO ___FILE_ONLY___ ═
397
+ 2022-09-14 13:43:32,250 INFO ___FILE_ONLY___ ═
398
+ 2022-09-14 13:43:32,250 INFO ___FILE_ONLY___ ═
399
+ 2022-09-14 13:43:32,251 INFO ___FILE_ONLY___ ═
400
+ 2022-09-14 13:43:32,251 INFO ___FILE_ONLY___ ═
401
+ 2022-09-14 13:43:32,251 INFO ___FILE_ONLY___ ═
402
+ 2022-09-14 13:43:32,251 INFO ___FILE_ONLY___ ═
403
+ 2022-09-14 13:43:32,251 INFO ___FILE_ONLY___ ═
404
+ 2022-09-14 13:43:32,251 INFO ___FILE_ONLY___ ═
405
+ 2022-09-14 13:43:32,251 INFO ___FILE_ONLY___ ═
406
+ 2022-09-14 13:43:32,252 INFO ___FILE_ONLY___ ═
407
+ 2022-09-14 13:43:32,252 INFO ___FILE_ONLY___ ═
408
+ 2022-09-14 13:43:32,288 INFO ___FILE_ONLY___ ═══════════════
409
+ 2022-09-14 13:43:32,289 INFO ___FILE_ONLY___ ═══════════════
410
+ 2022-09-14 13:43:32,289 INFO ___FILE_ONLY___ ╝
411
+
412
+ 2022-09-14 13:43:32,301 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
413
+
414
+ 2022-09-14 13:43:32,302 INFO ___FILE_ONLY___ ╠═ Installing: Google Cloud CRC32C Hash Tool ═╣
415
+
416
+ 2022-09-14 13:43:32,302 INFO ___FILE_ONLY___ ╚
417
+ 2022-09-14 13:43:32,307 INFO ___FILE_ONLY___ ════════════════════════════════════════════════════════════
418
+ 2022-09-14 13:43:32,307 INFO ___FILE_ONLY___ ╝
419
+
420
+ 2022-09-14 13:43:32,311 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
421
+
422
+ 2022-09-14 13:43:32,312 INFO ___FILE_ONLY___ ╠═ Installing: anthoscli ═╣
423
+
424
+ 2022-09-14 13:43:32,312 INFO ___FILE_ONLY___ ╚
425
+ 2022-09-14 13:43:32,316 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
426
+ 2022-09-14 13:43:32,408 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-anthoscli-linux-x86_64-20220826172526.tar.gz HTTP/1.1" 200 39436637
427
+ 2022-09-14 13:43:32,689 INFO ___FILE_ONLY___ ═
428
+ 2022-09-14 13:43:32,691 INFO ___FILE_ONLY___ ═
429
+ 2022-09-14 13:43:32,693 INFO ___FILE_ONLY___ ═
430
+ 2022-09-14 13:43:32,695 INFO ___FILE_ONLY___ ═
431
+ 2022-09-14 13:43:32,696 INFO ___FILE_ONLY___ ═
432
+ 2022-09-14 13:43:32,698 INFO ___FILE_ONLY___ ═
433
+ 2022-09-14 13:43:32,700 INFO ___FILE_ONLY___ ═
434
+ 2022-09-14 13:43:32,702 INFO ___FILE_ONLY___ ═
435
+ 2022-09-14 13:43:32,703 INFO ___FILE_ONLY___ ═
436
+ 2022-09-14 13:43:32,705 INFO ___FILE_ONLY___ ═
437
+ 2022-09-14 13:43:32,707 INFO ___FILE_ONLY___ ═
438
+ 2022-09-14 13:43:32,708 INFO ___FILE_ONLY___ ═
439
+ 2022-09-14 13:43:32,710 INFO ___FILE_ONLY___ ═
440
+ 2022-09-14 13:43:32,712 INFO ___FILE_ONLY___ ═
441
+ 2022-09-14 13:43:32,714 INFO ___FILE_ONLY___ ═
442
+ 2022-09-14 13:43:32,715 INFO ___FILE_ONLY___ ═
443
+ 2022-09-14 13:43:32,717 INFO ___FILE_ONLY___ ═
444
+ 2022-09-14 13:43:32,719 INFO ___FILE_ONLY___ ═
445
+ 2022-09-14 13:43:32,720 INFO ___FILE_ONLY___ ═
446
+ 2022-09-14 13:43:32,722 INFO ___FILE_ONLY___ ═
447
+ 2022-09-14 13:43:32,724 INFO ___FILE_ONLY___ ═
448
+ 2022-09-14 13:43:32,725 INFO ___FILE_ONLY___ ═
449
+ 2022-09-14 13:43:32,727 INFO ___FILE_ONLY___ ═
450
+ 2022-09-14 13:43:32,729 INFO ___FILE_ONLY___ ═
451
+ 2022-09-14 13:43:32,731 INFO ___FILE_ONLY___ ═
452
+ 2022-09-14 13:43:32,732 INFO ___FILE_ONLY___ ═
453
+ 2022-09-14 13:43:32,734 INFO ___FILE_ONLY___ ═
454
+ 2022-09-14 13:43:32,736 INFO ___FILE_ONLY___ ═
455
+ 2022-09-14 13:43:32,737 INFO ___FILE_ONLY___ ═
456
+ 2022-09-14 13:43:32,739 INFO ___FILE_ONLY___ ═
457
+ 2022-09-14 13:43:34,003 INFO ___FILE_ONLY___ ══════════
458
+ 2022-09-14 13:43:34,007 INFO ___FILE_ONLY___ ═════════
459
+ 2022-09-14 13:43:34,030 INFO ___FILE_ONLY___ ═══════════
460
+ 2022-09-14 13:43:34,030 INFO ___FILE_ONLY___ ╝
461
+
462
+ 2022-09-14 13:43:34,050 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
463
+
464
+ 2022-09-14 13:43:34,051 INFO ___FILE_ONLY___ ╠═ Installing: anthoscli ═╣
465
+
466
+ 2022-09-14 13:43:34,051 INFO ___FILE_ONLY___ ╚
467
+ 2022-09-14 13:43:34,056 INFO ___FILE_ONLY___ ════════════════════════════════════════════════════════════
468
+ 2022-09-14 13:43:34,056 INFO ___FILE_ONLY___ ╝
469
+
470
+ 2022-09-14 13:43:34,061 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
471
+
472
+ 2022-09-14 13:43:34,061 INFO ___FILE_ONLY___ ╠═ Installing: gcloud cli dependencies ═╣
473
+
474
+ 2022-09-14 13:43:34,061 INFO ___FILE_ONLY___ ╚
475
+ 2022-09-14 13:43:34,066 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
476
+ 2022-09-14 13:43:34,153 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-gcloud-deps-linux-x86_64-20210416153011.tar.gz HTTP/1.1" 200 104
477
+ 2022-09-14 13:43:34,154 INFO ___FILE_ONLY___ ══════════════════════════════
478
+ 2022-09-14 13:43:34,155 INFO ___FILE_ONLY___ ══════════════════════════════
479
+ 2022-09-14 13:43:34,155 INFO ___FILE_ONLY___ ╝
480
+
481
+ 2022-09-14 13:43:34,167 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
482
+
483
+ 2022-09-14 13:43:34,167 INFO ___FILE_ONLY___ ╠═ Creating backup and activating new installation ═╣
484
+
485
+ 2022-09-14 13:43:34,167 INFO ___FILE_ONLY___ ╚
486
+ 2022-09-14 13:43:34,168 DEBUG root Attempting to move directory [/tools/google-cloud-sdk] to [/tools/google-cloud-sdk.staging/.install/.backup]
487
+ 2022-09-14 13:43:34,168 INFO ___FILE_ONLY___ ══════════════════════════════
488
+ 2022-09-14 13:43:34,168 DEBUG root Attempting to move directory [/tools/google-cloud-sdk.staging] to [/tools/google-cloud-sdk]
489
+ 2022-09-14 13:43:34,168 INFO ___FILE_ONLY___ ══════════════════════════════
490
+ 2022-09-14 13:43:34,168 INFO ___FILE_ONLY___ ╝
491
+
492
+ 2022-09-14 13:43:34,172 DEBUG root Updating notification cache...
493
+ 2022-09-14 13:43:34,173 INFO ___FILE_ONLY___
494
+
495
+ 2022-09-14 13:43:34,177 INFO ___FILE_ONLY___ Performing post processing steps...
496
+ 2022-09-14 13:43:34,178 DEBUG root Executing command: ['python3', '-S', '/tools/google-cloud-sdk/lib/gcloud.py', 'components', 'post-process']
497
+ 2022-09-14 13:43:53,462 DEBUG ___FILE_ONLY___
498
+ 2022-09-14 13:43:53,462 DEBUG ___FILE_ONLY___
499
+ 2022-09-14 13:43:53,708 INFO ___FILE_ONLY___
500
+ Update done!
501
+
502
+
503
+ 2022-09-14 13:43:53,712 INFO root Display format: "none"
.config/logs/2022.09.14/13.43.34.754807.log ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ 2022-09-14 13:43:34,755 DEBUG root Loaded Command Group: ['gcloud', 'components']
2
+ 2022-09-14 13:43:34,759 DEBUG root Loaded Command Group: ['gcloud', 'components', 'post_process']
3
+ 2022-09-14 13:43:34,761 DEBUG root Running [gcloud.components.post-process] with arguments: []
4
+ 2022-09-14 13:43:53,353 INFO root Display format: "none"
.config/logs/2022.09.14/13.43.54.433400.log ADDED
@@ -0,0 +1,165 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ 2022-09-14 13:43:54,434 DEBUG root Loaded Command Group: ['gcloud', 'components']
2
+ 2022-09-14 13:43:54,436 DEBUG root Loaded Command Group: ['gcloud', 'components', 'update']
3
+ 2022-09-14 13:43:54,439 DEBUG root Running [gcloud.components.update] with arguments: [--quiet: "True", COMPONENT-IDS:8: "['gcloud', 'core', 'bq', 'gsutil', 'compute', 'preview', 'alpha', 'beta']"]
4
+ 2022-09-14 13:43:54,446 INFO ___FILE_ONLY___ Beginning update. This process may take several minutes.
5
+
6
+ 2022-09-14 13:43:54,454 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
7
+ 2022-09-14 13:43:54,485 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components-2.json HTTP/1.1" 200 197751
8
+ 2022-09-14 13:43:54,502 WARNING root Component [preview] no longer exists.
9
+ 2022-09-14 13:43:54,502 WARNING root Component [compute] no longer exists.
10
+ 2022-09-14 13:43:54,503 INFO ___FILE_ONLY___
11
+
12
+ 2022-09-14 13:43:54,504 INFO ___FILE_ONLY___
13
+ Your current Google Cloud CLI version is: 402.0.0
14
+
15
+ 2022-09-14 13:43:54,504 INFO ___FILE_ONLY___ Installing components from version: 402.0.0
16
+
17
+ 2022-09-14 13:43:54,504 INFO ___FILE_ONLY___
18
+
19
+ 2022-09-14 13:43:54,507 INFO ___FILE_ONLY___ ┌──────────────────────────────────────────────┐
20
+ 2022-09-14 13:43:54,507 INFO ___FILE_ONLY___
21
+
22
+ 2022-09-14 13:43:54,507 INFO ___FILE_ONLY___ │ These components will be installed. │
23
+ 2022-09-14 13:43:54,507 INFO ___FILE_ONLY___
24
+
25
+ 2022-09-14 13:43:54,507 INFO ___FILE_ONLY___ ├───────────────────────┬────────────┬─────────┤
26
+ 2022-09-14 13:43:54,507 INFO ___FILE_ONLY___
27
+
28
+ 2022-09-14 13:43:54,507 INFO ___FILE_ONLY___ │ Name │ Version │ Size │
29
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___
30
+
31
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___ ├───────────────────────┼────────────┼─────────┤
32
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___
33
+
34
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___ │
35
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___ gcloud Alpha Commands
36
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___
37
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___ │
38
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___ 2022.09.12
39
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___
40
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___ │
41
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___ < 1 MiB
42
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___
43
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___ │
44
+ 2022-09-14 13:43:54,508 INFO ___FILE_ONLY___
45
+
46
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___ │
47
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___ gcloud Beta Commands
48
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___
49
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___ │
50
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___ 2022.09.12
51
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___
52
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___ │
53
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___ < 1 MiB
54
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___
55
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___ │
56
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___
57
+
58
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___ └───────────────────────┴────────────┴─────────┘
59
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___
60
+
61
+ 2022-09-14 13:43:54,509 INFO ___FILE_ONLY___
62
+
63
+ 2022-09-14 13:43:54,512 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
64
+ 2022-09-14 13:43:54,601 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/RELEASE_NOTES HTTP/1.1" 200 867084
65
+ 2022-09-14 13:43:54,691 INFO ___FILE_ONLY___ For the latest full release notes, please visit:
66
+ https://cloud.google.com/sdk/release_notes
67
+
68
+
69
+ 2022-09-14 13:43:54,697 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
70
+
71
+ 2022-09-14 13:43:54,697 INFO ___FILE_ONLY___ ╠═ Creating update staging area ═╣
72
+
73
+ 2022-09-14 13:43:54,697 INFO ___FILE_ONLY___ ╚
74
+ 2022-09-14 13:43:54,697 INFO ___FILE_ONLY___ ══════
75
+ 2022-09-14 13:43:55,308 INFO ___FILE_ONLY___ ══════
76
+ 2022-09-14 13:43:55,308 INFO ___FILE_ONLY___ ══════
77
+ 2022-09-14 13:43:55,667 INFO ___FILE_ONLY___ ═
78
+ 2022-09-14 13:43:55,759 INFO ___FILE_ONLY___ ═
79
+ 2022-09-14 13:43:55,832 INFO ___FILE_ONLY___ ═
80
+ 2022-09-14 13:43:55,963 INFO ___FILE_ONLY___ ═
81
+ 2022-09-14 13:43:56,035 INFO ___FILE_ONLY___ ═
82
+ 2022-09-14 13:43:56,106 INFO ___FILE_ONLY___ ═
83
+ 2022-09-14 13:43:56,191 INFO ___FILE_ONLY___ ═
84
+ 2022-09-14 13:43:56,298 INFO ___FILE_ONLY___ ═
85
+ 2022-09-14 13:43:56,413 INFO ___FILE_ONLY___ ═
86
+ 2022-09-14 13:43:56,482 INFO ___FILE_ONLY___ ═
87
+ 2022-09-14 13:43:56,603 INFO ___FILE_ONLY___ ═
88
+ 2022-09-14 13:43:56,699 INFO ___FILE_ONLY___ ═
89
+ 2022-09-14 13:43:56,813 INFO ___FILE_ONLY___ ═
90
+ 2022-09-14 13:43:56,907 INFO ___FILE_ONLY___ ═
91
+ 2022-09-14 13:43:57,024 INFO ___FILE_ONLY___ ═
92
+ 2022-09-14 13:43:57,217 INFO ___FILE_ONLY___ ═
93
+ 2022-09-14 13:43:57,338 INFO ___FILE_ONLY___ ═
94
+ 2022-09-14 13:43:57,412 INFO ___FILE_ONLY___ ═
95
+ 2022-09-14 13:43:57,458 INFO ___FILE_ONLY___ ═
96
+ 2022-09-14 13:43:57,526 INFO ___FILE_ONLY___ ═
97
+ 2022-09-14 13:43:57,595 INFO ___FILE_ONLY___ ═
98
+ 2022-09-14 13:43:57,655 INFO ___FILE_ONLY___ ═
99
+ 2022-09-14 13:43:57,706 INFO ___FILE_ONLY___ ═
100
+ 2022-09-14 13:43:57,764 INFO ___FILE_ONLY___ ═
101
+ 2022-09-14 13:43:57,818 INFO ___FILE_ONLY___ ═
102
+ 2022-09-14 13:43:57,882 INFO ___FILE_ONLY___ ═
103
+ 2022-09-14 13:43:57,961 INFO ___FILE_ONLY___ ═
104
+ 2022-09-14 13:43:58,048 INFO ___FILE_ONLY___ ═
105
+ 2022-09-14 13:43:58,131 INFO ___FILE_ONLY___ ═
106
+ 2022-09-14 13:43:58,200 INFO ___FILE_ONLY___ ═
107
+ 2022-09-14 13:43:58,271 INFO ___FILE_ONLY___ ═
108
+ 2022-09-14 13:43:58,325 INFO ___FILE_ONLY___ ═
109
+ 2022-09-14 13:43:58,398 INFO ___FILE_ONLY___ ═
110
+ 2022-09-14 13:43:58,458 INFO ___FILE_ONLY___ ═
111
+ 2022-09-14 13:43:58,514 INFO ___FILE_ONLY___ ═
112
+ 2022-09-14 13:43:58,583 INFO ___FILE_ONLY___ ═
113
+ 2022-09-14 13:43:58,635 INFO ___FILE_ONLY___ ═
114
+ 2022-09-14 13:43:58,702 INFO ___FILE_ONLY___ ═
115
+ 2022-09-14 13:43:58,852 INFO ___FILE_ONLY___ ═
116
+ 2022-09-14 13:43:58,929 INFO ___FILE_ONLY___ ═
117
+ 2022-09-14 13:43:59,310 INFO ___FILE_ONLY___ ═
118
+ 2022-09-14 13:43:59,401 INFO ___FILE_ONLY___ ═
119
+ 2022-09-14 13:43:59,401 INFO ___FILE_ONLY___ ╝
120
+
121
+ 2022-09-14 13:44:01,735 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
122
+
123
+ 2022-09-14 13:44:01,735 INFO ___FILE_ONLY___ ╠═ Installing: gcloud Alpha Commands ═╣
124
+
125
+ 2022-09-14 13:44:01,735 INFO ___FILE_ONLY___ ╚
126
+ 2022-09-14 13:44:01,740 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
127
+ 2022-09-14 13:44:01,824 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-alpha-20220912133630.tar.gz HTTP/1.1" 200 800
128
+ 2022-09-14 13:44:01,825 INFO ___FILE_ONLY___ ══════════════════════════════
129
+ 2022-09-14 13:44:01,827 INFO ___FILE_ONLY___ ══════════════════════════════
130
+ 2022-09-14 13:44:01,827 INFO ___FILE_ONLY___ ╝
131
+
132
+ 2022-09-14 13:44:01,838 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
133
+
134
+ 2022-09-14 13:44:01,839 INFO ___FILE_ONLY___ ╠═ Installing: gcloud Beta Commands ═╣
135
+
136
+ 2022-09-14 13:44:01,839 INFO ___FILE_ONLY___ ╚
137
+ 2022-09-14 13:44:01,843 DEBUG urllib3.connectionpool Starting new HTTPS connection (1): dl.google.com:443
138
+ 2022-09-14 13:44:01,872 DEBUG urllib3.connectionpool https://dl.google.com:443 "GET /dl/cloudsdk/channels/rapid/components/google-cloud-sdk-beta-20220912133630.tar.gz HTTP/1.1" 200 797
139
+ 2022-09-14 13:44:01,873 INFO ___FILE_ONLY___ ══════════════════════════════
140
+ 2022-09-14 13:44:01,874 INFO ___FILE_ONLY___ ══════════════════════════════
141
+ 2022-09-14 13:44:01,874 INFO ___FILE_ONLY___ ╝
142
+
143
+ 2022-09-14 13:44:01,885 INFO ___FILE_ONLY___ ╔════════════════════════════════════════════════════════════╗
144
+
145
+ 2022-09-14 13:44:01,885 INFO ___FILE_ONLY___ ╠═ Creating backup and activating new installation ═╣
146
+
147
+ 2022-09-14 13:44:01,885 INFO ___FILE_ONLY___ ╚
148
+ 2022-09-14 13:44:01,886 DEBUG root Attempting to move directory [/tools/google-cloud-sdk] to [/tools/google-cloud-sdk.staging/.install/.backup]
149
+ 2022-09-14 13:44:01,886 INFO ___FILE_ONLY___ ══════════════════════════════
150
+ 2022-09-14 13:44:01,886 DEBUG root Attempting to move directory [/tools/google-cloud-sdk.staging] to [/tools/google-cloud-sdk]
151
+ 2022-09-14 13:44:01,886 INFO ___FILE_ONLY___ ══════════════════════════════
152
+ 2022-09-14 13:44:01,886 INFO ___FILE_ONLY___ ╝
153
+
154
+ 2022-09-14 13:44:01,890 DEBUG root Updating notification cache...
155
+ 2022-09-14 13:44:01,891 INFO ___FILE_ONLY___
156
+
157
+ 2022-09-14 13:44:01,896 INFO ___FILE_ONLY___ Performing post processing steps...
158
+ 2022-09-14 13:44:01,897 DEBUG root Executing command: ['python3', '-S', '/tools/google-cloud-sdk/lib/gcloud.py', 'components', 'post-process']
159
+ 2022-09-14 13:44:21,935 DEBUG ___FILE_ONLY___
160
+ 2022-09-14 13:44:21,936 DEBUG ___FILE_ONLY___
161
+ 2022-09-14 13:44:22,177 INFO ___FILE_ONLY___
162
+ Update done!
163
+
164
+
165
+ 2022-09-14 13:44:22,181 INFO root Display format: "none"
.config/logs/2022.09.14/13.44.02.480346.log ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ 2022-09-14 13:44:02,481 DEBUG root Loaded Command Group: ['gcloud', 'components']
2
+ 2022-09-14 13:44:02,483 DEBUG root Loaded Command Group: ['gcloud', 'components', 'post_process']
3
+ 2022-09-14 13:44:02,485 DEBUG root Running [gcloud.components.post-process] with arguments: []
4
+ 2022-09-14 13:44:21,837 INFO root Display format: "none"
.config/logs/2022.09.14/13.44.22.887116.log ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ 2022-09-14 13:44:22,889 DEBUG root Loaded Command Group: ['gcloud', 'config']
2
+ 2022-09-14 13:44:22,922 DEBUG root Loaded Command Group: ['gcloud', 'config', 'set']
3
+ 2022-09-14 13:44:22,925 DEBUG root Running [gcloud.config.set] with arguments: [SECTION/PROPERTY: "component_manager/disable_update_check", VALUE: "true"]
4
+ 2022-09-14 13:44:22,926 INFO ___FILE_ONLY___ Updated property [component_manager/disable_update_check].
5
+
6
+ 2022-09-14 13:44:22,927 INFO root Display format: "default"
7
+ 2022-09-14 13:44:22,928 DEBUG root SDK update checks are disabled.
.config/logs/2022.09.14/13.44.23.609544.log ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ 2022-09-14 13:44:23,611 DEBUG root Loaded Command Group: ['gcloud', 'config']
2
+ 2022-09-14 13:44:23,638 DEBUG root Loaded Command Group: ['gcloud', 'config', 'set']
3
+ 2022-09-14 13:44:23,641 DEBUG root Running [gcloud.config.set] with arguments: [SECTION/PROPERTY: "compute/gce_metadata_read_timeout_sec", VALUE: "0"]
4
+ 2022-09-14 13:44:23,642 INFO ___FILE_ONLY___ Updated property [compute/gce_metadata_read_timeout_sec].
5
+
6
+ 2022-09-14 13:44:23,644 INFO root Display format: "default"
7
+ 2022-09-14 13:44:23,645 DEBUG root SDK update checks are disabled.
README.md CHANGED
@@ -1,12 +1,136 @@
1
- ---
2
- title: Captcha Pixelplanet
3
- emoji: 👁
4
- colorFrom: purple
5
- colorTo: pink
6
- sdk: gradio
7
- sdk_version: 3.3.1
8
- app_file: app.py
9
- pinned: false
10
- ---
11
-
12
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # IterVM: Iterative Vision Modeling Module for Scene Text Recognition
2
+
3
+ The official code of [IterNet](https://arxiv.org/abs/2204.02630).
4
+
5
+ We propose IterVM, an iterative approach for visual feature extraction which can significantly improve scene text recognition accuracy.
6
+ IterVM repeatedly uses the high-level visual feature extracted at the previous iteration to enhance the multi-level features extracted at the subsequent iteration.
7
+
8
+
9
+ ![framework](./figures/framework.png)
10
+
11
+
12
+ ## Runtime Environment
13
+ ```
14
+ pip install -r requirements.txt
15
+ ```
16
+ Note: `fastai==1.0.60` is required.
17
+
18
+ ## Datasets
19
+ <details>
20
+ <summary>Training datasets (Click to expand) </summary>
21
+ 1. [MJSynth](http://www.robots.ox.ac.uk/~vgg/data/text/) (MJ):
22
+ - Use `tools/create_lmdb_dataset.py` to convert images into LMDB dataset
23
+ - [LMDB dataset BaiduNetdisk(passwd:n23k)](https://pan.baidu.com/s/1mgnTiyoR8f6Cm655rFI4HQ)
24
+ 2. [SynthText](http://www.robots.ox.ac.uk/~vgg/data/scenetext/) (ST):
25
+ - Use `tools/crop_by_word_bb.py` to crop images from original [SynthText](http://www.robots.ox.ac.uk/~vgg/data/scenetext/) dataset, and convert images into LMDB dataset by `tools/create_lmdb_dataset.py`
26
+ - [LMDB dataset BaiduNetdisk(passwd:n23k)](https://pan.baidu.com/s/1mgnTiyoR8f6Cm655rFI4HQ)
27
+ 3. [WikiText103](https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-v1.zip), which is only used for pre-trainig language models:
28
+ - Use `notebooks/prepare_wikitext103.ipynb` to convert text into CSV format.
29
+ - [CSV dataset BaiduNetdisk(passwd:dk01)](https://pan.baidu.com/s/1yabtnPYDKqhBb_Ie9PGFXA)
30
+ </details>
31
+
32
+ <details>
33
+ <summary>Evaluation datasets (Click to expand) </summary>
34
+ - Evaluation datasets, LMDB datasets can be downloaded from [BaiduNetdisk(passwd:1dbv)](https://pan.baidu.com/s/1RUg3Akwp7n8kZYJ55rU5LQ), [GoogleDrive](https://drive.google.com/file/d/1dTI0ipu14Q1uuK4s4z32DqbqF3dJPdkk/view?usp=sharing).
35
+ 1. ICDAR 2013 (IC13)
36
+ 2. ICDAR 2015 (IC15)
37
+ 3. IIIT5K Words (IIIT)
38
+ 4. Street View Text (SVT)
39
+ 5. Street View Text-Perspective (SVTP)
40
+ 6. CUTE80 (CUTE)
41
+ </details>
42
+
43
+ <details>
44
+ <summary>The structure of `data` directory (Click to expand) </summary>
45
+ - The structure of `data` directory is
46
+ ```
47
+ data
48
+ ├── charset_36.txt
49
+ ├── evaluation
50
+ │   ├── CUTE80
51
+ │   ├── IC13_857
52
+ │   ├── IC15_1811
53
+ │   ├── IIIT5k_3000
54
+ │   ├── SVT
55
+ │   └── SVTP
56
+ ├── training
57
+ │   ├── MJ
58
+ │   │   ├── MJ_test
59
+ │   │   ├── MJ_train
60
+ │   │   └── MJ_valid
61
+ │   └── ST
62
+ ├── WikiText-103.csv
63
+ └── WikiText-103_eval_d1.csv
64
+ ```
65
+ </details>
66
+
67
+ ## Pretrained Models
68
+
69
+ Get the pretrained models from [GoogleDrive](https://drive.google.com/drive/folders/1C8NMI8Od8mQUMlsnkHNLkYj73kbAQ7Bl?usp=sharing). Performances of the pretrained models are summaried as follows:
70
+
71
+ |Model|IC13|SVT|IIIT|IC15|SVTP|CUTE|AVG|
72
+ |-|-|-|-|-|-|-|-|
73
+ |IterNet|97.9|95.1|96.9|87.7|90.9|91.3|93.8|
74
+
75
+ ## Training
76
+
77
+ 1. Pre-train vision model
78
+ ```
79
+ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python main.py --config=configs/pretrain_vm.yaml
80
+ ```
81
+ 2. Pre-train language model
82
+ ```
83
+ CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py --config=configs/pretrain_language_model.yaml
84
+ ```
85
+ 3. Train IterNet
86
+ ```
87
+ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python main.py --config=configs/train_iternet.yaml
88
+ ```
89
+ Note:
90
+ - You can set the `checkpoint` path for vision model (vm) and language model separately for specific pretrained model, or set to `None` to train from scratch
91
+
92
+
93
+ ## Evaluation
94
+
95
+ ```
96
+ CUDA_VISIBLE_DEVICES=0 python main.py --config=configs/train_iternet.yaml --phase test --image_only
97
+ ```
98
+ Additional flags:
99
+ - `--checkpoint /path/to/checkpoint` set the path of evaluation model
100
+ - `--test_root /path/to/dataset` set the path of evaluation dataset
101
+ - `--model_eval [alignment|vision]` which sub-model to evaluate
102
+ - `--image_only` disable dumping visualization of attention masks
103
+
104
+ ## Run Demo
105
+ [<a href="https://colab.research.google.com/drive/1XmZGJzFF95uafmARtJMudPLLKBO2eXLv?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>](https://colab.research.google.com/drive/1XmZGJzFF95uafmARtJMudPLLKBO2eXLv?usp=sharing)
106
+
107
+ ```
108
+ python demo.py --config=configs/train_iternet.yaml --input=figures/demo
109
+ ```
110
+ Additional flags:
111
+ - `--config /path/to/config` set the path of configuration file
112
+ - `--input /path/to/image-directory` set the path of image directory or wildcard path, e.g, `--input='figs/test/*.png'`
113
+ - `--checkpoint /path/to/checkpoint` set the path of trained model
114
+ - `--cuda [-1|0|1|2|3...]` set the cuda id, by default -1 is set and stands for cpu
115
+ - `--model_eval [alignment|vision]` which sub-model to use
116
+ - `--image_only` disable dumping visualization of attention masks
117
+
118
+
119
+ ## Citation
120
+ If you find our method useful for your reserach, please cite
121
+ ```bash
122
+ @article{chu2022itervm,
123
+ title={IterVM: Iterative Vision Modeling Module for Scene Text Recognition},
124
+ author={Chu, Xiaojie and Wang, Yongtao},
125
+ journal={arXiv preprint arXiv:2204.02630},
126
+ year={2022}
127
+ }
128
+ ```
129
+
130
+ ## License
131
+ The project is only free for academic research purposes, but needs authorization for commerce. For commerce permission, please contact wyt@pku.edu.cn.
132
+
133
+ ## Acknowledgements
134
+ This project is based on [ABINet](https://github.com/FangShancheng/ABINet.git).
135
+ Thanks for their great works.
136
+
app.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import glob
2
+ import gradio as gr
3
+ from demo import get_model, preprocess, postprocess, load
4
+ from utils import Config, Logger, CharsetMapper
5
+
6
+ config = Config('configs/train_iternet.yaml')
7
+ config.model_vision_checkpoint = None
8
+ model = get_model(config)
9
+ model = load(model, 'workdir/train-iternet/best-train-iternet.pth')
10
+ charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)
11
+
12
+ def process_image(image):
13
+ img = image.convert('RGB')
14
+ img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
15
+ res = model(img)
16
+ return postprocess(res, charset, 'alignment')[0][0]
17
+
18
+ title = "Interactive demo: ABINet"
19
+ description = "Demo for ABINet, ABINet uses a vision model and an explicit language model to recognize text in the wild, which are trained in end-to-end way. The language model (BCN) achieves bidirectional language representation in simulating cloze test, additionally utilizing iterative correction strategy. To use it, simply upload a (single-text line) image or use one of the example images below and click 'submit'. Results will show up in a few seconds."
20
+ article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2103.06495.pdf'>Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition</a> | <a href='https://github.com/FangShancheng/ABINet'>Github Repo</a></p>"
21
+
22
+ iface = gr.Interface(fn=process_image,
23
+ inputs=gr.inputs.Image(type="pil"),
24
+ outputs=gr.outputs.Textbox(),
25
+ title=title,
26
+ description=description,
27
+ article=article,
28
+ examples=glob.glob('figs/test/*.png'))
29
+
30
+ iface.launch(debug=True)
callbacks.py ADDED
@@ -0,0 +1,360 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ import shutil
3
+ import time
4
+
5
+ import editdistance as ed
6
+ import torchvision.utils as vutils
7
+ from fastai.callbacks.tensorboard import (LearnerTensorboardWriter,
8
+ SummaryWriter, TBWriteRequest,
9
+ asyncTBWriter)
10
+ from fastai.vision import *
11
+ from torch.nn.parallel import DistributedDataParallel
12
+ from torchvision import transforms
13
+
14
+ import dataset
15
+ from utils import CharsetMapper, Timer, blend_mask
16
+
17
+
18
+ class IterationCallback(LearnerTensorboardWriter):
19
+ "A `TrackerCallback` that monitor in each iteration."
20
+ def __init__(self, learn:Learner, name:str='model', checpoint_keep_num=5,
21
+ show_iters:int=50, eval_iters:int=1000, save_iters:int=20000,
22
+ start_iters:int=0, stats_iters=20000):
23
+ #if self.learn.rank is not None: time.sleep(self.learn.rank) # keep all event files
24
+ super().__init__(learn, base_dir='.', name=learn.path, loss_iters=show_iters,
25
+ stats_iters=stats_iters, hist_iters=stats_iters)
26
+ self.name, self.bestname = Path(name).name, f'best-{Path(name).name}'
27
+ self.show_iters = show_iters
28
+ self.eval_iters = eval_iters
29
+ self.save_iters = save_iters
30
+ self.start_iters = start_iters
31
+ self.checpoint_keep_num = checpoint_keep_num
32
+ self.metrics_root = 'metrics/' # rewrite
33
+ self.timer = Timer()
34
+ self.host = self.learn.rank is None or self.learn.rank == 0
35
+
36
+ def _write_metrics(self, iteration:int, names:List[str], last_metrics:MetricsList)->None:
37
+ "Writes training metrics to Tensorboard."
38
+ for i, name in enumerate(names):
39
+ if last_metrics is None or len(last_metrics) < i+1: return
40
+ scalar_value = last_metrics[i]
41
+ self._write_scalar(name=name, scalar_value=scalar_value, iteration=iteration)
42
+
43
+ def _write_sub_loss(self, iteration:int, last_losses:dict)->None:
44
+ "Writes sub loss to Tensorboard."
45
+ for name, loss in last_losses.items():
46
+ scalar_value = to_np(loss)
47
+ tag = self.metrics_root + name
48
+ self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
49
+
50
+ def _save(self, name):
51
+ if isinstance(self.learn.model, DistributedDataParallel):
52
+ tmp = self.learn.model
53
+ self.learn.model = self.learn.model.module
54
+ self.learn.save(name)
55
+ self.learn.model = tmp
56
+ else: self.learn.save(name)
57
+
58
+ def _validate(self, dl=None, callbacks=None, metrics=None, keeped_items=False):
59
+ "Validate on `dl` with potential `callbacks` and `metrics`."
60
+ dl = ifnone(dl, self.learn.data.valid_dl)
61
+ metrics = ifnone(metrics, self.learn.metrics)
62
+ cb_handler = CallbackHandler(ifnone(callbacks, []), metrics)
63
+ cb_handler.on_train_begin(1, None, metrics); cb_handler.on_epoch_begin()
64
+ if keeped_items: cb_handler.state_dict.update(dict(keeped_items=[]))
65
+ val_metrics = validate(self.learn.model, dl, self.loss_func, cb_handler)
66
+ cb_handler.on_epoch_end(val_metrics)
67
+ if keeped_items: return cb_handler.state_dict['keeped_items']
68
+ else: return cb_handler.state_dict['last_metrics']
69
+
70
+ def jump_to_epoch_iter(self, epoch:int, iteration:int)->None:
71
+ try:
72
+ self.learn.load(f'{self.name}_{epoch}_{iteration}', purge=False)
73
+ logging.info(f'Loaded {self.name}_{epoch}_{iteration}')
74
+ except: logging.info(f'Model {self.name}_{epoch}_{iteration} not found.')
75
+
76
+ def on_train_begin(self, n_epochs, **kwargs):
77
+ # TODO: can not write graph here
78
+ # super().on_train_begin(**kwargs)
79
+ self.best = -float('inf')
80
+ self.timer.tic()
81
+ if self.host:
82
+ checkpoint_path = self.learn.path/'checkpoint.yaml'
83
+ if checkpoint_path.exists():
84
+ os.remove(checkpoint_path)
85
+ open(checkpoint_path, 'w').close()
86
+ return {'skip_validate': True, 'iteration':self.start_iters} # disable default validate
87
+
88
+ def on_batch_begin(self, **kwargs:Any)->None:
89
+ self.timer.toc_data()
90
+ super().on_batch_begin(**kwargs)
91
+
92
+ def on_batch_end(self, iteration, epoch, last_loss, smooth_loss, train, **kwargs):
93
+ super().on_batch_end(last_loss, iteration, train, **kwargs)
94
+ if iteration == 0: return
95
+
96
+ if iteration % self.loss_iters == 0:
97
+ last_losses = self.learn.loss_func.last_losses
98
+ self._write_sub_loss(iteration=iteration, last_losses=last_losses)
99
+ self.tbwriter.add_scalar(tag=self.metrics_root + 'lr',
100
+ scalar_value=self.opt.lr, global_step=iteration)
101
+
102
+ if iteration % self.show_iters == 0:
103
+ log_str = f'epoch {epoch} iter {iteration}: loss = {last_loss:6.4f}, ' \
104
+ f'smooth loss = {smooth_loss:6.4f}'
105
+ logging.info(log_str)
106
+ # log_str = f'data time = {self.timer.data_diff:.4f}s, runing time = {self.timer.running_diff:.4f}s'
107
+ # logging.info(log_str)
108
+
109
+ if iteration % self.eval_iters == 0:
110
+ # TODO: or remove time to on_epoch_end
111
+ # 1. Record time
112
+ log_str = f'average data time = {self.timer.average_data_time():.4f}s, ' \
113
+ f'average running time = {self.timer.average_running_time():.4f}s'
114
+ logging.info(log_str)
115
+
116
+ # 2. Call validate
117
+ last_metrics = self._validate()
118
+ self.learn.model.train()
119
+ log_str = f'epoch {epoch} iter {iteration}: eval loss = {last_metrics[0]:6.4f}, ' \
120
+ f'ccr = {last_metrics[1]:6.4f}, cwr = {last_metrics[2]:6.4f}, ' \
121
+ f'ted = {last_metrics[3]:6.4f}, ned = {last_metrics[4]:6.4f}, ' \
122
+ f'ted/w = {last_metrics[5]:6.4f}, '
123
+ logging.info(log_str)
124
+ names = ['eval_loss', 'ccr', 'cwr', 'ted', 'ned', 'ted/w']
125
+ self._write_metrics(iteration, names, last_metrics)
126
+
127
+ # 3. Save best model
128
+ current = last_metrics[2]
129
+ if current is not None and current > self.best:
130
+ logging.info(f'Better model found at epoch {epoch}, '\
131
+ f'iter {iteration} with accuracy value: {current:6.4f}.')
132
+ self.best = current
133
+ self._save(f'{self.bestname}')
134
+
135
+ if iteration % self.save_iters == 0 and self.host:
136
+ logging.info(f'Save model {self.name}_{epoch}_{iteration}')
137
+ filename = f'{self.name}_{epoch}_{iteration}'
138
+ self._save(filename)
139
+
140
+ checkpoint_path = self.learn.path/'checkpoint.yaml'
141
+ if not checkpoint_path.exists():
142
+ open(checkpoint_path, 'w').close()
143
+ with open(checkpoint_path, 'r') as file:
144
+ checkpoints = yaml.load(file, Loader=yaml.FullLoader) or dict()
145
+ checkpoints['all_checkpoints'] = (
146
+ checkpoints.get('all_checkpoints') or list())
147
+ checkpoints['all_checkpoints'].insert(0, filename)
148
+ if len(checkpoints['all_checkpoints']) > self.checpoint_keep_num:
149
+ removed_checkpoint = checkpoints['all_checkpoints'].pop()
150
+ removed_checkpoint = self.learn.path/self.learn.model_dir/f'{removed_checkpoint}.pth'
151
+ os.remove(removed_checkpoint)
152
+ checkpoints['current_checkpoint'] = filename
153
+ with open(checkpoint_path, 'w') as file:
154
+ yaml.dump(checkpoints, file)
155
+
156
+
157
+ self.timer.toc_running()
158
+
159
+ def on_train_end(self, **kwargs):
160
+ #self.learn.load(f'{self.bestname}', purge=False)
161
+ pass
162
+
163
+ def on_epoch_end(self, last_metrics:MetricsList, iteration:int, **kwargs)->None:
164
+ self._write_embedding(iteration=iteration)
165
+
166
+
167
+ class TextAccuracy(Callback):
168
+ _names = ['ccr', 'cwr', 'ted', 'ned', 'ted/w']
169
+ def __init__(self, charset_path, max_length, case_sensitive, model_eval):
170
+ self.charset_path = charset_path
171
+ self.max_length = max_length
172
+ self.case_sensitive = case_sensitive
173
+ self.charset = CharsetMapper(charset_path, self.max_length)
174
+ self.names = self._names
175
+
176
+ self.model_eval = model_eval or 'alignment'
177
+ assert self.model_eval in ['vision', 'language', 'alignment']
178
+
179
+ def on_epoch_begin(self, **kwargs):
180
+ self.total_num_char = 0.
181
+ self.total_num_word = 0.
182
+ self.correct_num_char = 0.
183
+ self.correct_num_word = 0.
184
+ self.total_ed = 0.
185
+ self.total_ned = 0.
186
+
187
+ def _get_output(self, last_output):
188
+ if isinstance(last_output, (tuple, list)):
189
+ for res in last_output:
190
+ if res['name'] == self.model_eval: output = res
191
+ else: output = last_output
192
+ return output
193
+
194
+ def _update_output(self, last_output, items):
195
+ if isinstance(last_output, (tuple, list)):
196
+ for res in last_output:
197
+ if res['name'] == self.model_eval: res.update(items)
198
+ else: last_output.update(items)
199
+ return last_output
200
+
201
+ def on_batch_end(self, last_output, last_target, **kwargs):
202
+ output = self._get_output(last_output)
203
+ logits, pt_lengths = output['logits'], output['pt_lengths']
204
+ pt_text, pt_scores, pt_lengths_ = self.decode(logits)
205
+ assert (pt_lengths == pt_lengths_).all(), f'{pt_lengths} != {pt_lengths_} for {pt_text}'
206
+ last_output = self._update_output(last_output, {'pt_text':pt_text, 'pt_scores':pt_scores})
207
+
208
+ pt_text = [self.charset.trim(t) for t in pt_text]
209
+ label = last_target[0]
210
+ if label.dim() == 3: label = label.argmax(dim=-1) # one-hot label
211
+ gt_text = [self.charset.get_text(l, trim=True) for l in label]
212
+
213
+ for i in range(len(gt_text)):
214
+ if not self.case_sensitive:
215
+ gt_text[i], pt_text[i] = gt_text[i].lower(), pt_text[i].lower()
216
+ distance = ed.eval(gt_text[i], pt_text[i])
217
+ self.total_ed += distance
218
+ self.total_ned += float(distance) / max(len(gt_text[i]), 1)
219
+
220
+ if gt_text[i] == pt_text[i]:
221
+ self.correct_num_word += 1
222
+ self.total_num_word += 1
223
+
224
+ for j in range(min(len(gt_text[i]), len(pt_text[i]))):
225
+ if gt_text[i][j] == pt_text[i][j]:
226
+ self.correct_num_char += 1
227
+ self.total_num_char += len(gt_text[i])
228
+
229
+ return {'last_output': last_output}
230
+
231
+ def on_epoch_end(self, last_metrics, **kwargs):
232
+ mets = [self.correct_num_char / self.total_num_char,
233
+ self.correct_num_word / self.total_num_word,
234
+ self.total_ed,
235
+ self.total_ned,
236
+ self.total_ed / self.total_num_word]
237
+ return add_metrics(last_metrics, mets)
238
+
239
+ def decode(self, logit):
240
+ """ Greed decode """
241
+ # TODO: test running time and decode on GPU
242
+ out = F.softmax(logit, dim=2)
243
+ pt_text, pt_scores, pt_lengths = [], [], []
244
+ for o in out:
245
+ text = self.charset.get_text(o.argmax(dim=1), padding=False, trim=False)
246
+ text = text.split(self.charset.null_char)[0] # end at end-token
247
+ pt_text.append(text)
248
+ pt_scores.append(o.max(dim=1)[0])
249
+ pt_lengths.append(min(len(text) + 1, self.max_length)) # one for end-token
250
+ pt_scores = torch.stack(pt_scores)
251
+ pt_lengths = pt_scores.new_tensor(pt_lengths, dtype=torch.long)
252
+ return pt_text, pt_scores, pt_lengths
253
+
254
+
255
+ class TopKTextAccuracy(TextAccuracy):
256
+ _names = ['ccr', 'cwr']
257
+ def __init__(self, k, charset_path, max_length, case_sensitive, model_eval):
258
+ self.k = k
259
+ self.charset_path = charset_path
260
+ self.max_length = max_length
261
+ self.case_sensitive = case_sensitive
262
+ self.charset = CharsetMapper(charset_path, self.max_length)
263
+ self.names = self._names
264
+
265
+ def on_epoch_begin(self, **kwargs):
266
+ self.total_num_char = 0.
267
+ self.total_num_word = 0.
268
+ self.correct_num_char = 0.
269
+ self.correct_num_word = 0.
270
+
271
+ def on_batch_end(self, last_output, last_target, **kwargs):
272
+ logits, pt_lengths = last_output['logits'], last_output['pt_lengths']
273
+ gt_labels, gt_lengths = last_target[:]
274
+
275
+ for logit, pt_length, label, length in zip(logits, pt_lengths, gt_labels, gt_lengths):
276
+ word_flag = True
277
+ for i in range(length):
278
+ char_logit = logit[i].topk(self.k)[1]
279
+ char_label = label[i].argmax(-1)
280
+ if char_label in char_logit: self.correct_num_char += 1
281
+ else: word_flag = False
282
+ self.total_num_char += 1
283
+ if pt_length == length and word_flag:
284
+ self.correct_num_word += 1
285
+ self.total_num_word += 1
286
+
287
+ def on_epoch_end(self, last_metrics, **kwargs):
288
+ mets = [self.correct_num_char / self.total_num_char,
289
+ self.correct_num_word / self.total_num_word,
290
+ 0., 0., 0.]
291
+ return add_metrics(last_metrics, mets)
292
+
293
+
294
+ class DumpPrediction(LearnerCallback):
295
+
296
+ def __init__(self, learn, dataset, charset_path, model_eval, image_only=False, debug=False):
297
+ super().__init__(learn=learn)
298
+ self.debug = debug
299
+ self.model_eval = model_eval or 'alignment'
300
+ self.image_only = image_only
301
+ assert self.model_eval in ['vision', 'language', 'alignment']
302
+
303
+ self.dataset, self.root = dataset, Path(self.learn.path)/f'{dataset}-{self.model_eval}'
304
+ self.attn_root = self.root/'attn'
305
+ self.charset = CharsetMapper(charset_path)
306
+ if self.root.exists(): shutil.rmtree(self.root)
307
+ self.root.mkdir(), self.attn_root.mkdir()
308
+
309
+ self.pil = transforms.ToPILImage()
310
+ self.tensor = transforms.ToTensor()
311
+ size = self.learn.data.img_h, self.learn.data.img_w
312
+ self.resize = transforms.Resize(size=size, interpolation=0)
313
+ self.c = 0
314
+
315
+ def on_batch_end(self, last_input, last_output, last_target, **kwargs):
316
+ if isinstance(last_output, (tuple, list)):
317
+ for res in last_output:
318
+ if res['name'] == self.model_eval: pt_text = res['pt_text']
319
+ if res['name'] == 'vision': attn_scores = res['attn_scores'].detach().cpu()
320
+ if res['name'] == self.model_eval: logits = res['logits']
321
+ else:
322
+ pt_text = last_output['pt_text']
323
+ attn_scores = last_output['attn_scores'].detach().cpu()
324
+ logits = last_output['logits']
325
+
326
+ images = last_input[0] if isinstance(last_input, (tuple, list)) else last_input
327
+ images = images.detach().cpu()
328
+ pt_text = [self.charset.trim(t) for t in pt_text]
329
+ gt_label = last_target[0]
330
+ if gt_label.dim() == 3: gt_label = gt_label.argmax(dim=-1) # one-hot label
331
+ gt_text = [self.charset.get_text(l, trim=True) for l in gt_label]
332
+
333
+ prediction, false_prediction = [], []
334
+ for gt, pt, image, attn, logit in zip(gt_text, pt_text, images, attn_scores, logits):
335
+ prediction.append(f'{gt}\t{pt}\n')
336
+ if gt != pt:
337
+ if self.debug:
338
+ scores = torch.softmax(logit, dim=-1)[:max(len(pt), len(gt)) + 1]
339
+ logging.info(f'{self.c} gt {gt}, pt {pt}, logit {logit.shape}, scores {scores.topk(5, dim=-1)}')
340
+ false_prediction.append(f'{gt}\t{pt}\n')
341
+
342
+ image = self.learn.data.denorm(image)
343
+ if not self.image_only:
344
+ image_np = np.array(self.pil(image))
345
+ attn_pil = [self.pil(a) for a in attn[:, None, :, :]]
346
+ attn = [self.tensor(self.resize(a)).repeat(3, 1, 1) for a in attn_pil]
347
+ attn_sum = np.array([np.array(a) for a in attn_pil[:len(pt)]]).sum(axis=0)
348
+ blended_sum = self.tensor(blend_mask(image_np, attn_sum))
349
+ blended = [self.tensor(blend_mask(image_np, np.array(a))) for a in attn_pil]
350
+ save_image = torch.stack([image] + attn + [blended_sum] + blended)
351
+ save_image = save_image.view(2, -1, *save_image.shape[1:])
352
+ save_image = save_image.permute(1, 0, 2, 3, 4).flatten(0, 1)
353
+ vutils.save_image(save_image, self.attn_root/f'{self.c}_{gt}_{pt}.jpg',
354
+ nrow=2, normalize=True, scale_each=True)
355
+ else:
356
+ self.pil(image).save(self.attn_root/f'{self.c}_{gt}_{pt}.jpg')
357
+ self.c += 1
358
+
359
+ with open(self.root/f'{self.model_eval}.txt', 'a') as f: f.writelines(prediction)
360
+ with open(self.root/f'{self.model_eval}-false.txt', 'a') as f: f.writelines(false_prediction)
.gitattributes → configs/.gitattributes RENAMED
File without changes
configs/train_iternet.yaml ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ global:
2
+ name: train-iternet
3
+ phase: train
4
+ stage: train-super
5
+ workdir: workdir
6
+ seed: ~
7
+
8
+ dataset:
9
+ train: {
10
+ roots: ['output_pixelplanet_dataset/'],
11
+ batch_size: 20
12
+ }
13
+ test: {
14
+ roots: ['output_pixelplanet_dataset/'],
15
+ batch_size: 20
16
+ }
17
+ data_aug: True
18
+ multiscales: False
19
+ num_workers: 8
20
+
21
+ training:
22
+ epochs: 1000
23
+ show_iters: 500
24
+ eval_iters: 500
25
+ # save_iters: 1
26
+
27
+ optimizer:
28
+ type: Adam
29
+ true_wd: False
30
+ wd: 0.0
31
+ bn_wd: False
32
+ clip_grad: 20
33
+ lr: 0.0001
34
+ args: {
35
+ betas: !!python/tuple [0.9, 0.999], # for default Adam
36
+ }
37
+ scheduler: {
38
+ periods: [6, 4],
39
+ gamma: 0.1,
40
+ }
41
+
42
+ model:
43
+ name: 'modules.model_iternet.IterNet'
44
+ iter_size: 3
45
+ ensemble: ''
46
+ use_vision: False
47
+ vision: {
48
+ checkpoint: workdir/train-iternet/best-train-iternet.pth,
49
+ loss_weight: 1.,
50
+ attention: 'position',
51
+ backbone: 'transformer',
52
+ backbone_ln: 3,
53
+ iter_size: 3,
54
+ backbone_alpha_d: 0.5,
55
+ }
56
+ # language: {
57
+ # checkpoint: workdir/pretrain-language-model/pretrain-language-model.pth,
58
+ # num_layers: 4,
59
+ # loss_weight: 1.,
60
+ # detach: True,
61
+ # use_self_attn: False
62
+ # }
63
+ alignment: {
64
+ loss_weight: 1.,
65
+ }
dataset.py ADDED
@@ -0,0 +1,278 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ import re
3
+
4
+ import cv2
5
+ import lmdb
6
+ import six
7
+ from fastai.vision import *
8
+ from torchvision import transforms
9
+
10
+ from transforms import CVColorJitter, CVDeterioration, CVGeometry
11
+ from utils import CharsetMapper, onehot
12
+
13
+
14
+ class ImageDataset(Dataset):
15
+ "`ImageDataset` read data from LMDB database."
16
+
17
+ def __init__(self,
18
+ path:PathOrStr,
19
+ is_training:bool=True,
20
+ img_h:int=32,
21
+ img_w:int=100,
22
+ max_length:int=25,
23
+ check_length:bool=True,
24
+ case_sensitive:bool=False,
25
+ charset_path:str='data/charset_36.txt',
26
+ convert_mode:str='RGB',
27
+ data_aug:bool=True,
28
+ deteriorate_ratio:float=0.,
29
+ multiscales:bool=True,
30
+ one_hot_y:bool=True,
31
+ return_idx:bool=False,
32
+ return_raw:bool=False,
33
+ **kwargs):
34
+ self.path, self.name = Path(path), Path(path).name
35
+ assert self.path.is_dir() and self.path.exists(), f"{path} is not a valid directory."
36
+ self.convert_mode, self.check_length = convert_mode, check_length
37
+ self.img_h, self.img_w = img_h, img_w
38
+ self.max_length, self.one_hot_y = max_length, one_hot_y
39
+ self.return_idx, self.return_raw = return_idx, return_raw
40
+ self.case_sensitive, self.is_training = case_sensitive, is_training
41
+ self.data_aug, self.multiscales = data_aug, multiscales
42
+ self.charset = CharsetMapper(charset_path, max_length=max_length+1)
43
+ self.c = self.charset.num_classes
44
+
45
+ self.env = lmdb.open(str(path), readonly=True, lock=False, readahead=False, meminit=False)
46
+ assert self.env, f'Cannot open LMDB dataset from {path}.'
47
+ with self.env.begin(write=False) as txn:
48
+ self.length = int(txn.get('num-samples'.encode()))
49
+
50
+ if self.is_training and self.data_aug:
51
+ self.augment_tfs = transforms.Compose([
52
+ CVGeometry(degrees=45, translate=(0.0, 0.0), scale=(0.5, 2.), shear=(45, 15), distortion=0.5, p=0.5),
53
+ CVDeterioration(var=20, degrees=6, factor=4, p=0.25),
54
+ CVColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.25)
55
+ ])
56
+ self.totensor = transforms.ToTensor()
57
+
58
+ def __len__(self): return self.length
59
+
60
+ def _next_image(self, index):
61
+ next_index = random.randint(0, len(self) - 1)
62
+ return self.get(next_index)
63
+
64
+ def _check_image(self, x, pixels=6):
65
+ if x.size[0] <= pixels or x.size[1] <= pixels: return False
66
+ else: return True
67
+
68
+ def resize_multiscales(self, img, borderType=cv2.BORDER_CONSTANT):
69
+ def _resize_ratio(img, ratio, fix_h=True):
70
+ if ratio * self.img_w < self.img_h:
71
+ if fix_h: trg_h = self.img_h
72
+ else: trg_h = int(ratio * self.img_w)
73
+ trg_w = self.img_w
74
+ else: trg_h, trg_w = self.img_h, int(self.img_h / ratio)
75
+ img = cv2.resize(img, (trg_w, trg_h))
76
+ pad_h, pad_w = (self.img_h - trg_h) / 2, (self.img_w - trg_w) / 2
77
+ top, bottom = math.ceil(pad_h), math.floor(pad_h)
78
+ left, right = math.ceil(pad_w), math.floor(pad_w)
79
+ img = cv2.copyMakeBorder(img, top, bottom, left, right, borderType)
80
+ return img
81
+
82
+ if self.is_training:
83
+ if random.random() < 0.5:
84
+ base, maxh, maxw = self.img_h, self.img_h, self.img_w
85
+ h, w = random.randint(base, maxh), random.randint(base, maxw)
86
+ return _resize_ratio(img, h/w)
87
+ else: return _resize_ratio(img, img.shape[0] / img.shape[1]) # keep aspect ratio
88
+ else: return _resize_ratio(img, img.shape[0] / img.shape[1]) # keep aspect ratio
89
+
90
+ def resize(self, img):
91
+ if self.multiscales: return self.resize_multiscales(img, cv2.BORDER_REPLICATE)
92
+ else: return cv2.resize(img, (self.img_w, self.img_h))
93
+
94
+ def get(self, idx):
95
+ with self.env.begin(write=False) as txn:
96
+ image_key, label_key = f'image-{idx+1:09d}', f'label-{idx+1:09d}'
97
+ try:
98
+ label = str(txn.get(label_key.encode()), 'utf-8') # label
99
+ label = re.sub('[^0-9a-zA-Z]+', '', label)
100
+ if self.check_length and self.max_length > 0:
101
+ if len(label) > self.max_length or len(label) <= 0:
102
+ #logging.info(f'Long or short text image is found: {self.name}, {idx}, {label}, {len(label)}')
103
+ return self._next_image(idx)
104
+ label = label[:self.max_length]
105
+
106
+ imgbuf = txn.get(image_key.encode()) # image
107
+ buf = six.BytesIO()
108
+ buf.write(imgbuf)
109
+ buf.seek(0)
110
+ with warnings.catch_warnings():
111
+ warnings.simplefilter("ignore", UserWarning) # EXIF warning from TiffPlugin
112
+ image = PIL.Image.open(buf).convert(self.convert_mode)
113
+ if self.is_training and not self._check_image(image):
114
+ #logging.info(f'Invalid image is found: {self.name}, {idx}, {label}, {len(label)}')
115
+ return self._next_image(idx)
116
+ except:
117
+ import traceback
118
+ traceback.print_exc()
119
+ logging.info(f'Corrupted image is found: {self.name}, {idx}, {label}, {len(label)}')
120
+ return self._next_image(idx)
121
+ return image, label, idx
122
+
123
+ def _process_training(self, image):
124
+ if self.data_aug: image = self.augment_tfs(image)
125
+ image = self.resize(np.array(image))
126
+ return image
127
+
128
+ def _process_test(self, image):
129
+ return self.resize(np.array(image)) # TODO:move is_training to here
130
+
131
+ def __getitem__(self, idx):
132
+ image, text, idx_new = self.get(idx)
133
+ if not self.is_training: assert idx == idx_new, f'idx {idx} != idx_new {idx_new} during testing.'
134
+
135
+ if self.is_training: image = self._process_training(image)
136
+ else: image = self._process_test(image)
137
+ if self.return_raw: return image, text
138
+ image = self.totensor(image)
139
+
140
+ length = tensor(len(text) + 1).to(dtype=torch.long) # one for end token
141
+ label = self.charset.get_labels(text, case_sensitive=self.case_sensitive)
142
+ label = tensor(label).to(dtype=torch.long)
143
+ if self.one_hot_y: label = onehot(label, self.charset.num_classes)
144
+
145
+ if self.return_idx: y = [label, length, idx_new]
146
+ else: y = [label, length]
147
+ return image, y
148
+
149
+
150
+ class TextDataset(Dataset):
151
+ def __init__(self,
152
+ path:PathOrStr,
153
+ delimiter:str='\t',
154
+ max_length:int=25,
155
+ charset_path:str='data/charset_36.txt',
156
+ case_sensitive=False,
157
+ one_hot_x=True,
158
+ one_hot_y=True,
159
+ is_training=True,
160
+ smooth_label=False,
161
+ smooth_factor=0.2,
162
+ use_sm=False,
163
+ **kwargs):
164
+ self.path = Path(path)
165
+ self.case_sensitive, self.use_sm = case_sensitive, use_sm
166
+ self.smooth_factor, self.smooth_label = smooth_factor, smooth_label
167
+ self.charset = CharsetMapper(charset_path, max_length=max_length+1)
168
+ self.one_hot_x, self.one_hot_y, self.is_training = one_hot_x, one_hot_y, is_training
169
+ if self.is_training and self.use_sm: self.sm = SpellingMutation(charset=self.charset)
170
+
171
+ dtype = {'inp': str, 'gt': str}
172
+ self.df = pd.read_csv(self.path, dtype=dtype, delimiter=delimiter, na_filter=False)
173
+ self.inp_col, self.gt_col = 0, 1
174
+
175
+ def __len__(self): return len(self.df)
176
+
177
+ def __getitem__(self, idx):
178
+ text_x = self.df.iloc[idx, self.inp_col]
179
+ text_x = re.sub('[^0-9a-zA-Z]+', '', text_x)
180
+ if not self.case_sensitive: text_x = text_x.lower()
181
+ if self.is_training and self.use_sm: text_x = self.sm(text_x)
182
+
183
+ length_x = tensor(len(text_x) + 1).to(dtype=torch.long) # one for end token
184
+ label_x = self.charset.get_labels(text_x, case_sensitive=self.case_sensitive)
185
+ label_x = tensor(label_x)
186
+ if self.one_hot_x:
187
+ label_x = onehot(label_x, self.charset.num_classes)
188
+ if self.is_training and self.smooth_label:
189
+ label_x = torch.stack([self.prob_smooth_label(l) for l in label_x])
190
+ x = [label_x, length_x]
191
+
192
+ text_y = self.df.iloc[idx, self.gt_col]
193
+ text_y = re.sub('[^0-9a-zA-Z]+', '', text_y)
194
+ if not self.case_sensitive: text_y = text_y.lower()
195
+ length_y = tensor(len(text_y) + 1).to(dtype=torch.long) # one for end token
196
+ label_y = self.charset.get_labels(text_y, case_sensitive=self.case_sensitive)
197
+ label_y = tensor(label_y)
198
+ if self.one_hot_y: label_y = onehot(label_y, self.charset.num_classes)
199
+ y = [label_y, length_y]
200
+
201
+ return x, y
202
+
203
+ def prob_smooth_label(self, one_hot):
204
+ one_hot = one_hot.float()
205
+ delta = torch.rand([]) * self.smooth_factor
206
+ num_classes = len(one_hot)
207
+ noise = torch.rand(num_classes)
208
+ noise = noise / noise.sum() * delta
209
+ one_hot = one_hot * (1 - delta) + noise
210
+ return one_hot
211
+
212
+
213
+ class SpellingMutation(object):
214
+ def __init__(self, pn0=0.7, pn1=0.85, pn2=0.95, pt0=0.7, pt1=0.85, charset=None):
215
+ """
216
+ Args:
217
+ pn0: the prob of not modifying characters is (pn0)
218
+ pn1: the prob of modifying one characters is (pn1 - pn0)
219
+ pn2: the prob of modifying two characters is (pn2 - pn1),
220
+ and three (1 - pn2)
221
+ pt0: the prob of replacing operation is pt0.
222
+ pt1: the prob of inserting operation is (pt1 - pt0),
223
+ and deleting operation is (1 - pt1)
224
+ """
225
+ super().__init__()
226
+ self.pn0, self.pn1, self.pn2 = pn0, pn1, pn2
227
+ self.pt0, self.pt1 = pt0, pt1
228
+ self.charset = charset
229
+ logging.info(f'the probs: pn0={self.pn0}, pn1={self.pn1} ' +
230
+ f'pn2={self.pn2}, pt0={self.pt0}, pt1={self.pt1}')
231
+
232
+ def is_digit(self, text, ratio=0.5):
233
+ length = max(len(text), 1)
234
+ digit_num = sum([t in self.charset.digits for t in text])
235
+ if digit_num / length < ratio: return False
236
+ return True
237
+
238
+ def is_unk_char(self, char):
239
+ # return char == self.charset.unk_char
240
+ return (char not in self.charset.digits) and (char not in self.charset.alphabets)
241
+
242
+ def get_num_to_modify(self, length):
243
+ prob = random.random()
244
+ if prob < self.pn0: num_to_modify = 0
245
+ elif prob < self.pn1: num_to_modify = 1
246
+ elif prob < self.pn2: num_to_modify = 2
247
+ else: num_to_modify = 3
248
+
249
+ if length <= 1: num_to_modify = 0
250
+ elif length >= 2 and length <= 4: num_to_modify = min(num_to_modify, 1)
251
+ else: num_to_modify = min(num_to_modify, length // 2) # smaller than length // 2
252
+ return num_to_modify
253
+
254
+ def __call__(self, text, debug=False):
255
+ if self.is_digit(text): return text
256
+ length = len(text)
257
+ num_to_modify = self.get_num_to_modify(length)
258
+ if num_to_modify <= 0: return text
259
+
260
+ chars = []
261
+ index = np.arange(0, length)
262
+ random.shuffle(index)
263
+ index = index[: num_to_modify]
264
+ if debug: self.index = index
265
+ for i, t in enumerate(text):
266
+ if i not in index: chars.append(t)
267
+ elif self.is_unk_char(t): chars.append(t)
268
+ else:
269
+ prob = random.random()
270
+ if prob < self.pt0: # replace
271
+ chars.append(random.choice(self.charset.alphabets))
272
+ elif prob < self.pt1: # insert
273
+ chars.append(random.choice(self.charset.alphabets))
274
+ chars.append(t)
275
+ else: # delete
276
+ continue
277
+ new_text = ''.join(chars[: self.charset.max_length-1])
278
+ return new_text if len(new_text) >= 1 else text
demo.py ADDED
@@ -0,0 +1,109 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import logging
3
+ import os
4
+ import glob
5
+ import tqdm
6
+ import torch
7
+ import PIL
8
+ import cv2
9
+ import numpy as np
10
+ import torch.nn.functional as F
11
+ from torchvision import transforms
12
+ from utils import Config, Logger, CharsetMapper
13
+
14
+ def get_model(config):
15
+ import importlib
16
+ names = config.model_name.split('.')
17
+ module_name, class_name = '.'.join(names[:-1]), names[-1]
18
+ cls = getattr(importlib.import_module(module_name), class_name)
19
+ model = cls(config)
20
+ logging.info(model)
21
+ model = model.eval()
22
+ return model
23
+
24
+ def preprocess(img, width, height):
25
+ img = cv2.resize(np.array(img), (width, height))
26
+ img = transforms.ToTensor()(img).unsqueeze(0)
27
+ mean = torch.tensor([0.485, 0.456, 0.406])
28
+ std = torch.tensor([0.229, 0.224, 0.225])
29
+ return (img-mean[...,None,None]) / std[...,None,None]
30
+
31
+ def postprocess(output, charset, model_eval):
32
+ def _get_output(last_output, model_eval):
33
+ if isinstance(last_output, (tuple, list)):
34
+ for res in last_output:
35
+ if res['name'] == model_eval: output = res
36
+ else: output = last_output
37
+ return output
38
+
39
+ def _decode(logit):
40
+ """ Greed decode """
41
+ out = F.softmax(logit, dim=2)
42
+ pt_text, pt_scores, pt_lengths = [], [], []
43
+ for o in out:
44
+ text = charset.get_text(o.argmax(dim=1), padding=False, trim=False)
45
+ text = text.split(charset.null_char)[0] # end at end-token
46
+ pt_text.append(text)
47
+ pt_scores.append(o.max(dim=1)[0])
48
+ pt_lengths.append(min(len(text) + 1, charset.max_length)) # one for end-token
49
+ return pt_text, pt_scores, pt_lengths
50
+
51
+ output = _get_output(output, model_eval)
52
+ logits, pt_lengths = output['logits'], output['pt_lengths']
53
+ pt_text, pt_scores, pt_lengths_ = _decode(logits)
54
+
55
+ return pt_text, pt_scores, pt_lengths_
56
+
57
+ def load(model, file, device=None, strict=True):
58
+ if device is None: device = 'cpu'
59
+ elif isinstance(device, int): device = torch.device('cuda', device)
60
+ assert os.path.isfile(file)
61
+ state = torch.load(file, map_location=device)
62
+ if set(state.keys()) == {'model', 'opt'}:
63
+ state = state['model']
64
+ model.load_state_dict(state, strict=strict)
65
+ return model
66
+
67
+ def main():
68
+ parser = argparse.ArgumentParser()
69
+ parser.add_argument('--config', type=str, default='configs/train_iternet.yaml',
70
+ help='path to config file')
71
+ parser.add_argument('--input', type=str, default='figures/demo')
72
+ parser.add_argument('--cuda', type=int, default=-1)
73
+ parser.add_argument('--checkpoint', type=str, default='workdir/train-iternet/best-train-iternet.pth')
74
+ parser.add_argument('--model_eval', type=str, default='alignment',
75
+ choices=['alignment', 'vision', 'language'])
76
+ args = parser.parse_args()
77
+ config = Config(args.config)
78
+ if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
79
+ if args.model_eval is not None: config.model_eval = args.model_eval
80
+ config.global_phase = 'test'
81
+ config.model_vision_checkpoint, config.model_language_checkpoint = None, None
82
+ device = 'cpu' if args.cuda < 0 else f'cuda:{args.cuda}'
83
+
84
+ Logger.init(config.global_workdir, config.global_name, config.global_phase)
85
+ Logger.enable_file()
86
+ logging.info(config)
87
+
88
+ logging.info('Construct model.')
89
+ model = get_model(config).to(device)
90
+ model = load(model, config.model_checkpoint, device=device)
91
+ charset = CharsetMapper(filename=config.dataset_charset_path,
92
+ max_length=config.dataset_max_length + 1)
93
+
94
+ if os.path.isdir(args.input):
95
+ paths = [os.path.join(args.input, fname) for fname in os.listdir(args.input)]
96
+ else:
97
+ paths = glob.glob(os.path.expanduser(args.input))
98
+ assert paths, "The input path(s) was not found"
99
+ paths = sorted(paths)
100
+ for path in tqdm.tqdm(paths):
101
+ img = PIL.Image.open(path).convert('RGB')
102
+ img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
103
+ img = img.to(device)
104
+ res = model(img)
105
+ pt_text, _, __ = postprocess(res, charset, config.model_eval)
106
+ logging.info(f'{path}: {pt_text[0]}')
107
+
108
+ if __name__ == '__main__':
109
+ main()
losses.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastai.vision import *
2
+
3
+ from modules.model import Model
4
+
5
+
6
+ class MultiLosses(nn.Module):
7
+ def __init__(self, one_hot=True):
8
+ super().__init__()
9
+ self.ce = SoftCrossEntropyLoss() if one_hot else torch.nn.CrossEntropyLoss()
10
+ self.bce = torch.nn.BCELoss()
11
+
12
+ @property
13
+ def last_losses(self):
14
+ return self.losses
15
+
16
+ def _flatten(self, sources, lengths):
17
+ return torch.cat([t[:l] for t, l in zip(sources, lengths)])
18
+
19
+ def _merge_list(self, all_res):
20
+ if not isinstance(all_res, (list, tuple)):
21
+ return all_res
22
+ def merge(items):
23
+ if isinstance(items[0], torch.Tensor): return torch.cat(items, dim=0)
24
+ else: return items[0]
25
+ res = dict()
26
+ for key in all_res[0].keys():
27
+ items = [r[key] for r in all_res]
28
+ res[key] = merge(items)
29
+ return res
30
+
31
+ def _ce_loss(self, output, gt_labels, gt_lengths, idx=None, record=True):
32
+ loss_name = output.get('name')
33
+ pt_logits, weight = output['logits'], output['loss_weight']
34
+
35
+ assert pt_logits.shape[0] % gt_labels.shape[0] == 0
36
+ iter_size = pt_logits.shape[0] // gt_labels.shape[0]
37
+ if iter_size > 1:
38
+ gt_labels = gt_labels.repeat(iter_size, 1, 1)
39
+ gt_lengths = gt_lengths.repeat(iter_size)
40
+ flat_gt_labels = self._flatten(gt_labels, gt_lengths)
41
+ flat_pt_logits = self._flatten(pt_logits, gt_lengths)
42
+
43
+ nll = output.get('nll')
44
+ if nll is not None:
45
+ loss = self.ce(flat_pt_logits, flat_gt_labels, softmax=False) * weight
46
+ else:
47
+ loss = self.ce(flat_pt_logits, flat_gt_labels) * weight
48
+ if record and loss_name is not None: self.losses[f'{loss_name}_loss'] = loss
49
+
50
+ return loss
51
+
52
+ def forward(self, outputs, *args):
53
+ self.losses = {}
54
+ if isinstance(outputs, (tuple, list)):
55
+ outputs = [self._merge_list(o) for o in outputs]
56
+ return sum([self._ce_loss(o, *args) for o in outputs if o['loss_weight'] > 0.])
57
+ else:
58
+ return self._ce_loss(outputs, *args, record=False)
59
+
60
+
61
+ class SoftCrossEntropyLoss(nn.Module):
62
+ def __init__(self, reduction="mean"):
63
+ super().__init__()
64
+ self.reduction = reduction
65
+
66
+ def forward(self, input, target, softmax=True):
67
+ if softmax: log_prob = F.log_softmax(input, dim=-1)
68
+ else: log_prob = torch.log(input)
69
+ loss = -(target * log_prob).sum(dim=-1)
70
+ if self.reduction == "mean": return loss.mean()
71
+ elif self.reduction == "sum": return loss.sum()
72
+ else: return loss
main.py ADDED
@@ -0,0 +1,246 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import logging
3
+ import os
4
+ import random
5
+
6
+ import torch
7
+ from fastai.callbacks.general_sched import GeneralScheduler, TrainingPhase
8
+ from fastai.distributed import *
9
+ from fastai.vision import *
10
+ from torch.backends import cudnn
11
+
12
+ from callbacks import DumpPrediction, IterationCallback, TextAccuracy, TopKTextAccuracy
13
+ from dataset import ImageDataset, TextDataset
14
+ from losses import MultiLosses
15
+ from utils import Config, Logger, MyDataParallel, MyConcatDataset
16
+
17
+
18
+ def _set_random_seed(seed):
19
+ if seed is not None:
20
+ random.seed(seed)
21
+ torch.manual_seed(seed)
22
+ cudnn.deterministic = True
23
+ logging.warning('You have chosen to seed training. '
24
+ 'This will slow down your training!')
25
+
26
+ def _get_training_phases(config, n):
27
+ lr = np.array(config.optimizer_lr)
28
+ periods = config.optimizer_scheduler_periods
29
+ sigma = [config.optimizer_scheduler_gamma ** i for i in range(len(periods))]
30
+ phases = [TrainingPhase(n * periods[i]).schedule_hp('lr', lr * sigma[i])
31
+ for i in range(len(periods))]
32
+ return phases
33
+
34
+ def _get_dataset(ds_type, paths, is_training, config, **kwargs):
35
+ kwargs.update({
36
+ 'img_h': config.dataset_image_height,
37
+ 'img_w': config.dataset_image_width,
38
+ 'max_length': config.dataset_max_length,
39
+ 'case_sensitive': config.dataset_case_sensitive,
40
+ 'charset_path': config.dataset_charset_path,
41
+ 'data_aug': config.dataset_data_aug,
42
+ 'deteriorate_ratio': config.dataset_deteriorate_ratio,
43
+ 'is_training': is_training,
44
+ 'multiscales': config.dataset_multiscales,
45
+ 'one_hot_y': config.dataset_one_hot_y,
46
+ })
47
+ datasets = [ds_type(p, **kwargs) for p in paths]
48
+ if len(datasets) > 1: return MyConcatDataset(datasets)
49
+ else: return datasets[0]
50
+
51
+
52
+ def _get_language_databaunch(config):
53
+ kwargs = {
54
+ 'max_length': config.dataset_max_length,
55
+ 'case_sensitive': config.dataset_case_sensitive,
56
+ 'charset_path': config.dataset_charset_path,
57
+ 'smooth_label': config.dataset_smooth_label,
58
+ 'smooth_factor': config.dataset_smooth_factor,
59
+ 'one_hot_y': config.dataset_one_hot_y,
60
+ 'use_sm': config.dataset_use_sm,
61
+ }
62
+ train_ds = TextDataset(config.dataset_train_roots[0], is_training=True, **kwargs)
63
+ valid_ds = TextDataset(config.dataset_test_roots[0], is_training=False, **kwargs)
64
+ data = DataBunch.create(
65
+ path=train_ds.path,
66
+ train_ds=train_ds,
67
+ valid_ds=valid_ds,
68
+ bs=config.dataset_train_batch_size,
69
+ val_bs=config.dataset_test_batch_size,
70
+ num_workers=config.dataset_num_workers,
71
+ pin_memory=config.dataset_pin_memory)
72
+ logging.info(f'{len(data.train_ds)} training items found.')
73
+ if not data.empty_val:
74
+ logging.info(f'{len(data.valid_ds)} valid items found.')
75
+ return data
76
+
77
+ def _get_databaunch(config):
78
+ # An awkward way to reduce loadding data time during test
79
+ if config.global_phase == 'test': config.dataset_train_roots = config.dataset_test_roots
80
+ train_ds = _get_dataset(ImageDataset, config.dataset_train_roots, True, config)
81
+ valid_ds = _get_dataset(ImageDataset, config.dataset_test_roots, False, config)
82
+ data = ImageDataBunch.create(
83
+ train_ds=train_ds,
84
+ valid_ds=valid_ds,
85
+ bs=config.dataset_train_batch_size,
86
+ val_bs=config.dataset_test_batch_size,
87
+ num_workers=config.dataset_num_workers,
88
+ pin_memory=config.dataset_pin_memory).normalize(imagenet_stats)
89
+ ar_tfm = lambda x: ((x[0], x[1]), x[1]) # auto-regression only for dtd
90
+ data.add_tfm(ar_tfm)
91
+
92
+ logging.info(f'{len(data.train_ds)} training items found.')
93
+ if not data.empty_val:
94
+ logging.info(f'{len(data.valid_ds)} valid items found.')
95
+
96
+ return data
97
+
98
+ def _get_model(config):
99
+ import importlib
100
+ names = config.model_name.split('.')
101
+ module_name, class_name = '.'.join(names[:-1]), names[-1]
102
+ cls = getattr(importlib.import_module(module_name), class_name)
103
+ model = cls(config)
104
+ logging.info(model)
105
+ return model
106
+
107
+
108
+ def _get_learner(config, data, model, local_rank=None):
109
+ strict = ifnone(config.model_strict, True)
110
+ if config.global_stage == 'pretrain-language':
111
+ metrics = [TopKTextAccuracy(
112
+ k=ifnone(config.model_k, 5),
113
+ charset_path=config.dataset_charset_path,
114
+ max_length=config.dataset_max_length + 1,
115
+ case_sensitive=config.dataset_eval_case_sensisitves,
116
+ model_eval=config.model_eval)]
117
+ else:
118
+ metrics = [TextAccuracy(
119
+ charset_path=config.dataset_charset_path,
120
+ max_length=config.dataset_max_length + 1,
121
+ case_sensitive=config.dataset_eval_case_sensisitves,
122
+ model_eval=config.model_eval)]
123
+ opt_type = getattr(torch.optim, config.optimizer_type)
124
+ learner = Learner(data, model, silent=True, model_dir='.',
125
+ true_wd=config.optimizer_true_wd,
126
+ wd=config.optimizer_wd,
127
+ bn_wd=config.optimizer_bn_wd,
128
+ path=config.global_workdir,
129
+ metrics=metrics,
130
+ opt_func=partial(opt_type, **config.optimizer_args or dict()),
131
+ loss_func=MultiLosses(one_hot=config.dataset_one_hot_y))
132
+ learner.split(lambda m: children(m))
133
+
134
+ if config.global_phase == 'train':
135
+ num_replicas = 1 if local_rank is None else torch.distributed.get_world_size()
136
+ phases = _get_training_phases(config, len(learner.data.train_dl)//num_replicas)
137
+ learner.callback_fns += [
138
+ partial(GeneralScheduler, phases=phases),
139
+ partial(GradientClipping, clip=config.optimizer_clip_grad),
140
+ partial(IterationCallback, name=config.global_name,
141
+ show_iters=config.training_show_iters,
142
+ eval_iters=config.training_eval_iters,
143
+ save_iters=config.training_save_iters,
144
+ start_iters=config.training_start_iters,
145
+ stats_iters=config.training_stats_iters)]
146
+ else:
147
+ learner.callbacks += [
148
+ DumpPrediction(learn=learner,
149
+ dataset='-'.join([Path(p).name for p in config.dataset_test_roots]),charset_path=config.dataset_charset_path,
150
+ model_eval=config.model_eval,
151
+ debug=config.global_debug,
152
+ image_only=config.global_image_only)]
153
+
154
+ learner.rank = local_rank
155
+ if local_rank is not None:
156
+ logging.info(f'Set model to distributed with rank {local_rank}.')
157
+ learner.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(learner.model)
158
+ learner.model.to(local_rank)
159
+ learner = learner.to_distributed(local_rank)
160
+
161
+ if torch.cuda.device_count() > 1 and local_rank is None:
162
+ logging.info(f'Use {torch.cuda.device_count()} GPUs.')
163
+ learner.model = MyDataParallel(learner.model)
164
+
165
+ if config.model_checkpoint:
166
+ if Path(config.model_checkpoint).exists():
167
+ with open(config.model_checkpoint, 'rb') as f:
168
+ buffer = io.BytesIO(f.read())
169
+ learner.load(buffer, strict=strict)
170
+ else:
171
+ from distutils.dir_util import copy_tree
172
+ src = Path('/data/fangsc/model')/config.global_name
173
+ trg = Path('/output')/config.global_name
174
+ if src.exists(): copy_tree(str(src), str(trg))
175
+ learner.load(config.model_checkpoint, strict=strict)
176
+ logging.info(f'Read model from {config.model_checkpoint}')
177
+ elif config.global_phase == 'test':
178
+ learner.load(f'best-{config.global_name}', strict=strict)
179
+ logging.info(f'Read model from best-{config.global_name}')
180
+
181
+ if learner.opt_func.func.__name__ == 'Adadelta': # fastai bug, fix after 1.0.60
182
+ learner.fit(epochs=0, lr=config.optimizer_lr)
183
+ learner.opt.mom = 0.
184
+
185
+ return learner
186
+
187
+ def main():
188
+ parser = argparse.ArgumentParser()
189
+ parser.add_argument('--config', type=str, required=True,
190
+ help='path to config file')
191
+ parser.add_argument('--phase', type=str, default=None, choices=['train', 'test'])
192
+ parser.add_argument('--name', type=str, default=None)
193
+ parser.add_argument('--checkpoint', type=str, default=None)
194
+ parser.add_argument('--test_root', type=str, default=None)
195
+ parser.add_argument("--local_rank", type=int, default=None)
196
+ parser.add_argument('--debug', action='store_true', default=None)
197
+ parser.add_argument('--image_only', action='store_true', default=None)
198
+ parser.add_argument('--model_strict', action='store_false', default=None)
199
+ parser.add_argument('--model_eval', type=str, default=None,
200
+ choices=['alignment', 'vision', 'language'])
201
+ args = parser.parse_args()
202
+ config = Config(args.config)
203
+ if args.name is not None: config.global_name = args.name
204
+ if args.phase is not None: config.global_phase = args.phase
205
+ if args.test_root is not None: config.dataset_test_roots = [args.test_root]
206
+ if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
207
+ if args.debug is not None: config.global_debug = args.debug
208
+ if args.image_only is not None: config.global_image_only = args.image_only
209
+ if args.model_eval is not None: config.model_eval = args.model_eval
210
+ if args.model_strict is not None: config.model_strict = args.model_strict
211
+
212
+ Logger.init(config.global_workdir, config.global_name, config.global_phase)
213
+ Logger.enable_file()
214
+ _set_random_seed(config.global_seed)
215
+ logging.info(config)
216
+
217
+ if args.local_rank is not None:
218
+ logging.info(f'Init distribution training at device {args.local_rank}.')
219
+ torch.cuda.set_device(args.local_rank)
220
+ torch.distributed.init_process_group(backend='nccl', init_method='env://')
221
+
222
+ logging.info('Construct dataset.')
223
+ if config.global_stage == 'pretrain-language': data = _get_language_databaunch(config)
224
+ else: data = _get_databaunch(config)
225
+
226
+ logging.info('Construct model.')
227
+ model = _get_model(config)
228
+
229
+ logging.info('Construct learner.')
230
+ learner = _get_learner(config, data, model, args.local_rank)
231
+
232
+ if config.global_phase == 'train':
233
+ logging.info('Start training.')
234
+ learner.fit(epochs=config.training_epochs,
235
+ lr=config.optimizer_lr)
236
+ else:
237
+ logging.info('Start validate')
238
+ last_metrics = learner.validate()
239
+ log_str = f'eval loss = {last_metrics[0]:6.3f}, ' \
240
+ f'ccr = {last_metrics[1]:6.3f}, cwr = {last_metrics[2]:6.3f}, ' \
241
+ f'ted = {last_metrics[3]:6.3f}, ned = {last_metrics[4]:6.0f}, ' \
242
+ f'ted/w = {last_metrics[5]:6.3f}, '
243
+ logging.info(log_str)
244
+
245
+ if __name__ == '__main__':
246
+ main()
modules/__init__.py ADDED
File without changes
modules/attention.py ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ from .transformer import PositionalEncoding
4
+
5
+ class Attention(nn.Module):
6
+ def __init__(self, in_channels=512, max_length=25, n_feature=256):
7
+ super().__init__()
8
+ self.max_length = max_length
9
+
10
+ self.f0_embedding = nn.Embedding(max_length, in_channels)
11
+ self.w0 = nn.Linear(max_length, n_feature)
12
+ self.wv = nn.Linear(in_channels, in_channels)
13
+ self.we = nn.Linear(in_channels, max_length)
14
+
15
+ self.active = nn.Tanh()
16
+ self.softmax = nn.Softmax(dim=2)
17
+
18
+ def forward(self, enc_output):
19
+ enc_output = enc_output.permute(0, 2, 3, 1).flatten(1, 2)
20
+ reading_order = torch.arange(self.max_length, dtype=torch.long, device=enc_output.device)
21
+ reading_order = reading_order.unsqueeze(0).expand(enc_output.size(0), -1) # (S,) -> (B, S)
22
+ reading_order_embed = self.f0_embedding(reading_order) # b,25,512
23
+
24
+ t = self.w0(reading_order_embed.permute(0, 2, 1)) # b,512,256
25
+ t = self.active(t.permute(0, 2, 1) + self.wv(enc_output)) # b,256,512
26
+
27
+ attn = self.we(t) # b,256,25
28
+ attn = self.softmax(attn.permute(0, 2, 1)) # b,25,256
29
+ g_output = torch.bmm(attn, enc_output) # b,25,512
30
+ return g_output, attn.view(*attn.shape[:2], 8, 32)
31
+
32
+
33
+ def encoder_layer(in_c, out_c, k=3, s=2, p=1):
34
+ return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),
35
+ nn.BatchNorm2d(out_c),
36
+ nn.ReLU(True))
37
+
38
+ def decoder_layer(in_c, out_c, k=3, s=1, p=1, mode='nearest', scale_factor=None, size=None):
39
+ align_corners = None if mode=='nearest' else True
40
+ return nn.Sequential(nn.Upsample(size=size, scale_factor=scale_factor,
41
+ mode=mode, align_corners=align_corners),
42
+ nn.Conv2d(in_c, out_c, k, s, p),
43
+ nn.BatchNorm2d(out_c),
44
+ nn.ReLU(True))
45
+
46
+
47
+ class PositionAttention(nn.Module):
48
+ def __init__(self, max_length, in_channels=512, num_channels=64,
49
+ h=8, w=32, mode='nearest', **kwargs):
50
+ super().__init__()
51
+ self.max_length = max_length
52
+ self.k_encoder = nn.Sequential(
53
+ encoder_layer(in_channels, num_channels, s=(1, 2)),
54
+ encoder_layer(num_channels, num_channels, s=(2, 2)),
55
+ encoder_layer(num_channels, num_channels, s=(2, 2)),
56
+ encoder_layer(num_channels, num_channels, s=(2, 2))
57
+ )
58
+ self.k_decoder = nn.Sequential(
59
+ decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
60
+ decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
61
+ decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
62
+ decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)
63
+ )
64
+
65
+ self.pos_encoder = PositionalEncoding(in_channels, dropout=0, max_len=max_length)
66
+ self.project = nn.Linear(in_channels, in_channels)
67
+
68
+ def forward(self, x):
69
+ N, E, H, W = x.size()
70
+ k, v = x, x # (N, E, H, W)
71
+
72
+ # calculate key vector
73
+ features = []
74
+ for i in range(0, len(self.k_encoder)):
75
+ k = self.k_encoder[i](k)
76
+ features.append(k)
77
+ for i in range(0, len(self.k_decoder) - 1):
78
+ k = self.k_decoder[i](k)
79
+ k = k + features[len(self.k_decoder) - 2 - i]
80
+ k = self.k_decoder[-1](k)
81
+
82
+ # calculate query vector
83
+ # TODO q=f(q,k)
84
+ zeros = x.new_zeros((self.max_length, N, E)) # (T, N, E)
85
+ q = self.pos_encoder(zeros) # (T, N, E)
86
+ q = q.permute(1, 0, 2) # (N, T, E)
87
+ q = self.project(q) # (N, T, E)
88
+
89
+ # calculate attention
90
+ attn_scores = torch.bmm(q, k.flatten(2, 3)) # (N, T, (H*W))
91
+ attn_scores = attn_scores / (E ** 0.5)
92
+ attn_scores = torch.softmax(attn_scores, dim=-1)
93
+
94
+ v = v.permute(0, 2, 3, 1).view(N, -1, E) # (N, (H*W), E)
95
+ attn_vecs = torch.bmm(attn_scores, v) # (N, T, E)
96
+
97
+ return attn_vecs, attn_scores.view(N, -1, H, W)
modules/backbone.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ from fastai.vision import *
4
+
5
+ from modules.model import _default_tfmer_cfg
6
+ from modules.resnet import resnet45
7
+ from modules.transformer import (PositionalEncoding,
8
+ TransformerEncoder,
9
+ TransformerEncoderLayer)
10
+
11
+ class ResTranformer(nn.Module):
12
+ def __init__(self, config):
13
+ super().__init__()
14
+ alpha_d = ifnone(config.model_vision_backbone_alpha_d, 1.)
15
+
16
+ self.d_model = ifnone(config.model_vision_d_model, _default_tfmer_cfg['d_model'])
17
+
18
+ self.resnet = resnet45(alpha_d, output_channels=self.d_model)
19
+
20
+ nhead = ifnone(config.model_vision_nhead, _default_tfmer_cfg['nhead'])
21
+ d_inner = ifnone(config.model_vision_d_inner, _default_tfmer_cfg['d_inner'])
22
+ dropout = ifnone(config.model_vision_dropout, _default_tfmer_cfg['dropout'])
23
+ activation = ifnone(config.model_vision_activation, _default_tfmer_cfg['activation'])
24
+ num_layers = ifnone(config.model_vision_backbone_ln, 2)
25
+
26
+ self.pos_encoder = PositionalEncoding(self.d_model, max_len=8*32)
27
+ encoder_layer = TransformerEncoderLayer(d_model=self.d_model, nhead=nhead,
28
+ dim_feedforward=d_inner, dropout=dropout, activation=activation)
29
+ self.transformer = TransformerEncoder(encoder_layer, num_layers)
30
+
31
+ def forward_transformer(self, feature):
32
+ n, c, h, w = feature.shape
33
+ feature = feature.view(n, c, -1).permute(2, 0, 1)
34
+ feature = self.pos_encoder(feature)
35
+ feature = self.transformer(feature)
36
+ feature = feature.permute(1, 2, 0).view(n, c, h, w)
37
+ return feature
38
+
39
+ def forward(self, images, **kwargs):
40
+ feature = self.resnet(images, **kwargs)
41
+ feature = self.forward_transformer(feature)
42
+ return feature
modules/model.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+
4
+ from utils import CharsetMapper
5
+
6
+
7
+ _default_tfmer_cfg = dict(d_model=512, nhead=8, d_inner=2048, # 1024
8
+ dropout=0.1, activation='relu')
9
+
10
+ class Model(nn.Module):
11
+
12
+ def __init__(self, config):
13
+ super().__init__()
14
+ self.max_length = config.dataset_max_length + 1
15
+ self.charset = CharsetMapper(config.dataset_charset_path, max_length=self.max_length)
16
+
17
+ def load(self, source, device='cpu', strict=True):
18
+ state = torch.load(source, map_location=device)
19
+ self.load_state_dict(state['model'], strict=strict)
20
+
21
+ def _get_length(self, logit, dim=-1):
22
+ """ Greed decoder to obtain length from logit"""
23
+ out = (logit.argmax(dim=-1) == self.charset.null_label)
24
+ abn = out.any(dim)
25
+ out = ((out.cumsum(dim) == 1) & out).max(dim)[1]
26
+ out = out + 1 # additional end token
27
+ out = torch.where(abn, out, out.new_tensor(logit.shape[1]))
28
+ return out
29
+
30
+ @staticmethod
31
+ def _get_padding_mask(length, max_length):
32
+ length = length.unsqueeze(-1)
33
+ grid = torch.arange(0, max_length, device=length.device).unsqueeze(0)
34
+ return grid >= length
35
+
36
+ @staticmethod
37
+ def _get_square_subsequent_mask(sz, device, diagonal=0, fw=True):
38
+ r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
39
+ Unmasked positions are filled with float(0.0).
40
+ """
41
+ mask = (torch.triu(torch.ones(sz, sz, device=device), diagonal=diagonal) == 1)
42
+ if fw: mask = mask.transpose(0, 1)
43
+ mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
44
+ return mask
45
+
46
+ @staticmethod
47
+ def _get_location_mask(sz, device=None):
48
+ mask = torch.eye(sz, device=device)
49
+ mask = mask.float().masked_fill(mask == 1, float('-inf'))
50
+ return mask
modules/model_alignment.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ from fastai.vision import *
4
+
5
+ from modules.model import Model, _default_tfmer_cfg
6
+
7
+
8
+ class BaseAlignment(Model):
9
+ def __init__(self, config):
10
+ super().__init__(config)
11
+ d_model = ifnone(config.model_alignment_d_model, _default_tfmer_cfg['d_model'])
12
+
13
+ self.loss_weight = ifnone(config.model_alignment_loss_weight, 1.0)
14
+ self.max_length = config.dataset_max_length + 1 # additional stop token
15
+ self.w_att = nn.Linear(2 * d_model, d_model)
16
+ self.cls = nn.Linear(d_model, self.charset.num_classes)
17
+
18
+ def forward(self, l_feature, v_feature):
19
+ """
20
+ Args:
21
+ l_feature: (N, T, E) where T is length, N is batch size and d is dim of model
22
+ v_feature: (N, T, E) shape the same as l_feature
23
+ l_lengths: (N,)
24
+ v_lengths: (N,)
25
+ """
26
+ f = torch.cat((l_feature, v_feature), dim=2)
27
+ f_att = torch.sigmoid(self.w_att(f))
28
+ output = f_att * v_feature + (1 - f_att) * l_feature
29
+
30
+ logits = self.cls(output) # (N, T, C)
31
+ pt_lengths = self._get_length(logits)
32
+
33
+ return {'logits': logits, 'pt_lengths': pt_lengths, 'loss_weight':self.loss_weight,
34
+ 'name': 'alignment'}
modules/model_iternet.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ from fastai.vision import *
4
+
5
+ from .model_vision import BaseIterVision
6
+ from .model_language import BCNLanguage
7
+ from .model_alignment import BaseAlignment
8
+
9
+ class IterNet(nn.Module):
10
+ def __init__(self, config):
11
+ super().__init__()
12
+ self.iter_size = ifnone(config.model_iter_size, 1)
13
+ self.max_length = config.dataset_max_length + 1 # additional stop token
14
+ self.vision = BaseIterVision(config)
15
+ self.language = BCNLanguage(config)
16
+ self.alignment = BaseAlignment(config)
17
+ self.deep_supervision = ifnone(config.model_deep_supervision, True)
18
+
19
+ def forward(self, images, *args):
20
+ list_v_res = self.vision(images)
21
+ if not isinstance(list_v_res, (list, tuple)):
22
+ list_v_res = [list_v_res]
23
+ all_l_res, all_a_res = [], []
24
+
25
+ for v_res in list_v_res:
26
+ a_res = v_res
27
+ for _ in range(self.iter_size):
28
+ tokens = torch.softmax(a_res['logits'], dim=-1)
29
+ lengths = a_res['pt_lengths']
30
+ lengths.clamp_(2, self.max_length) # TODO:move to langauge model
31
+ l_res = self.language(tokens, lengths)
32
+ all_l_res.append(l_res)
33
+ a_res = self.alignment(l_res['feature'], v_res['feature'])
34
+ all_a_res.append(a_res)
35
+ if self.training and self.deep_supervision:
36
+ return all_a_res, all_l_res, list_v_res
37
+ else:
38
+ return a_res, all_l_res[-1], list_v_res[-1]
modules/model_language.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ import torch.nn as nn
3
+ from fastai.vision import *
4
+
5
+ from modules.model import _default_tfmer_cfg
6
+ from modules.model import Model
7
+ from modules.transformer import (PositionalEncoding,
8
+ TransformerDecoder,
9
+ TransformerDecoderLayer)
10
+
11
+
12
+ class BCNLanguage(Model):
13
+ def __init__(self, config):
14
+ super().__init__(config)
15
+ d_model = ifnone(config.model_language_d_model, _default_tfmer_cfg['d_model'])
16
+ nhead = ifnone(config.model_language_nhead, _default_tfmer_cfg['nhead'])
17
+ d_inner = ifnone(config.model_language_d_inner, _default_tfmer_cfg['d_inner'])
18
+ dropout = ifnone(config.model_language_dropout, _default_tfmer_cfg['dropout'])
19
+ activation = ifnone(config.model_language_activation, _default_tfmer_cfg['activation'])
20
+ num_layers = ifnone(config.model_language_num_layers, 4)
21
+ self.d_model = d_model
22
+ self.detach = ifnone(config.model_language_detach, True)
23
+ self.use_self_attn = ifnone(config.model_language_use_self_attn, False)
24
+ self.loss_weight = ifnone(config.model_language_loss_weight, 1.0)
25
+ self.max_length = config.dataset_max_length + 1 # additional stop token
26
+ self.debug = ifnone(config.global_debug, False)
27
+
28
+ self.proj = nn.Linear(self.charset.num_classes, d_model, False)
29
+ self.token_encoder = PositionalEncoding(d_model, max_len=self.max_length)
30
+ self.pos_encoder = PositionalEncoding(d_model, dropout=0, max_len=self.max_length)
31
+ decoder_layer = TransformerDecoderLayer(d_model, nhead, d_inner, dropout,
32
+ activation, self_attn=self.use_self_attn, debug=self.debug)
33
+ self.model = TransformerDecoder(decoder_layer, num_layers)
34
+
35
+ self.cls = nn.Linear(d_model, self.charset.num_classes)
36
+
37
+ if config.model_language_checkpoint is not None:
38
+ logging.info(f'Read language model from {config.model_language_checkpoint}.')
39
+ self.load(config.model_language_checkpoint)
40
+
41
+ def forward(self, tokens, lengths):
42
+ """
43
+ Args:
44
+ tokens: (N, T, C) where T is length, N is batch size and C is classes number
45
+ lengths: (N,)
46
+ """
47
+ if self.detach: tokens = tokens.detach()
48
+ embed = self.proj(tokens) # (N, T, E)
49
+ embed = embed.permute(1, 0, 2) # (T, N, E)
50
+ embed = self.token_encoder(embed) # (T, N, E)
51
+ padding_mask = self._get_padding_mask(lengths, self.max_length)
52
+
53
+ zeros = embed.new_zeros(*embed.shape)
54
+ qeury = self.pos_encoder(zeros)
55
+ location_mask = self._get_location_mask(self.max_length, tokens.device)
56
+ output = self.model(qeury, embed,
57
+ tgt_key_padding_mask=padding_mask,
58
+ memory_mask=location_mask,
59
+ memory_key_padding_mask=padding_mask) # (T, N, E)
60
+ output = output.permute(1, 0, 2) # (N, T, E)
61
+
62
+ logits = self.cls(output) # (N, T, C)
63
+ pt_lengths = self._get_length(logits)
64
+
65
+ res = {'feature': output, 'logits': logits, 'pt_lengths': pt_lengths,
66
+ 'loss_weight':self.loss_weight, 'name': 'language'}
67
+ return res
modules/model_vision.py ADDED
@@ -0,0 +1,151 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ import torch.nn as nn
3
+ from fastai.vision import *
4
+
5
+ from modules.attention import *
6
+ from modules.backbone import ResTranformer
7
+ from modules.model import Model
8
+ from modules.resnet import resnet45
9
+
10
+
11
+ class BaseVision(Model):
12
+ def __init__(self, config):
13
+ super().__init__(config)
14
+ self.loss_weight = ifnone(config.model_vision_loss_weight, 1.0)
15
+ self.out_channels = ifnone(config.model_vision_d_model, 512)
16
+
17
+ if config.model_vision_backbone == 'transformer':
18
+ self.backbone = ResTranformer(config)
19
+ else: self.backbone = resnet45()
20
+
21
+ if config.model_vision_attention == 'position':
22
+ mode = ifnone(config.model_vision_attention_mode, 'nearest')
23
+ self.attention = PositionAttention(
24
+ in_channels=self.out_channels,
25
+ max_length=config.dataset_max_length + 1, # additional stop token
26
+ mode=mode,
27
+ )
28
+ elif config.model_vision_attention == 'attention':
29
+ self.attention = Attention(
30
+ in_channels=self.out_channels,
31
+ max_length=config.dataset_max_length + 1, # additional stop token
32
+ n_feature=8*32,
33
+ )
34
+ else:
35
+ raise Exception(f'{config.model_vision_attention} is not valid.')
36
+ self.cls = nn.Linear(self.out_channels, self.charset.num_classes)
37
+
38
+ if config.model_vision_checkpoint is not None:
39
+ logging.info(f'Read vision model from {config.model_vision_checkpoint}.')
40
+ self.load(config.model_vision_checkpoint)
41
+
42
+ def _forward(self, b_features):
43
+ attn_vecs, attn_scores = self.attention(b_features) # (N, T, E), (N, T, H, W)
44
+ logits = self.cls(attn_vecs) # (N, T, C)
45
+ pt_lengths = self._get_length(logits)
46
+
47
+ return {'feature': attn_vecs, 'logits': logits, 'pt_lengths': pt_lengths,
48
+ 'attn_scores': attn_scores, 'loss_weight':self.loss_weight, 'name': 'vision', 'b_features':b_features}
49
+
50
+ def forward(self, images, *args, **kwargs):
51
+ features = self.backbone(images, **kwargs) # (N, E, H, W)
52
+ return self._forward(features)
53
+
54
+
55
+ class BaseIterVision(BaseVision):
56
+ def __init__(self, config):
57
+ super().__init__(config)
58
+ assert config.model_vision_backbone == 'transformer'
59
+ self.iter_size = ifnone(config.model_vision_iter_size, 1)
60
+ self.share_weights = ifnone(config.model_vision_share_weights, False)
61
+ self.share_cnns = ifnone(config.model_vision_share_cnns, False)
62
+ self.add_transformer = ifnone(config.model_vision_add_transformer, False)
63
+ self.simple_trans = ifnone(config.model_vision_simple_trans, False)
64
+ self.deep_supervision = ifnone(config.model_vision_deep_supervision, True)
65
+ self.backbones = nn.ModuleList()
66
+ self.trans = nn.ModuleList()
67
+ for i in range(self.iter_size-1):
68
+ B = None if self.share_weights else ResTranformer(config)
69
+ if self.share_cnns:
70
+ del B.resnet
71
+ self.backbones.append(B)
72
+ output_channel = self.out_channels
73
+ if self.add_transformer:
74
+ self.split_sizes = [output_channel]
75
+ elif self.simple_trans:
76
+ # self.split_sizes=[output_channel//16] + [0] * 5
77
+ # self.split_sizes= [output_channel//16, output_channel//16, output_channel//8, output_channel//4, output_channel//2] + [0]
78
+ self.split_sizes= [output_channel//16, output_channel//16, 0, output_channel//4, output_channel//2, output_channel]
79
+ else:
80
+ self.split_sizes=[output_channel//16, output_channel//16, output_channel//8, output_channel//4, output_channel//2, output_channel]
81
+ self.trans.append(nn.Conv2d(output_channel, sum(self.split_sizes), 1))
82
+ torch.nn.init.zeros_(self.trans[-1].weight)
83
+
84
+ if config.model_vision_checkpoint is not None:
85
+ logging.info(f'Read vision model from {config.model_vision_checkpoint}.')
86
+ self.load(config.model_vision_checkpoint)
87
+ cb_init = ifnone(config.model_vision_cb_init, True)
88
+ if cb_init:
89
+ self.cb_init()
90
+
91
+ def load(self, source, device=None, strict=False):
92
+ state = torch.load(source, map_location=device)
93
+ msg = self.load_state_dict(state['model'], strict=strict)
94
+ print(msg)
95
+
96
+ def cb_init(self):
97
+ model_state_dict = self.backbone.state_dict()
98
+
99
+ for m in self.backbones:
100
+ if m:
101
+ print('cb_init')
102
+ msg = m.load_state_dict(model_state_dict, strict=False)
103
+ print(msg)
104
+
105
+ def forward_test(self, images, *args):
106
+ l_feats = self.backbone.resnet(images)
107
+ b_feats = self.backbone.forward_transformer(l_feats)
108
+ cnt = len(self.backbones)
109
+ if cnt == 0:
110
+ v_res = super()._forward(b_feats)
111
+ for B,T in zip(self.backbones, self.trans):
112
+ cnt -= 1
113
+ extra_feats = T(b_feats).split(self.split_sizes, dim=1)
114
+ if self.share_weights:
115
+ v_res = super().forward(images, extra_feats=extra_feats)
116
+ else:
117
+ if self.add_transformer:
118
+ if not self.share_cnns:
119
+ l_feats = B.resnet(images)
120
+ b_feats = B.forward_transformer(extra_feats[-1] + l_feats)
121
+ else:
122
+ b_feats = B(images, extra_feats=extra_feats)
123
+ v_res = super()._forward(b_feats) if cnt==0 else None
124
+ return v_res
125
+
126
+ def forward_train(self, images, *args):
127
+ l_feats = self.backbone.resnet(images)
128
+ b_feats = self.backbone.forward_transformer(l_feats)
129
+ v_res = super()._forward(b_feats)
130
+ # v_res = super().forward(images)
131
+ all_v_res = [v_res]
132
+ for B,T in zip(self.backbones, self.trans):
133
+ extra_feats = T(v_res['b_features']).split(self.split_sizes, dim=1)
134
+ if self.share_weights:
135
+ v_res = super().forward(images, extra_feats=extra_feats)
136
+ else:
137
+ if self.add_transformer:
138
+ if not self.share_cnns:
139
+ l_feats = B.resnet(images)
140
+ b_feats = B.forward_transformer(extra_feats[-1] + l_feats)
141
+ else:
142
+ b_feats = B(images, extra_feats=extra_feats)
143
+ v_res = super()._forward(b_feats)
144
+ all_v_res.append(v_res)
145
+ return all_v_res
146
+
147
+ def forward(self, images, *args):
148
+ if self.training and self.deep_supervision:
149
+ return self.forward_train(images, *args)
150
+ else:
151
+ return self.forward_test(images, *args)
modules/resnet.py ADDED
@@ -0,0 +1,133 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch.nn as nn
4
+ import torch.nn.functional as F
5
+ import torch.utils.model_zoo as model_zoo
6
+ import torch.utils.checkpoint as cp
7
+
8
+ def conv1x1(in_planes, out_planes, stride=1):
9
+ return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
10
+
11
+
12
+ def conv3x3(in_planes, out_planes, stride=1):
13
+ "3x3 convolution with padding"
14
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
15
+ padding=1, bias=False)
16
+
17
+
18
+ class BasicBlock(nn.Module):
19
+ expansion = 1
20
+
21
+ def __init__(self, inplanes, planes, stride=1, downsample=None, with_cp=False):
22
+ super(BasicBlock, self).__init__()
23
+ self.conv1 = conv1x1(inplanes, planes)
24
+ self.bn1 = nn.BatchNorm2d(planes)
25
+ self.relu = nn.ReLU(inplace=True)
26
+ self.conv2 = conv3x3(planes, planes, stride)
27
+ self.bn2 = nn.BatchNorm2d(planes)
28
+ self.downsample = downsample
29
+ self.stride = stride
30
+ self.with_cp = with_cp
31
+
32
+ def forward(self, x):
33
+ def _inner_forward(x):
34
+ residual = x
35
+
36
+ out = self.conv1(x)
37
+ out = self.bn1(out)
38
+ out = self.relu(out)
39
+
40
+ out = self.conv2(out)
41
+ out = self.bn2(out)
42
+
43
+ if self.downsample is not None:
44
+ residual = self.downsample(x)
45
+
46
+ out += residual
47
+
48
+ return out
49
+
50
+ if self.with_cp and x.requires_grad:
51
+ out = cp.checkpoint(_inner_forward, x)
52
+ else:
53
+ out = _inner_forward(x)
54
+
55
+ out = self.relu(out)
56
+
57
+ return out
58
+
59
+
60
+ class ResNet(nn.Module):
61
+
62
+ def __init__(self, block, layers, output_channels=512):
63
+ super(ResNet, self).__init__()
64
+ channels = [output_channels//(2**i) for i in reversed(range(5))]
65
+ self.inplanes = channels[0]
66
+ self.conv1 = nn.Conv2d(3, channels[0], kernel_size=3, stride=1, padding=1,
67
+ bias=False)
68
+ self.bn1 = nn.BatchNorm2d(channels[0])
69
+ self.relu = nn.ReLU(inplace=True)
70
+
71
+ self.layer1 = self._make_layer(block, channels[0], layers[0], stride=2)
72
+ self.layer2 = self._make_layer(block, channels[1], layers[1], stride=1)
73
+ self.layer3 = self._make_layer(block, channels[2], layers[2], stride=2)
74
+ self.layer4 = self._make_layer(block, channels[3], layers[3], stride=1)
75
+ self.layer5 = self._make_layer(block, channels[4], layers[4], stride=1)
76
+
77
+ for m in self.modules():
78
+ if isinstance(m, nn.Conv2d):
79
+ n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
80
+ m.weight.data.normal_(0, math.sqrt(2. / n))
81
+ elif isinstance(m, nn.BatchNorm2d):
82
+ m.weight.data.fill_(1)
83
+ m.bias.data.zero_()
84
+
85
+ def _make_layer(self, block, planes, blocks, stride=1):
86
+ downsample = None
87
+ if stride != 1 or self.inplanes != planes * block.expansion:
88
+ downsample = nn.Sequential(
89
+ nn.Conv2d(self.inplanes, planes * block.expansion,
90
+ kernel_size=1, stride=stride, bias=False),
91
+ nn.BatchNorm2d(planes * block.expansion),
92
+ )
93
+
94
+ layers = []
95
+ layers.append(block(self.inplanes, planes, stride, downsample))
96
+ self.inplanes = planes * block.expansion
97
+ for i in range(1, blocks):
98
+ layers.append(block(self.inplanes, planes))
99
+
100
+ return nn.Sequential(*layers)
101
+
102
+ def forward(self, x, extra_feats=None):
103
+ x = self.conv1(x)
104
+ x = self.bn1(x)
105
+ x = self.relu(x)
106
+ if extra_feats is not None:
107
+ if extra_feats[0].shape[1]>0:
108
+ x = x+F.interpolate(extra_feats[0], x.shape[2:], mode='nearest')
109
+ x = self.layer1(x)
110
+ if extra_feats is not None:
111
+ if extra_feats[1].shape[1]>0:
112
+ x = x+F.interpolate(extra_feats[1], x.shape[2:], mode='nearest')
113
+ x = self.layer2(x)
114
+ if extra_feats is not None:
115
+ if extra_feats[2].shape[1]>0:
116
+ x = x+F.interpolate(extra_feats[2], x.shape[2:], mode='nearest')
117
+ x = self.layer3(x)
118
+ if extra_feats is not None:
119
+ if extra_feats[3].shape[1]>0:
120
+ x = x+F.interpolate(extra_feats[3], x.shape[2:], mode='nearest')
121
+ x = self.layer4(x)
122
+ if extra_feats is not None:
123
+ if extra_feats[4].shape[1]>0:
124
+ x = x+F.interpolate(extra_feats[4], x.shape[2:], mode='nearest')
125
+ x = self.layer5(x)
126
+ if extra_feats is not None:
127
+ if extra_feats[5].shape[1]>0:
128
+ x = x+F.interpolate(extra_feats[5], x.shape[2:], mode='nearest')
129
+ return x
130
+
131
+ def resnet45(alpha_d, output_channels=512):
132
+ layers = [int(round(x*alpha_d)) for x in [3, 4, 6, 6, 3]]
133
+ return ResNet(BasicBlock, layers, output_channels=output_channels)
modules/transformer.py ADDED
@@ -0,0 +1,901 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # pytorch 1.5.0
2
+ import copy
3
+ import math
4
+ import warnings
5
+ from typing import Optional
6
+
7
+ import torch
8
+ import torch.nn as nn
9
+ from torch import Tensor
10
+ from torch.nn import Dropout, LayerNorm, Linear, Module, ModuleList, Parameter
11
+ from torch.nn import functional as F
12
+ from torch.nn.init import constant_, xavier_uniform_
13
+
14
+
15
+ def multi_head_attention_forward(query, # type: Tensor
16
+ key, # type: Tensor
17
+ value, # type: Tensor
18
+ embed_dim_to_check, # type: int
19
+ num_heads, # type: int
20
+ in_proj_weight, # type: Tensor
21
+ in_proj_bias, # type: Tensor
22
+ bias_k, # type: Optional[Tensor]
23
+ bias_v, # type: Optional[Tensor]
24
+ add_zero_attn, # type: bool
25
+ dropout_p, # type: float
26
+ out_proj_weight, # type: Tensor
27
+ out_proj_bias, # type: Tensor
28
+ training=True, # type: bool
29
+ key_padding_mask=None, # type: Optional[Tensor]
30
+ need_weights=True, # type: bool
31
+ attn_mask=None, # type: Optional[Tensor]
32
+ use_separate_proj_weight=False, # type: bool
33
+ q_proj_weight=None, # type: Optional[Tensor]
34
+ k_proj_weight=None, # type: Optional[Tensor]
35
+ v_proj_weight=None, # type: Optional[Tensor]
36
+ static_k=None, # type: Optional[Tensor]
37
+ static_v=None # type: Optional[Tensor]
38
+ ):
39
+ # type: (...) -> Tuple[Tensor, Optional[Tensor]]
40
+ r"""
41
+ Args:
42
+ query, key, value: map a query and a set of key-value pairs to an output.
43
+ See "Attention Is All You Need" for more details.
44
+ embed_dim_to_check: total dimension of the model.
45
+ num_heads: parallel attention heads.
46
+ in_proj_weight, in_proj_bias: input projection weight and bias.
47
+ bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
48
+ add_zero_attn: add a new batch of zeros to the key and
49
+ value sequences at dim=1.
50
+ dropout_p: probability of an element to be zeroed.
51
+ out_proj_weight, out_proj_bias: the output projection weight and bias.
52
+ training: apply dropout if is ``True``.
53
+ key_padding_mask: if provided, specified padding elements in the key will
54
+ be ignored by the attention. This is an binary mask. When the value is True,
55
+ the corresponding value on the attention layer will be filled with -inf.
56
+ need_weights: output attn_output_weights.
57
+ attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
58
+ the batches while a 3D mask allows to specify a different mask for the entries of each batch.
59
+ use_separate_proj_weight: the function accept the proj. weights for query, key,
60
+ and value in different forms. If false, in_proj_weight will be used, which is
61
+ a combination of q_proj_weight, k_proj_weight, v_proj_weight.
62
+ q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
63
+ static_k, static_v: static key and value used for attention operators.
64
+ Shape:
65
+ Inputs:
66
+ - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
67
+ the embedding dimension.
68
+ - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
69
+ the embedding dimension.
70
+ - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
71
+ the embedding dimension.
72
+ - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
73
+ If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
74
+ will be unchanged. If a BoolTensor is provided, the positions with the
75
+ value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
76
+ - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
77
+ 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
78
+ S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
79
+ positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
80
+ while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
81
+ are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
82
+ is provided, it will be added to the attention weight.
83
+ - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
84
+ N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
85
+ - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
86
+ N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
87
+ Outputs:
88
+ - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
89
+ E is the embedding dimension.
90
+ - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
91
+ L is the target sequence length, S is the source sequence length.
92
+ """
93
+ # if not torch.jit.is_scripting():
94
+ # tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v,
95
+ # out_proj_weight, out_proj_bias)
96
+ # if any([type(t) is not Tensor for t in tens_ops]) and has_torch_function(tens_ops):
97
+ # return handle_torch_function(
98
+ # multi_head_attention_forward, tens_ops, query, key, value,
99
+ # embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias,
100
+ # bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight,
101
+ # out_proj_bias, training=training, key_padding_mask=key_padding_mask,
102
+ # need_weights=need_weights, attn_mask=attn_mask,
103
+ # use_separate_proj_weight=use_separate_proj_weight,
104
+ # q_proj_weight=q_proj_weight, k_proj_weight=k_proj_weight,
105
+ # v_proj_weight=v_proj_weight, static_k=static_k, static_v=static_v)
106
+ tgt_len, bsz, embed_dim = query.size()
107
+ assert embed_dim == embed_dim_to_check
108
+ assert key.size() == value.size()
109
+
110
+ head_dim = embed_dim // num_heads
111
+ assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
112
+ scaling = float(head_dim) ** -0.5
113
+
114
+ if not use_separate_proj_weight:
115
+ if torch.equal(query, key) and torch.equal(key, value):
116
+ # self-attention
117
+ q, k, v = F.linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)
118
+
119
+ elif torch.equal(key, value):
120
+ # encoder-decoder attention
121
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
122
+ _b = in_proj_bias
123
+ _start = 0
124
+ _end = embed_dim
125
+ _w = in_proj_weight[_start:_end, :]
126
+ if _b is not None:
127
+ _b = _b[_start:_end]
128
+ q = F.linear(query, _w, _b)
129
+
130
+ if key is None:
131
+ assert value is None
132
+ k = None
133
+ v = None
134
+ else:
135
+
136
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
137
+ _b = in_proj_bias
138
+ _start = embed_dim
139
+ _end = None
140
+ _w = in_proj_weight[_start:, :]
141
+ if _b is not None:
142
+ _b = _b[_start:]
143
+ k, v = F.linear(key, _w, _b).chunk(2, dim=-1)
144
+
145
+ else:
146
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
147
+ _b = in_proj_bias
148
+ _start = 0
149
+ _end = embed_dim
150
+ _w = in_proj_weight[_start:_end, :]
151
+ if _b is not None:
152
+ _b = _b[_start:_end]
153
+ q = F.linear(query, _w, _b)
154
+
155
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
156
+ _b = in_proj_bias
157
+ _start = embed_dim
158
+ _end = embed_dim * 2
159
+ _w = in_proj_weight[_start:_end, :]
160
+ if _b is not None:
161
+ _b = _b[_start:_end]
162
+ k = F.linear(key, _w, _b)
163
+
164
+ # This is inline in_proj function with in_proj_weight and in_proj_bias
165
+ _b = in_proj_bias
166
+ _start = embed_dim * 2
167
+ _end = None
168
+ _w = in_proj_weight[_start:, :]
169
+ if _b is not None:
170
+ _b = _b[_start:]
171
+ v = F.linear(value, _w, _b)
172
+ else:
173
+ q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)
174
+ len1, len2 = q_proj_weight_non_opt.size()
175
+ assert len1 == embed_dim and len2 == query.size(-1)
176
+
177
+ k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)
178
+ len1, len2 = k_proj_weight_non_opt.size()
179
+ assert len1 == embed_dim and len2 == key.size(-1)
180
+
181
+ v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)
182
+ len1, len2 = v_proj_weight_non_opt.size()
183
+ assert len1 == embed_dim and len2 == value.size(-1)
184
+
185
+ if in_proj_bias is not None:
186
+ q = F.linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])
187
+ k = F.linear(key, k_proj_weight_non_opt, in_proj_bias[embed_dim:(embed_dim * 2)])
188
+ v = F.linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2):])
189
+ else:
190
+ q = F.linear(query, q_proj_weight_non_opt, in_proj_bias)
191
+ k = F.linear(key, k_proj_weight_non_opt, in_proj_bias)
192
+ v = F.linear(value, v_proj_weight_non_opt, in_proj_bias)
193
+ q = q * scaling
194
+
195
+ if attn_mask is not None:
196
+ assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \
197
+ attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \
198
+ 'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)
199
+ if attn_mask.dtype == torch.uint8:
200
+ warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
201
+ attn_mask = attn_mask.to(torch.bool)
202
+
203
+ if attn_mask.dim() == 2:
204
+ attn_mask = attn_mask.unsqueeze(0)
205
+ if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
206
+ raise RuntimeError('The size of the 2D attn_mask is not correct.')
207
+ elif attn_mask.dim() == 3:
208
+ if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
209
+ raise RuntimeError('The size of the 3D attn_mask is not correct.')
210
+ else:
211
+ raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
212
+ # attn_mask's dim is 3 now.
213
+
214
+ # # convert ByteTensor key_padding_mask to bool
215
+ # if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
216
+ # warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
217
+ # key_padding_mask = key_padding_mask.to(torch.bool)
218
+
219
+ if bias_k is not None and bias_v is not None:
220
+ if static_k is None and static_v is None:
221
+ k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
222
+ v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
223
+ if attn_mask is not None:
224
+ attn_mask = pad(attn_mask, (0, 1))
225
+ if key_padding_mask is not None:
226
+ key_padding_mask = pad(key_padding_mask, (0, 1))
227
+ else:
228
+ assert static_k is None, "bias cannot be added to static key."
229
+ assert static_v is None, "bias cannot be added to static value."
230
+ else:
231
+ assert bias_k is None
232
+ assert bias_v is None
233
+
234
+ q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
235
+ if k is not None:
236
+ k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
237
+ if v is not None:
238
+ v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
239
+
240
+ if static_k is not None:
241
+ assert static_k.size(0) == bsz * num_heads
242
+ assert static_k.size(2) == head_dim
243
+ k = static_k
244
+
245
+ if static_v is not None:
246
+ assert static_v.size(0) == bsz * num_heads
247
+ assert static_v.size(2) == head_dim
248
+ v = static_v
249
+
250
+ src_len = k.size(1)
251
+
252
+ if key_padding_mask is not None:
253
+ assert key_padding_mask.size(0) == bsz
254
+ assert key_padding_mask.size(1) == src_len
255
+
256
+ if add_zero_attn:
257
+ src_len += 1
258
+ k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)
259
+ v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)
260
+ if attn_mask is not None:
261
+ attn_mask = pad(attn_mask, (0, 1))
262
+ if key_padding_mask is not None:
263
+ key_padding_mask = pad(key_padding_mask, (0, 1))
264
+
265
+ attn_output_weights = torch.bmm(q, k.transpose(1, 2))
266
+ assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
267
+
268
+ if attn_mask is not None:
269
+ if attn_mask.dtype == torch.bool:
270
+ attn_output_weights.masked_fill_(attn_mask, float('-inf'))
271
+ else:
272
+ attn_output_weights += attn_mask
273
+
274
+
275
+ if key_padding_mask is not None:
276
+ attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
277
+ attn_output_weights = attn_output_weights.masked_fill(
278
+ key_padding_mask.unsqueeze(1).unsqueeze(2),
279
+ float('-inf'),
280
+ )
281
+ attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)
282
+
283
+ attn_output_weights = F.softmax(
284
+ attn_output_weights, dim=-1)
285
+ attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)
286
+
287
+ attn_output = torch.bmm(attn_output_weights, v)
288
+ assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
289
+ attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
290
+ attn_output = F.linear(attn_output, out_proj_weight, out_proj_bias)
291
+
292
+ if need_weights:
293
+ # average attention weights over heads
294
+ attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
295
+ return attn_output, attn_output_weights.sum(dim=1) / num_heads
296
+ else:
297
+ return attn_output, None
298
+
299
+ class MultiheadAttention(Module):
300
+ r"""Allows the model to jointly attend to information
301
+ from different representation subspaces.
302
+ See reference: Attention Is All You Need
303
+ .. math::
304
+ \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
305
+ \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
306
+ Args:
307
+ embed_dim: total dimension of the model.
308
+ num_heads: parallel attention heads.
309
+ dropout: a Dropout layer on attn_output_weights. Default: 0.0.
310
+ bias: add bias as module parameter. Default: True.
311
+ add_bias_kv: add bias to the key and value sequences at dim=0.
312
+ add_zero_attn: add a new batch of zeros to the key and
313
+ value sequences at dim=1.
314
+ kdim: total number of features in key. Default: None.
315
+ vdim: total number of features in value. Default: None.
316
+ Note: if kdim and vdim are None, they will be set to embed_dim such that
317
+ query, key, and value have the same number of features.
318
+ Examples::
319
+ >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
320
+ >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
321
+ """
322
+ # __annotations__ = {
323
+ # 'bias_k': torch._jit_internal.Optional[torch.Tensor],
324
+ # 'bias_v': torch._jit_internal.Optional[torch.Tensor],
325
+ # }
326
+ __constants__ = ['q_proj_weight', 'k_proj_weight', 'v_proj_weight', 'in_proj_weight']
327
+
328
+ def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
329
+ super(MultiheadAttention, self).__init__()
330
+ self.embed_dim = embed_dim
331
+ self.kdim = kdim if kdim is not None else embed_dim
332
+ self.vdim = vdim if vdim is not None else embed_dim
333
+ self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
334
+
335
+ self.num_heads = num_heads
336
+ self.dropout = dropout
337
+ self.head_dim = embed_dim // num_heads
338
+ assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
339
+
340
+ if self._qkv_same_embed_dim is False:
341
+ self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
342
+ self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
343
+ self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
344
+ self.register_parameter('in_proj_weight', None)
345
+ else:
346
+ self.in_proj_weight = Parameter(torch.empty(3 * embed_dim, embed_dim))
347
+ self.register_parameter('q_proj_weight', None)
348
+ self.register_parameter('k_proj_weight', None)
349
+ self.register_parameter('v_proj_weight', None)
350
+
351
+ if bias:
352
+ self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))
353
+ else:
354
+ self.register_parameter('in_proj_bias', None)
355
+ self.out_proj = Linear(embed_dim, embed_dim, bias=bias)
356
+
357
+ if add_bias_kv:
358
+ self.bias_k = Parameter(torch.empty(1, 1, embed_dim))
359
+ self.bias_v = Parameter(torch.empty(1, 1, embed_dim))
360
+ else:
361
+ self.bias_k = self.bias_v = None
362
+
363
+ self.add_zero_attn = add_zero_attn
364
+
365
+ self._reset_parameters()
366
+
367
+ def _reset_parameters(self):
368
+ if self._qkv_same_embed_dim:
369
+ xavier_uniform_(self.in_proj_weight)
370
+ else:
371
+ xavier_uniform_(self.q_proj_weight)
372
+ xavier_uniform_(self.k_proj_weight)
373
+ xavier_uniform_(self.v_proj_weight)
374
+
375
+ if self.in_proj_bias is not None:
376
+ constant_(self.in_proj_bias, 0.)
377
+ constant_(self.out_proj.bias, 0.)
378
+ if self.bias_k is not None:
379
+ xavier_normal_(self.bias_k)
380
+ if self.bias_v is not None:
381
+ xavier_normal_(self.bias_v)
382
+
383
+ def __setstate__(self, state):
384
+ # Support loading old MultiheadAttention checkpoints generated by v1.1.0
385
+ if '_qkv_same_embed_dim' not in state:
386
+ state['_qkv_same_embed_dim'] = True
387
+
388
+ super(MultiheadAttention, self).__setstate__(state)
389
+
390
+ def forward(self, query, key, value, key_padding_mask=None,
391
+ need_weights=True, attn_mask=None):
392
+ # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
393
+ r"""
394
+ Args:
395
+ query, key, value: map a query and a set of key-value pairs to an output.
396
+ See "Attention Is All You Need" for more details.
397
+ key_padding_mask: if provided, specified padding elements in the key will
398
+ be ignored by the attention. This is an binary mask. When the value is True,
399
+ the corresponding value on the attention layer will be filled with -inf.
400
+ need_weights: output attn_output_weights.
401
+ attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
402
+ the batches while a 3D mask allows to specify a different mask for the entries of each batch.
403
+ Shape:
404
+ - Inputs:
405
+ - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
406
+ the embedding dimension.
407
+ - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
408
+ the embedding dimension.
409
+ - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
410
+ the embedding dimension.
411
+ - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
412
+ If a ByteTensor is provided, the non-zero positions will be ignored while the position
413
+ with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
414
+ value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
415
+ - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
416
+ 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
417
+ S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
418
+ positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
419
+ while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
420
+ is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
421
+ is provided, it will be added to the attention weight.
422
+ - Outputs:
423
+ - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
424
+ E is the embedding dimension.
425
+ - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
426
+ L is the target sequence length, S is the source sequence length.
427
+ """
428
+ if not self._qkv_same_embed_dim:
429
+ return multi_head_attention_forward(
430
+ query, key, value, self.embed_dim, self.num_heads,
431
+ self.in_proj_weight, self.in_proj_bias,
432
+ self.bias_k, self.bias_v, self.add_zero_attn,
433
+ self.dropout, self.out_proj.weight, self.out_proj.bias,
434
+ training=self.training,
435
+ key_padding_mask=key_padding_mask, need_weights=need_weights,
436
+ attn_mask=attn_mask, use_separate_proj_weight=True,
437
+ q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
438
+ v_proj_weight=self.v_proj_weight)
439
+ else:
440
+ return multi_head_attention_forward(
441
+ query, key, value, self.embed_dim, self.num_heads,
442
+ self.in_proj_weight, self.in_proj_bias,
443
+ self.bias_k, self.bias_v, self.add_zero_attn,
444
+ self.dropout, self.out_proj.weight, self.out_proj.bias,
445
+ training=self.training,
446
+ key_padding_mask=key_padding_mask, need_weights=need_weights,
447
+ attn_mask=attn_mask)
448
+
449
+
450
+ class Transformer(Module):
451
+ r"""A transformer model. User is able to modify the attributes as needed. The architecture
452
+ is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
453
+ Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
454
+ Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
455
+ Processing Systems, pages 6000-6010. Users can build the BERT(https://arxiv.org/abs/1810.04805)
456
+ model with corresponding parameters.
457
+
458
+ Args:
459
+ d_model: the number of expected features in the encoder/decoder inputs (default=512).
460
+ nhead: the number of heads in the multiheadattention models (default=8).
461
+ num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
462
+ num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
463
+ dim_feedforward: the dimension of the feedforward network model (default=2048).
464
+ dropout: the dropout value (default=0.1).
465
+ activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
466
+ custom_encoder: custom encoder (default=None).
467
+ custom_decoder: custom decoder (default=None).
468
+
469
+ Examples::
470
+ >>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
471
+ >>> src = torch.rand((10, 32, 512))
472
+ >>> tgt = torch.rand((20, 32, 512))
473
+ >>> out = transformer_model(src, tgt)
474
+
475
+ Note: A full example to apply nn.Transformer module for the word language model is available in
476
+ https://github.com/pytorch/examples/tree/master/word_language_model
477
+ """
478
+
479
+ def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
480
+ num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
481
+ activation="relu", custom_encoder=None, custom_decoder=None):
482
+ super(Transformer, self).__init__()
483
+
484
+ if custom_encoder is not None:
485
+ self.encoder = custom_encoder
486
+ else:
487
+ encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
488
+ encoder_norm = LayerNorm(d_model)
489
+ self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
490
+
491
+ if custom_decoder is not None:
492
+ self.decoder = custom_decoder
493
+ else:
494
+ decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation)
495
+ decoder_norm = LayerNorm(d_model)
496
+ self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
497
+
498
+ self._reset_parameters()
499
+
500
+ self.d_model = d_model
501
+ self.nhead = nhead
502
+
503
+ def forward(self, src, tgt, src_mask=None, tgt_mask=None,
504
+ memory_mask=None, src_key_padding_mask=None,
505
+ tgt_key_padding_mask=None, memory_key_padding_mask=None):
506
+ # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor # noqa
507
+ r"""Take in and process masked source/target sequences.
508
+
509
+ Args:
510
+ src: the sequence to the encoder (required).
511
+ tgt: the sequence to the decoder (required).
512
+ src_mask: the additive mask for the src sequence (optional).
513
+ tgt_mask: the additive mask for the tgt sequence (optional).
514
+ memory_mask: the additive mask for the encoder output (optional).
515
+ src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).
516
+ tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).
517
+ memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).
518
+
519
+ Shape:
520
+ - src: :math:`(S, N, E)`.
521
+ - tgt: :math:`(T, N, E)`.
522
+ - src_mask: :math:`(S, S)`.
523
+ - tgt_mask: :math:`(T, T)`.
524
+ - memory_mask: :math:`(T, S)`.
525
+ - src_key_padding_mask: :math:`(N, S)`.
526
+ - tgt_key_padding_mask: :math:`(N, T)`.
527
+ - memory_key_padding_mask: :math:`(N, S)`.
528
+
529
+ Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked
530
+ positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
531
+ while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
532
+ are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
533
+ is provided, it will be added to the attention weight.
534
+ [src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
535
+ the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero
536
+ positions will be unchanged. If a BoolTensor is provided, the positions with the
537
+ value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
538
+
539
+ - output: :math:`(T, N, E)`.
540
+
541
+ Note: Due to the multi-head attention architecture in the transformer model,
542
+ the output sequence length of a transformer is same as the input sequence
543
+ (i.e. target) length of the decode.
544
+
545
+ where S is the source sequence length, T is the target sequence length, N is the
546
+ batch size, E is the feature number
547
+
548
+ Examples:
549
+ >>> output = transformer_model(src, tgt, src_mask=src_mask, tgt_mask=tgt_mask)
550
+ """
551
+
552
+ if src.size(1) != tgt.size(1):
553
+ raise RuntimeError("the batch number of src and tgt must be equal")
554
+
555
+ if src.size(2) != self.d_model or tgt.size(2) != self.d_model:
556
+ raise RuntimeError("the feature number of src and tgt must be equal to d_model")
557
+
558
+ memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
559
+ output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
560
+ tgt_key_padding_mask=tgt_key_padding_mask,
561
+ memory_key_padding_mask=memory_key_padding_mask)
562
+ return output
563
+
564
+ def generate_square_subsequent_mask(self, sz):
565
+ r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
566
+ Unmasked positions are filled with float(0.0).
567
+ """
568
+ mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
569
+ mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
570
+ return mask
571
+
572
+ def _reset_parameters(self):
573
+ r"""Initiate parameters in the transformer model."""
574
+
575
+ for p in self.parameters():
576
+ if p.dim() > 1:
577
+ xavier_uniform_(p)
578
+
579
+
580
+ class TransformerEncoder(Module):
581
+ r"""TransformerEncoder is a stack of N encoder layers
582
+
583
+ Args:
584
+ encoder_layer: an instance of the TransformerEncoderLayer() class (required).
585
+ num_layers: the number of sub-encoder-layers in the encoder (required).
586
+ norm: the layer normalization component (optional).
587
+
588
+ Examples::
589
+ >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
590
+ >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
591
+ >>> src = torch.rand(10, 32, 512)
592
+ >>> out = transformer_encoder(src)
593
+ """
594
+ __constants__ = ['norm']
595
+
596
+ def __init__(self, encoder_layer, num_layers, norm=None):
597
+ super(TransformerEncoder, self).__init__()
598
+ self.layers = _get_clones(encoder_layer, num_layers)
599
+ self.num_layers = num_layers
600
+ self.norm = norm
601
+
602
+ def forward(self, src, mask=None, src_key_padding_mask=None):
603
+ # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
604
+ r"""Pass the input through the encoder layers in turn.
605
+
606
+ Args:
607
+ src: the sequence to the encoder (required).
608
+ mask: the mask for the src sequence (optional).
609
+ src_key_padding_mask: the mask for the src keys per batch (optional).
610
+
611
+ Shape:
612
+ see the docs in Transformer class.
613
+ """
614
+ output = src
615
+
616
+ for i, mod in enumerate(self.layers):
617
+ output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
618
+
619
+ if self.norm is not None:
620
+ output = self.norm(output)
621
+
622
+ return output
623
+
624
+
625
+ class TransformerDecoder(Module):
626
+ r"""TransformerDecoder is a stack of N decoder layers
627
+
628
+ Args:
629
+ decoder_layer: an instance of the TransformerDecoderLayer() class (required).
630
+ num_layers: the number of sub-decoder-layers in the decoder (required).
631
+ norm: the layer normalization component (optional).
632
+
633
+ Examples::
634
+ >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
635
+ >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
636
+ >>> memory = torch.rand(10, 32, 512)
637
+ >>> tgt = torch.rand(20, 32, 512)
638
+ >>> out = transformer_decoder(tgt, memory)
639
+ """
640
+ __constants__ = ['norm']
641
+
642
+ def __init__(self, decoder_layer, num_layers, norm=None):
643
+ super(TransformerDecoder, self).__init__()
644
+ self.layers = _get_clones(decoder_layer, num_layers)
645
+ self.num_layers = num_layers
646
+ self.norm = norm
647
+
648
+ def forward(self, tgt, memory, memory2=None, tgt_mask=None,
649
+ memory_mask=None, memory_mask2=None, tgt_key_padding_mask=None,
650
+ memory_key_padding_mask=None, memory_key_padding_mask2=None):
651
+ # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor
652
+ r"""Pass the inputs (and mask) through the decoder layer in turn.
653
+
654
+ Args:
655
+ tgt: the sequence to the decoder (required).
656
+ memory: the sequence from the last layer of the encoder (required).
657
+ tgt_mask: the mask for the tgt sequence (optional).
658
+ memory_mask: the mask for the memory sequence (optional).
659
+ tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
660
+ memory_key_padding_mask: the mask for the memory keys per batch (optional).
661
+
662
+ Shape:
663
+ see the docs in Transformer class.
664
+ """
665
+ output = tgt
666
+
667
+ for mod in self.layers:
668
+ output = mod(output, memory, memory2=memory2, tgt_mask=tgt_mask,
669
+ memory_mask=memory_mask, memory_mask2=memory_mask2,
670
+ tgt_key_padding_mask=tgt_key_padding_mask,
671
+ memory_key_padding_mask=memory_key_padding_mask,
672
+ memory_key_padding_mask2=memory_key_padding_mask2)
673
+
674
+ if self.norm is not None:
675
+ output = self.norm(output)
676
+
677
+ return output
678
+
679
+ class TransformerEncoderLayer(Module):
680
+ r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
681
+ This standard encoder layer is based on the paper "Attention Is All You Need".
682
+ Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
683
+ Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
684
+ Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
685
+ in a different way during application.
686
+
687
+ Args:
688
+ d_model: the number of expected features in the input (required).
689
+ nhead: the number of heads in the multiheadattention models (required).
690
+ dim_feedforward: the dimension of the feedforward network model (default=2048).
691
+ dropout: the dropout value (default=0.1).
692
+ activation: the activation function of intermediate layer, relu or gelu (default=relu).
693
+
694
+ Examples::
695
+ >>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
696
+ >>> src = torch.rand(10, 32, 512)
697
+ >>> out = encoder_layer(src)
698
+ """
699
+
700
+ def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
701
+ activation="relu", debug=False):
702
+ super(TransformerEncoderLayer, self).__init__()
703
+ self.debug = debug
704
+ self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
705
+ # Implementation of Feedforward model
706
+ self.linear1 = Linear(d_model, dim_feedforward)
707
+ self.dropout = Dropout(dropout)
708
+ self.linear2 = Linear(dim_feedforward, d_model)
709
+
710
+ self.norm1 = LayerNorm(d_model)
711
+ self.norm2 = LayerNorm(d_model)
712
+ self.dropout1 = Dropout(dropout)
713
+ self.dropout2 = Dropout(dropout)
714
+
715
+ self.activation = _get_activation_fn(activation)
716
+
717
+ def __setstate__(self, state):
718
+ if 'activation' not in state:
719
+ state['activation'] = F.relu
720
+ super(TransformerEncoderLayer, self).__setstate__(state)
721
+
722
+ def forward(self, src, src_mask=None, src_key_padding_mask=None):
723
+ # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
724
+ r"""Pass the input through the encoder layer.
725
+
726
+ Args:
727
+ src: the sequence to the encoder layer (required).
728
+ src_mask: the mask for the src sequence (optional).
729
+ src_key_padding_mask: the mask for the src keys per batch (optional).
730
+
731
+ Shape:
732
+ see the docs in Transformer class.
733
+ """
734
+ src2, attn = self.self_attn(src, src, src, attn_mask=src_mask,
735
+ key_padding_mask=src_key_padding_mask)
736
+ if self.debug: self.attn = attn
737
+ src = src + self.dropout1(src2)
738
+ src = self.norm1(src)
739
+ src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
740
+ src = src + self.dropout2(src2)
741
+ src = self.norm2(src)
742
+
743
+ return src
744
+
745
+
746
+ class TransformerDecoderLayer(Module):
747
+ r"""TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
748
+ This standard decoder layer is based on the paper "Attention Is All You Need".
749
+ Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
750
+ Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
751
+ Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
752
+ in a different way during application.
753
+
754
+ Args:
755
+ d_model: the number of expected features in the input (required).
756
+ nhead: the number of heads in the multiheadattention models (required).
757
+ dim_feedforward: the dimension of the feedforward network model (default=2048).
758
+ dropout: the dropout value (default=0.1).
759
+ activation: the activation function of intermediate layer, relu or gelu (default=relu).
760
+
761
+ Examples::
762
+ >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
763
+ >>> memory = torch.rand(10, 32, 512)
764
+ >>> tgt = torch.rand(20, 32, 512)
765
+ >>> out = decoder_layer(tgt, memory)
766
+ """
767
+
768
+ def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
769
+ activation="relu", self_attn=True, siamese=False, debug=False):
770
+ super(TransformerDecoderLayer, self).__init__()
771
+ self.has_self_attn, self.siamese = self_attn, siamese
772
+ self.debug = debug
773
+ if self.has_self_attn:
774
+ self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
775
+ self.norm1 = LayerNorm(d_model)
776
+ self.dropout1 = Dropout(dropout)
777
+ self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
778
+ # Implementation of Feedforward model
779
+ self.linear1 = Linear(d_model, dim_feedforward)
780
+ self.dropout = Dropout(dropout)
781
+ self.linear2 = Linear(dim_feedforward, d_model)
782
+
783
+ self.norm2 = LayerNorm(d_model)
784
+ self.norm3 = LayerNorm(d_model)
785
+ self.dropout2 = Dropout(dropout)
786
+ self.dropout3 = Dropout(dropout)
787
+ if self.siamese:
788
+ self.multihead_attn2 = MultiheadAttention(d_model, nhead, dropout=dropout)
789
+
790
+ self.activation = _get_activation_fn(activation)
791
+
792
+ def __setstate__(self, state):
793
+ if 'activation' not in state:
794
+ state['activation'] = F.relu
795
+ super(TransformerDecoderLayer, self).__setstate__(state)
796
+
797
+ def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
798
+ tgt_key_padding_mask=None, memory_key_padding_mask=None,
799
+ memory2=None, memory_mask2=None, memory_key_padding_mask2=None):
800
+ # type: (Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Tensor]) -> Tensor
801
+ r"""Pass the inputs (and mask) through the decoder layer.
802
+
803
+ Args:
804
+ tgt: the sequence to the decoder layer (required).
805
+ memory: the sequence from the last layer of the encoder (required).
806
+ tgt_mask: the mask for the tgt sequence (optional).
807
+ memory_mask: the mask for the memory sequence (optional).
808
+ tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
809
+ memory_key_padding_mask: the mask for the memory keys per batch (optional).
810
+
811
+ Shape:
812
+ see the docs in Transformer class.
813
+ """
814
+ if self.has_self_attn:
815
+ tgt2, attn = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
816
+ key_padding_mask=tgt_key_padding_mask)
817
+ tgt = tgt + self.dropout1(tgt2)
818
+ tgt = self.norm1(tgt)
819
+ if self.debug: self.attn = attn
820
+ tgt2, attn2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
821
+ key_padding_mask=memory_key_padding_mask)
822
+ if self.debug: self.attn2 = attn2
823
+
824
+ if self.siamese:
825
+ tgt3, attn3 = self.multihead_attn2(tgt, memory2, memory2, attn_mask=memory_mask2,
826
+ key_padding_mask=memory_key_padding_mask2)
827
+ tgt = tgt + self.dropout2(tgt3)
828
+ if self.debug: self.attn3 = attn3
829
+
830
+ tgt = tgt + self.dropout2(tgt2)
831
+ tgt = self.norm2(tgt)
832
+ tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
833
+ tgt = tgt + self.dropout3(tgt2)
834
+ tgt = self.norm3(tgt)
835
+
836
+ return tgt
837
+
838
+
839
+ def _get_clones(module, N):
840
+ return ModuleList([copy.deepcopy(module) for i in range(N)])
841
+
842
+
843
+ def _get_activation_fn(activation):
844
+ if activation == "relu":
845
+ return F.relu
846
+ elif activation == "gelu":
847
+ return F.gelu
848
+
849
+ raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
850
+
851
+
852
+ class PositionalEncoding(nn.Module):
853
+ r"""Inject some information about the relative or absolute position of the tokens
854
+ in the sequence. The positional encodings have the same dimension as
855
+ the embeddings, so that the two can be summed. Here, we use sine and cosine
856
+ functions of different frequencies.
857
+ .. math::
858
+ \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
859
+ \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
860
+ \text{where pos is the word position and i is the embed idx)
861
+ Args:
862
+ d_model: the embed dim (required).
863
+ dropout: the dropout value (default=0.1).
864
+ max_len: the max. length of the incoming sequence (default=5000).
865
+ Examples:
866
+ >>> pos_encoder = PositionalEncoding(d_model)
867
+ """
868
+
869
+ def __init__(self, d_model, dropout=0.1, max_len=5000):
870
+ super(PositionalEncoding, self).__init__()
871
+ self.dropout = nn.Dropout(p=dropout)
872
+
873
+ pe = torch.zeros(max_len, d_model)
874
+ position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
875
+ div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
876
+ pe[:, 0::2] = torch.sin(position * div_term)
877
+ pe[:, 1::2] = torch.cos(position * div_term)
878
+ pe = pe.unsqueeze(0).transpose(0, 1)
879
+ self.register_buffer('pe', pe)
880
+
881
+ def forward(self, x):
882
+ r"""Inputs of forward function
883
+ Args:
884
+ x: the sequence fed to the positional encoder model (required).
885
+ Shape:
886
+ x: [sequence length, batch size, embed dim]
887
+ output: [sequence length, batch size, embed dim]
888
+ Examples:
889
+ >>> output = pos_encoder(x)
890
+ """
891
+
892
+ x = x + self.pe[:x.size(0), :]
893
+ return self.dropout(x)
894
+
895
+
896
+ if __name__ == '__main__':
897
+ transformer_model = Transformer(nhead=16, num_encoder_layers=12)
898
+ src = torch.rand((10, 32, 512))
899
+ tgt = torch.rand((20, 32, 512))
900
+ out = transformer_model(src, tgt)
901
+ print(out)
packages.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ python3-opencv
requirements.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ fastai==1.0.60
4
+ LMDB
5
+ Pillow
6
+ opencv-python
7
+ tensorboardX
8
+ editdistance
transforms.py ADDED
@@ -0,0 +1,329 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import numbers
3
+ import random
4
+
5
+ import cv2
6
+ import numpy as np
7
+ from PIL import Image
8
+ from torchvision import transforms
9
+ from torchvision.transforms import Compose
10
+
11
+
12
+ def sample_asym(magnitude, size=None):
13
+ return np.random.beta(1, 4, size) * magnitude
14
+
15
+ def sample_sym(magnitude, size=None):
16
+ return (np.random.beta(4, 4, size=size) - 0.5) * 2 * magnitude
17
+
18
+ def sample_uniform(low, high, size=None):
19
+ return np.random.uniform(low, high, size=size)
20
+
21
+ def get_interpolation(type='random'):
22
+ if type == 'random':
23
+ choice = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA]
24
+ interpolation = choice[random.randint(0, len(choice)-1)]
25
+ elif type == 'nearest': interpolation = cv2.INTER_NEAREST
26
+ elif type == 'linear': interpolation = cv2.INTER_LINEAR
27
+ elif type == 'cubic': interpolation = cv2.INTER_CUBIC
28
+ elif type == 'area': interpolation = cv2.INTER_AREA
29
+ else: raise TypeError('Interpolation types only nearest, linear, cubic, area are supported!')
30
+ return interpolation
31
+
32
+ class CVRandomRotation(object):
33
+ def __init__(self, degrees=15):
34
+ assert isinstance(degrees, numbers.Number), "degree should be a single number."
35
+ assert degrees >= 0, "degree must be positive."
36
+ self.degrees = degrees
37
+
38
+ @staticmethod
39
+ def get_params(degrees):
40
+ return sample_sym(degrees)
41
+
42
+ def __call__(self, img):
43
+ angle = self.get_params(self.degrees)
44
+ src_h, src_w = img.shape[:2]
45
+ M = cv2.getRotationMatrix2D(center=(src_w/2, src_h/2), angle=angle, scale=1.0)
46
+ abs_cos, abs_sin = abs(M[0,0]), abs(M[0,1])
47
+ dst_w = int(src_h * abs_sin + src_w * abs_cos)
48
+ dst_h = int(src_h * abs_cos + src_w * abs_sin)
49
+ M[0, 2] += (dst_w - src_w)/2
50
+ M[1, 2] += (dst_h - src_h)/2
51
+
52
+ flags = get_interpolation()
53
+ return cv2.warpAffine(img, M, (dst_w, dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE)
54
+
55
+ class CVRandomAffine(object):
56
+ def __init__(self, degrees, translate=None, scale=None, shear=None):
57
+ assert isinstance(degrees, numbers.Number), "degree should be a single number."
58
+ assert degrees >= 0, "degree must be positive."
59
+ self.degrees = degrees
60
+
61
+ if translate is not None:
62
+ assert isinstance(translate, (tuple, list)) and len(translate) == 2, \
63
+ "translate should be a list or tuple and it must be of length 2."
64
+ for t in translate:
65
+ if not (0.0 <= t <= 1.0):
66
+ raise ValueError("translation values should be between 0 and 1")
67
+ self.translate = translate
68
+
69
+ if scale is not None:
70
+ assert isinstance(scale, (tuple, list)) and len(scale) == 2, \
71
+ "scale should be a list or tuple and it must be of length 2."
72
+ for s in scale:
73
+ if s <= 0:
74
+ raise ValueError("scale values should be positive")
75
+ self.scale = scale
76
+
77
+ if shear is not None:
78
+ if isinstance(shear, numbers.Number):
79
+ if shear < 0:
80
+ raise ValueError("If shear is a single number, it must be positive.")
81
+ self.shear = [shear]
82
+ else:
83
+ assert isinstance(shear, (tuple, list)) and (len(shear) == 2), \
84
+ "shear should be a list or tuple and it must be of length 2."
85
+ self.shear = shear
86
+ else:
87
+ self.shear = shear
88
+
89
+ def _get_inverse_affine_matrix(self, center, angle, translate, scale, shear):
90
+ # https://github.com/pytorch/vision/blob/v0.4.0/torchvision/transforms/functional.py#L717
91
+ from numpy import sin, cos, tan
92
+
93
+ if isinstance(shear, numbers.Number):
94
+ shear = [shear, 0]
95
+
96
+ if not isinstance(shear, (tuple, list)) and len(shear) == 2:
97
+ raise ValueError(
98
+ "Shear should be a single value or a tuple/list containing " +
99
+ "two values. Got {}".format(shear))
100
+
101
+ rot = math.radians(angle)
102
+ sx, sy = [math.radians(s) for s in shear]
103
+
104
+ cx, cy = center
105
+ tx, ty = translate
106
+
107
+ # RSS without scaling
108
+ a = cos(rot - sy) / cos(sy)
109
+ b = -cos(rot - sy) * tan(sx) / cos(sy) - sin(rot)
110
+ c = sin(rot - sy) / cos(sy)
111
+ d = -sin(rot - sy) * tan(sx) / cos(sy) + cos(rot)
112
+
113
+ # Inverted rotation matrix with scale and shear
114
+ # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
115
+ M = [d, -b, 0,
116
+ -c, a, 0]
117
+ M = [x / scale for x in M]
118
+
119
+ # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
120
+ M[2] += M[0] * (-cx - tx) + M[1] * (-cy - ty)
121
+ M[5] += M[3] * (-cx - tx) + M[4] * (-cy - ty)
122
+
123
+ # Apply center translation: C * RSS^-1 * C^-1 * T^-1
124
+ M[2] += cx
125
+ M[5] += cy
126
+ return M
127
+
128
+ @staticmethod
129
+ def get_params(degrees, translate, scale_ranges, shears, height):
130
+ angle = sample_sym(degrees)
131
+ if translate is not None:
132
+ max_dx = translate[0] * height
133
+ max_dy = translate[1] * height
134
+ translations = (np.round(sample_sym(max_dx)), np.round(sample_sym(max_dy)))
135
+ else:
136
+ translations = (0, 0)
137
+
138
+ if scale_ranges is not None:
139
+ scale = sample_uniform(scale_ranges[0], scale_ranges[1])
140
+ else:
141
+ scale = 1.0
142
+
143
+ if shears is not None:
144
+ if len(shears) == 1:
145
+ shear = [sample_sym(shears[0]), 0.]
146
+ elif len(shears) == 2:
147
+ shear = [sample_sym(shears[0]), sample_sym(shears[1])]
148
+ else:
149
+ shear = 0.0
150
+
151
+ return angle, translations, scale, shear
152
+
153
+
154
+ def __call__(self, img):
155
+ src_h, src_w = img.shape[:2]
156
+ angle, translate, scale, shear = self.get_params(
157
+ self.degrees, self.translate, self.scale, self.shear, src_h)
158
+
159
+ M = self._get_inverse_affine_matrix((src_w/2, src_h/2), angle, (0, 0), scale, shear)
160
+ M = np.array(M).reshape(2,3)
161
+
162
+ startpoints = [(0, 0), (src_w - 1, 0), (src_w - 1, src_h - 1), (0, src_h - 1)]
163
+ project = lambda x, y, a, b, c: int(a*x + b*y + c)
164
+ endpoints = [(project(x, y, *M[0]), project(x, y, *M[1])) for x, y in startpoints]
165
+
166
+ rect = cv2.minAreaRect(np.array(endpoints))
167
+ bbox = cv2.boxPoints(rect).astype(dtype=np.int)
168
+ max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
169
+ min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
170
+
171
+ dst_w = int(max_x - min_x)
172
+ dst_h = int(max_y - min_y)
173
+ M[0, 2] += (dst_w - src_w) / 2
174
+ M[1, 2] += (dst_h - src_h) / 2
175
+
176
+ # add translate
177
+ dst_w += int(abs(translate[0]))
178
+ dst_h += int(abs(translate[1]))
179
+ if translate[0] < 0: M[0, 2] += abs(translate[0])
180
+ if translate[1] < 0: M[1, 2] += abs(translate[1])
181
+
182
+ flags = get_interpolation()
183
+ return cv2.warpAffine(img, M, (dst_w , dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE)
184
+
185
+ class CVRandomPerspective(object):
186
+ def __init__(self, distortion=0.5):
187
+ self.distortion = distortion
188
+
189
+ def get_params(self, width, height, distortion):
190
+ offset_h = sample_asym(distortion * height / 2, size=4).astype(dtype=np.int)
191
+ offset_w = sample_asym(distortion * width / 2, size=4).astype(dtype=np.int)
192
+ topleft = ( offset_w[0], offset_h[0])
193
+ topright = (width - 1 - offset_w[1], offset_h[1])
194
+ botright = (width - 1 - offset_w[2], height - 1 - offset_h[2])
195
+ botleft = ( offset_w[3], height - 1 - offset_h[3])
196
+
197
+ startpoints = [(0, 0), (width - 1, 0), (width - 1, height - 1), (0, height - 1)]
198
+ endpoints = [topleft, topright, botright, botleft]
199
+ return np.array(startpoints, dtype=np.float32), np.array(endpoints, dtype=np.float32)
200
+
201
+ def __call__(self, img):
202
+ height, width = img.shape[:2]
203
+ startpoints, endpoints = self.get_params(width, height, self.distortion)
204
+ M = cv2.getPerspectiveTransform(startpoints, endpoints)
205
+
206
+ # TODO: more robust way to crop image
207
+ rect = cv2.minAreaRect(endpoints)
208
+ bbox = cv2.boxPoints(rect).astype(dtype=np.int)
209
+ max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
210
+ min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
211
+ min_x, min_y = max(min_x, 0), max(min_y, 0)
212
+
213
+ flags = get_interpolation()
214
+ img = cv2.warpPerspective(img, M, (max_x, max_y), flags=flags, borderMode=cv2.BORDER_REPLICATE)
215
+ img = img[min_y:, min_x:]
216
+ return img
217
+
218
+ class CVRescale(object):
219
+
220
+ def __init__(self, factor=4, base_size=(128, 512)):
221
+ """ Define image scales using gaussian pyramid and rescale image to target scale.
222
+
223
+ Args:
224
+ factor: the decayed factor from base size, factor=4 keeps target scale by default.
225
+ base_size: base size the build the bottom layer of pyramid
226
+ """
227
+ if isinstance(factor, numbers.Number):
228
+ self.factor = round(sample_uniform(0, factor))
229
+ elif isinstance(factor, (tuple, list)) and len(factor) == 2:
230
+ self.factor = round(sample_uniform(factor[0], factor[1]))
231
+ else:
232
+ raise Exception('factor must be number or list with length 2')
233
+ # assert factor is valid
234
+ self.base_h, self.base_w = base_size[:2]
235
+
236
+ def __call__(self, img):
237
+ if self.factor == 0: return img
238
+ src_h, src_w = img.shape[:2]
239
+ cur_w, cur_h = self.base_w, self.base_h
240
+ scale_img = cv2.resize(img, (cur_w, cur_h), interpolation=get_interpolation())
241
+ for _ in range(self.factor):
242
+ scale_img = cv2.pyrDown(scale_img)
243
+ scale_img = cv2.resize(scale_img, (src_w, src_h), interpolation=get_interpolation())
244
+ return scale_img
245
+
246
+ class CVGaussianNoise(object):
247
+ def __init__(self, mean=0, var=20):
248
+ self.mean = mean
249
+ if isinstance(var, numbers.Number):
250
+ self.var = max(int(sample_asym(var)), 1)
251
+ elif isinstance(var, (tuple, list)) and len(var) == 2:
252
+ self.var = int(sample_uniform(var[0], var[1]))
253
+ else:
254
+ raise Exception('degree must be number or list with length 2')
255
+
256
+ def __call__(self, img):
257
+ noise = np.random.normal(self.mean, self.var**0.5, img.shape)
258
+ img = np.clip(img + noise, 0, 255).astype(np.uint8)
259
+ return img
260
+
261
+ class CVMotionBlur(object):
262
+ def __init__(self, degrees=12, angle=90):
263
+ if isinstance(degrees, numbers.Number):
264
+ self.degree = max(int(sample_asym(degrees)), 1)
265
+ elif isinstance(degrees, (tuple, list)) and len(degrees) == 2:
266
+ self.degree = int(sample_uniform(degrees[0], degrees[1]))
267
+ else:
268
+ raise Exception('degree must be number or list with length 2')
269
+ self.angle = sample_uniform(-angle, angle)
270
+
271
+ def __call__(self, img):
272
+ M = cv2.getRotationMatrix2D((self.degree // 2, self.degree // 2), self.angle, 1)
273
+ motion_blur_kernel = np.zeros((self.degree, self.degree))
274
+ motion_blur_kernel[self.degree // 2, :] = 1
275
+ motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, (self.degree, self.degree))
276
+ motion_blur_kernel = motion_blur_kernel / self.degree
277
+ img = cv2.filter2D(img, -1, motion_blur_kernel)
278
+ img = np.clip(img, 0, 255).astype(np.uint8)
279
+ return img
280
+
281
+ class CVGeometry(object):
282
+ def __init__(self, degrees=15, translate=(0.3, 0.3), scale=(0.5, 2.),
283
+ shear=(45, 15), distortion=0.5, p=0.5):
284
+ self.p = p
285
+ type_p = random.random()
286
+ if type_p < 0.33:
287
+ self.transforms = CVRandomRotation(degrees=degrees)
288
+ elif type_p < 0.66:
289
+ self.transforms = CVRandomAffine(degrees=degrees, translate=translate, scale=scale, shear=shear)
290
+ else:
291
+ self.transforms = CVRandomPerspective(distortion=distortion)
292
+
293
+ def __call__(self, img):
294
+ if random.random() < self.p:
295
+ img = np.array(img)
296
+ return Image.fromarray(self.transforms(img))
297
+ else: return img
298
+
299
+ class CVDeterioration(object):
300
+ def __init__(self, var, degrees, factor, p=0.5):
301
+ self.p = p
302
+ transforms = []
303
+ if var is not None:
304
+ transforms.append(CVGaussianNoise(var=var))
305
+ if degrees is not None:
306
+ transforms.append(CVMotionBlur(degrees=degrees))
307
+ if factor is not None:
308
+ transforms.append(CVRescale(factor=factor))
309
+
310
+ random.shuffle(transforms)
311
+ transforms = Compose(transforms)
312
+ self.transforms = transforms
313
+
314
+ def __call__(self, img):
315
+ if random.random() < self.p:
316
+ img = np.array(img)
317
+ return Image.fromarray(self.transforms(img))
318
+ else: return img
319
+
320
+
321
+ class CVColorJitter(object):
322
+ def __init__(self, brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.5):
323
+ self.p = p
324
+ self.transforms = transforms.ColorJitter(brightness=brightness, contrast=contrast,
325
+ saturation=saturation, hue=hue)
326
+
327
+ def __call__(self, img):
328
+ if random.random() < self.p: return self.transforms(img)
329
+ else: return img
utils.py ADDED
@@ -0,0 +1,304 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ import os
3
+ import time
4
+
5
+ import cv2
6
+ import numpy as np
7
+ import torch
8
+ import yaml
9
+ from matplotlib import colors
10
+ from matplotlib import pyplot as plt
11
+ from torch import Tensor, nn
12
+ from torch.utils.data import ConcatDataset
13
+
14
+ class CharsetMapper(object):
15
+ """A simple class to map ids into strings.
16
+
17
+ It works only when the character set is 1:1 mapping between individual
18
+ characters and individual ids.
19
+ """
20
+
21
+ def __init__(self,
22
+ filename='',
23
+ max_length=30,
24
+ null_char=u'\u2591'):
25
+ """Creates a lookup table.
26
+
27
+ Args:
28
+ filename: Path to charset file which maps characters to ids.
29
+ max_sequence_length: The max length of ids and string.
30
+ null_char: A unicode character used to replace '<null>' character.
31
+ the default value is a light shade block '░'.
32
+ """
33
+ self.null_char = null_char
34
+ self.max_length = max_length
35
+
36
+ self.label_to_char = self._read_charset(filename)
37
+ self.char_to_label = dict(map(reversed, self.label_to_char.items()))
38
+ self.num_classes = len(self.label_to_char)
39
+
40
+ def _read_charset(self, filename):
41
+ """Reads a charset definition from a tab separated text file.
42
+
43
+ Args:
44
+ filename: a path to the charset file.
45
+
46
+ Returns:
47
+ a dictionary with keys equal to character codes and values - unicode
48
+ characters.
49
+ """
50
+ import re
51
+ pattern = re.compile(r'(\d+)\t(.+)')
52
+ charset = {}
53
+ self.null_label = 0
54
+ charset[self.null_label] = self.null_char
55
+ with open(filename, 'r') as f:
56
+ for i, line in enumerate(f):
57
+ m = pattern.match(line)
58
+ assert m, f'Incorrect charset file. line #{i}: {line}'
59
+ label = int(m.group(1)) + 1
60
+ char = m.group(2)
61
+ charset[label] = char
62
+ return charset
63
+
64
+ def trim(self, text):
65
+ assert isinstance(text, str)
66
+ return text.replace(self.null_char, '')
67
+
68
+ def get_text(self, labels, length=None, padding=True, trim=False):
69
+ """ Returns a string corresponding to a sequence of character ids.
70
+ """
71
+ length = length if length else self.max_length
72
+ labels = [l.item() if isinstance(l, Tensor) else int(l) for l in labels]
73
+ if padding:
74
+ labels = labels + [self.null_label] * (length-len(labels))
75
+ text = ''.join([self.label_to_char[label] for label in labels])
76
+ if trim: text = self.trim(text)
77
+ return text
78
+
79
+ def get_labels(self, text, length=None, padding=True, case_sensitive=False):
80
+ """ Returns the labels of the corresponding text.
81
+ """
82
+ length = length if length else self.max_length
83
+ if padding:
84
+ text = text + self.null_char * (length - len(text))
85
+ if not case_sensitive:
86
+ text = text.lower()
87
+ labels = [self.char_to_label[char] for char in text]
88
+ return labels
89
+
90
+ def pad_labels(self, labels, length=None):
91
+ length = length if length else self.max_length
92
+
93
+ return labels + [self.null_label] * (length - len(labels))
94
+
95
+ @property
96
+ def digits(self):
97
+ return '0123456789'
98
+
99
+ @property
100
+ def digit_labels(self):
101
+ return self.get_labels(self.digits, padding=False)
102
+
103
+ @property
104
+ def alphabets(self):
105
+ all_chars = list(self.char_to_label.keys())
106
+ valid_chars = []
107
+ for c in all_chars:
108
+ if c in 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ':
109
+ valid_chars.append(c)
110
+ return ''.join(valid_chars)
111
+
112
+ @property
113
+ def alphabet_labels(self):
114
+ return self.get_labels(self.alphabets, padding=False)
115
+
116
+
117
+ class Timer(object):
118
+ """A simple timer."""
119
+ def __init__(self):
120
+ self.data_time = 0.
121
+ self.data_diff = 0.
122
+ self.data_total_time = 0.
123
+ self.data_call = 0
124
+ self.running_time = 0.
125
+ self.running_diff = 0.
126
+ self.running_total_time = 0.
127
+ self.running_call = 0
128
+
129
+ def tic(self):
130
+ self.start_time = time.time()
131
+ self.running_time = self.start_time
132
+
133
+ def toc_data(self):
134
+ self.data_time = time.time()
135
+ self.data_diff = self.data_time - self.running_time
136
+ self.data_total_time += self.data_diff
137
+ self.data_call += 1
138
+
139
+ def toc_running(self):
140
+ self.running_time = time.time()
141
+ self.running_diff = self.running_time - self.data_time
142
+ self.running_total_time += self.running_diff
143
+ self.running_call += 1
144
+
145
+ def total_time(self):
146
+ return self.data_total_time + self.running_total_time
147
+
148
+ def average_time(self):
149
+ return self.average_data_time() + self.average_running_time()
150
+
151
+ def average_data_time(self):
152
+ return self.data_total_time / (self.data_call or 1)
153
+
154
+ def average_running_time(self):
155
+ return self.running_total_time / (self.running_call or 1)
156
+
157
+
158
+ class Logger(object):
159
+ _handle = None
160
+ _root = None
161
+
162
+ @staticmethod
163
+ def init(output_dir, name, phase):
164
+ format = '[%(asctime)s %(filename)s:%(lineno)d %(levelname)s {}] ' \
165
+ '%(message)s'.format(name)
166
+ logging.basicConfig(level=logging.INFO, format=format)
167
+
168
+ try: os.makedirs(output_dir)
169
+ except: pass
170
+ config_path = os.path.join(output_dir, f'{phase}.txt')
171
+ Logger._handle = logging.FileHandler(config_path)
172
+ Logger._root = logging.getLogger()
173
+
174
+ @staticmethod
175
+ def enable_file():
176
+ if Logger._handle is None or Logger._root is None:
177
+ raise Exception('Invoke Logger.init() first!')
178
+ Logger._root.addHandler(Logger._handle)
179
+
180
+ @staticmethod
181
+ def disable_file():
182
+ if Logger._handle is None or Logger._root is None:
183
+ raise Exception('Invoke Logger.init() first!')
184
+ Logger._root.removeHandler(Logger._handle)
185
+
186
+
187
+ class Config(object):
188
+
189
+ def __init__(self, config_path, host=True):
190
+ def __dict2attr(d, prefix=''):
191
+ for k, v in d.items():
192
+ if isinstance(v, dict):
193
+ __dict2attr(v, f'{prefix}{k}_')
194
+ else:
195
+ if k == 'phase':
196
+ assert v in ['train', 'test']
197
+ if k == 'stage':
198
+ assert v in ['pretrain-vision', 'pretrain-language',
199
+ 'train-semi-super', 'train-super']
200
+ self.__setattr__(f'{prefix}{k}', v)
201
+
202
+ assert os.path.exists(config_path), '%s does not exists!' % config_path
203
+ with open(config_path) as file:
204
+ config_dict = yaml.load(file, Loader=yaml.FullLoader)
205
+ with open('configs/template.yaml') as file:
206
+ default_config_dict = yaml.load(file, Loader=yaml.FullLoader)
207
+ __dict2attr(default_config_dict)
208
+ __dict2attr(config_dict)
209
+ self.global_workdir = os.path.join(self.global_workdir, self.global_name)
210
+
211
+ def __getattr__(self, item):
212
+ attr = self.__dict__.get(item)
213
+ if attr is None:
214
+ attr = dict()
215
+ prefix = f'{item}_'
216
+ for k, v in self.__dict__.items():
217
+ if k.startswith(prefix):
218
+ n = k.replace(prefix, '')
219
+ attr[n] = v
220
+ return attr if len(attr) > 0 else None
221
+ else:
222
+ return attr
223
+
224
+ def __repr__(self):
225
+ str = 'ModelConfig(\n'
226
+ for i, (k, v) in enumerate(sorted(vars(self).items())):
227
+ str += f'\t({i}): {k} = {v}\n'
228
+ str += ')'
229
+ return str
230
+
231
+ def blend_mask(image, mask, alpha=0.5, cmap='jet', color='b', color_alpha=1.0):
232
+ # normalize mask
233
+ mask = (mask-mask.min()) / (mask.max() - mask.min() + np.finfo(float).eps)
234
+ if mask.shape != image.shape:
235
+ mask = cv2.resize(mask,(image.shape[1], image.shape[0]))
236
+ # get color map
237
+ color_map = plt.get_cmap(cmap)
238
+ mask = color_map(mask)[:,:,:3]
239
+ # convert float to uint8
240
+ mask = (mask * 255).astype(dtype=np.uint8)
241
+
242
+ # set the basic color
243
+ basic_color = np.array(colors.to_rgb(color)) * 255
244
+ basic_color = np.tile(basic_color, [image.shape[0], image.shape[1], 1])
245
+ basic_color = basic_color.astype(dtype=np.uint8)
246
+ # blend with basic color
247
+ blended_img = cv2.addWeighted(image, color_alpha, basic_color, 1-color_alpha, 0)
248
+ # blend with mask
249
+ blended_img = cv2.addWeighted(blended_img, alpha, mask, 1-alpha, 0)
250
+
251
+ return blended_img
252
+
253
+ def onehot(label, depth, device=None):
254
+ """
255
+ Args:
256
+ label: shape (n1, n2, ..., )
257
+ depth: a scalar
258
+
259
+ Returns:
260
+ onehot: (n1, n2, ..., depth)
261
+ """
262
+ if not isinstance(label, torch.Tensor):
263
+ label = torch.tensor(label, device=device)
264
+ onehot = torch.zeros(label.size() + torch.Size([depth]), device=device)
265
+ onehot = onehot.scatter_(-1, label.unsqueeze(-1), 1)
266
+
267
+ return onehot
268
+
269
+ class MyDataParallel(nn.DataParallel):
270
+
271
+ def gather(self, outputs, target_device):
272
+ r"""
273
+ Gathers tensors from different GPUs on a specified device
274
+ (-1 means the CPU).
275
+ """
276
+ def gather_map(outputs):
277
+ out = outputs[0]
278
+ if isinstance(out, (str, int, float)):
279
+ return out
280
+ if isinstance(out, list) and isinstance(out[0], str):
281
+ return [o for out in outputs for o in out]
282
+ if isinstance(out, torch.Tensor):
283
+ return torch.nn.parallel._functions.Gather.apply(target_device, self.dim, *outputs)
284
+ if out is None:
285
+ return None
286
+ if isinstance(out, dict):
287
+ if not all((len(out) == len(d) for d in outputs)):
288
+ raise ValueError('All dicts must have the same number of keys')
289
+ return type(out)(((k, gather_map([d[k] for d in outputs]))
290
+ for k in out))
291
+ return type(out)(map(gather_map, zip(*outputs)))
292
+
293
+ # Recursive function calls like this create reference cycles.
294
+ # Setting the function to None clears the refcycle.
295
+ try:
296
+ res = gather_map(outputs)
297
+ finally:
298
+ gather_map = None
299
+ return res
300
+
301
+
302
+ class MyConcatDataset(ConcatDataset):
303
+ def __getattr__(self, k):
304
+ return getattr(self.datasets[0], k)