Spaces:
Running
on
A10G
Running
on
A10G
haoheliu
commited on
Commit
•
ea270ed
1
Parent(s):
9cfb0dd
first commit
Browse filesThis view is limited to 50 files because it contains too many changes.
See raw diff
- .gitignore +9 -0
- LICENSE +437 -0
- MANIFEST.in +2 -0
- README.md +16 -6
- app.py +361 -0
- audioldm2/__init__.py +2 -0
- audioldm2/audiomae_gen/__init__.py +1 -0
- audioldm2/audiomae_gen/sequence_input.py +429 -0
- audioldm2/audiomae_gen/utils.py +27 -0
- audioldm2/clap/__init__.py +0 -0
- audioldm2/clap/open_clip/__init__.py +25 -0
- audioldm2/clap/open_clip/bpe_simple_vocab_16e6.txt.gz +3 -0
- audioldm2/clap/open_clip/factory.py +276 -0
- audioldm2/clap/open_clip/feature_fusion.py +192 -0
- audioldm2/clap/open_clip/htsat.py +1304 -0
- audioldm2/clap/open_clip/loss.py +397 -0
- audioldm2/clap/open_clip/model.py +931 -0
- audioldm2/clap/open_clip/model_configs/HTSAT-base.json +23 -0
- audioldm2/clap/open_clip/model_configs/HTSAT-large.json +23 -0
- audioldm2/clap/open_clip/model_configs/HTSAT-tiny-win-1536.json +23 -0
- audioldm2/clap/open_clip/model_configs/HTSAT-tiny.json +23 -0
- audioldm2/clap/open_clip/model_configs/PANN-10.json +23 -0
- audioldm2/clap/open_clip/model_configs/PANN-14-fmax-18k.json +23 -0
- audioldm2/clap/open_clip/model_configs/PANN-14-fmax-8k-20s.json +23 -0
- audioldm2/clap/open_clip/model_configs/PANN-14-tiny-transformer.json +23 -0
- audioldm2/clap/open_clip/model_configs/PANN-14-win-1536.json +23 -0
- audioldm2/clap/open_clip/model_configs/PANN-14.json +23 -0
- audioldm2/clap/open_clip/model_configs/PANN-6.json +23 -0
- audioldm2/clap/open_clip/model_configs/RN101-quickgelu.json +22 -0
- audioldm2/clap/open_clip/model_configs/RN101.json +21 -0
- audioldm2/clap/open_clip/model_configs/RN50-quickgelu.json +22 -0
- audioldm2/clap/open_clip/model_configs/RN50.json +21 -0
- audioldm2/clap/open_clip/model_configs/RN50x16.json +21 -0
- audioldm2/clap/open_clip/model_configs/RN50x4.json +21 -0
- audioldm2/clap/open_clip/model_configs/ViT-B-16.json +16 -0
- audioldm2/clap/open_clip/model_configs/ViT-B-32-quickgelu.json +17 -0
- audioldm2/clap/open_clip/model_configs/ViT-B-32.json +16 -0
- audioldm2/clap/open_clip/model_configs/ViT-L-14.json +16 -0
- audioldm2/clap/open_clip/openai.py +156 -0
- audioldm2/clap/open_clip/pann_model.py +697 -0
- audioldm2/clap/open_clip/pretrained.py +167 -0
- audioldm2/clap/open_clip/timm_model.py +112 -0
- audioldm2/clap/open_clip/tokenizer.py +197 -0
- audioldm2/clap/open_clip/transform.py +45 -0
- audioldm2/clap/open_clip/utils.py +356 -0
- audioldm2/clap/training/__init__.py +0 -0
- audioldm2/clap/training/audioset_textmap.npy +3 -0
- audioldm2/clap/training/bpe_simple_vocab_16e6.txt.gz +3 -0
- audioldm2/clap/training/data.py +865 -0
- audioldm2/clap/training/params.py +563 -0
.gitignore
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
*.pyc
|
2 |
+
__pycache__
|
3 |
+
test.py
|
4 |
+
flagged
|
5 |
+
output
|
6 |
+
gradio_cached*
|
7 |
+
dist*
|
8 |
+
*egg-info
|
9 |
+
build*
|
LICENSE
ADDED
@@ -0,0 +1,437 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
Attribution-NonCommercial-ShareAlike 4.0 International
|
2 |
+
|
3 |
+
=======================================================================
|
4 |
+
|
5 |
+
Creative Commons Corporation ("Creative Commons") is not a law firm and
|
6 |
+
does not provide legal services or legal advice. Distribution of
|
7 |
+
Creative Commons public licenses does not create a lawyer-client or
|
8 |
+
other relationship. Creative Commons makes its licenses and related
|
9 |
+
information available on an "as-is" basis. Creative Commons gives no
|
10 |
+
warranties regarding its licenses, any material licensed under their
|
11 |
+
terms and conditions, or any related information. Creative Commons
|
12 |
+
disclaims all liability for damages resulting from their use to the
|
13 |
+
fullest extent possible.cp
|
14 |
+
|
15 |
+
Using Creative Commons Public Licenses
|
16 |
+
|
17 |
+
Creative Commons public licenses provide a standard set of terms and
|
18 |
+
conditions that creators and other rights holders may use to share
|
19 |
+
original works of authorship and other material subject to copyright
|
20 |
+
and certain other rights specified in the public license below. The
|
21 |
+
following considerations are for informational purposes only, are not
|
22 |
+
exhaustive, and do not form part of our licenses.
|
23 |
+
|
24 |
+
Considerations for licensors: Our public licenses are
|
25 |
+
intended for use by those authorized to give the public
|
26 |
+
permission to use material in ways otherwise restricted by
|
27 |
+
copyright and certain other rights. Our licenses are
|
28 |
+
irrevocable. Licensors should read and understand the terms
|
29 |
+
and conditions of the license they choose before applying it.
|
30 |
+
Licensors should also secure all rights necessary before
|
31 |
+
applying our licenses so that the public can reuse the
|
32 |
+
material as expected. Licensors should clearly mark any
|
33 |
+
material not subject to the license. This includes other CC-
|
34 |
+
licensed material, or material used under an exception or
|
35 |
+
limitation to copyright. More considerations for licensors:
|
36 |
+
wiki.creativecommons.org/Considerations_for_licensors
|
37 |
+
|
38 |
+
Considerations for the public: By using one of our public
|
39 |
+
licenses, a licensor grants the public permission to use the
|
40 |
+
licensed material under specified terms and conditions. If
|
41 |
+
the licensor's permission is not necessary for any reason--for
|
42 |
+
example, because of any applicable exception or limitation to
|
43 |
+
copyright--then that use is not regulated by the license. Our
|
44 |
+
licenses grant only permissions under copyright and certain
|
45 |
+
other rights that a licensor has authority to grant. Use of
|
46 |
+
the licensed material may still be restricted for other
|
47 |
+
reasons, including because others have copyright or other
|
48 |
+
rights in the material. A licensor may make special requests,
|
49 |
+
such as asking that all changes be marked or described.
|
50 |
+
Although not required by our licenses, you are encouraged to
|
51 |
+
respect those requests where reasonable. More_considerations
|
52 |
+
for the public:
|
53 |
+
wiki.creativecommons.org/Considerations_for_licensees
|
54 |
+
|
55 |
+
=======================================================================
|
56 |
+
|
57 |
+
Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
|
58 |
+
Public License
|
59 |
+
|
60 |
+
By exercising the Licensed Rights (defined below), You accept and agree
|
61 |
+
to be bound by the terms and conditions of this Creative Commons
|
62 |
+
Attribution-NonCommercial-ShareAlike 4.0 International Public License
|
63 |
+
("Public License"). To the extent this Public License may be
|
64 |
+
interpreted as a contract, You are granted the Licensed Rights in
|
65 |
+
consideration of Your acceptance of these terms and conditions, and the
|
66 |
+
Licensor grants You such rights in consideration of benefits the
|
67 |
+
Licensor receives from making the Licensed Material available under
|
68 |
+
these terms and conditions.
|
69 |
+
|
70 |
+
|
71 |
+
Section 1 -- Definitions.
|
72 |
+
|
73 |
+
a. Adapted Material means material subject to Copyright and Similar
|
74 |
+
Rights that is derived from or based upon the Licensed Material
|
75 |
+
and in which the Licensed Material is translated, altered,
|
76 |
+
arranged, transformed, or otherwise modified in a manner requiring
|
77 |
+
permission under the Copyright and Similar Rights held by the
|
78 |
+
Licensor. For purposes of this Public License, where the Licensed
|
79 |
+
Material is a musical work, performance, or sound recording,
|
80 |
+
Adapted Material is always produced where the Licensed Material is
|
81 |
+
synched in timed relation with a moving image.
|
82 |
+
|
83 |
+
b. Adapter's License means the license You apply to Your Copyright
|
84 |
+
and Similar Rights in Your contributions to Adapted Material in
|
85 |
+
accordance with the terms and conditions of this Public License.
|
86 |
+
|
87 |
+
c. BY-NC-SA Compatible License means a license listed at
|
88 |
+
creativecommons.org/compatiblelicenses, approved by Creative
|
89 |
+
Commons as essentially the equivalent of this Public License.
|
90 |
+
|
91 |
+
d. Copyright and Similar Rights means copyright and/or similar rights
|
92 |
+
closely related to copyright including, without limitation,
|
93 |
+
performance, broadcast, sound recording, and Sui Generis Database
|
94 |
+
Rights, without regard to how the rights are labeled or
|
95 |
+
categorized. For purposes of this Public License, the rights
|
96 |
+
specified in Section 2(b)(1)-(2) are not Copyright and Similar
|
97 |
+
Rights.
|
98 |
+
|
99 |
+
e. Effective Technological Measures means those measures that, in the
|
100 |
+
absence of proper authority, may not be circumvented under laws
|
101 |
+
fulfilling obligations under Article 11 of the WIPO Copyright
|
102 |
+
Treaty adopted on December 20, 1996, and/or similar international
|
103 |
+
agreements.
|
104 |
+
|
105 |
+
f. Exceptions and Limitations means fair use, fair dealing, and/or
|
106 |
+
any other exception or limitation to Copyright and Similar Rights
|
107 |
+
that applies to Your use of the Licensed Material.
|
108 |
+
|
109 |
+
g. License Elements means the license attributes listed in the name
|
110 |
+
of a Creative Commons Public License. The License Elements of this
|
111 |
+
Public License are Attribution, NonCommercial, and ShareAlike.
|
112 |
+
|
113 |
+
h. Licensed Material means the artistic or literary work, database,
|
114 |
+
or other material to which the Licensor applied this Public
|
115 |
+
License.
|
116 |
+
|
117 |
+
i. Licensed Rights means the rights granted to You subject to the
|
118 |
+
terms and conditions of this Public License, which are limited to
|
119 |
+
all Copyright and Similar Rights that apply to Your use of the
|
120 |
+
Licensed Material and that the Licensor has authority to license.
|
121 |
+
|
122 |
+
j. Licensor means the individual(s) or entity(ies) granting rights
|
123 |
+
under this Public License.
|
124 |
+
|
125 |
+
k. NonCommercial means not primarily intended for or directed towards
|
126 |
+
commercial advantage or monetary compensation. For purposes of
|
127 |
+
this Public License, the exchange of the Licensed Material for
|
128 |
+
other material subject to Copyright and Similar Rights by digital
|
129 |
+
file-sharing or similar means is NonCommercial provided there is
|
130 |
+
no payment of monetary compensation in connection with the
|
131 |
+
exchange.
|
132 |
+
|
133 |
+
l. Share means to provide material to the public by any means or
|
134 |
+
process that requires permission under the Licensed Rights, such
|
135 |
+
as reproduction, public display, public performance, distribution,
|
136 |
+
dissemination, communication, or importation, and to make material
|
137 |
+
available to the public including in ways that members of the
|
138 |
+
public may access the material from a place and at a time
|
139 |
+
individually chosen by them.
|
140 |
+
|
141 |
+
m. Sui Generis Database Rights means rights other than copyright
|
142 |
+
resulting from Directive 96/9/EC of the European Parliament and of
|
143 |
+
the Council of 11 March 1996 on the legal protection of databases,
|
144 |
+
as amended and/or succeeded, as well as other essentially
|
145 |
+
equivalent rights anywhere in the world.
|
146 |
+
|
147 |
+
n. You means the individual or entity exercising the Licensed Rights
|
148 |
+
under this Public License. Your has a corresponding meaning.
|
149 |
+
|
150 |
+
|
151 |
+
Section 2 -- Scope.
|
152 |
+
|
153 |
+
a. License grant.
|
154 |
+
|
155 |
+
1. Subject to the terms and conditions of this Public License,
|
156 |
+
the Licensor hereby grants You a worldwide, royalty-free,
|
157 |
+
non-sublicensable, non-exclusive, irrevocable license to
|
158 |
+
exercise the Licensed Rights in the Licensed Material to:
|
159 |
+
|
160 |
+
a. reproduce and Share the Licensed Material, in whole or
|
161 |
+
in part, for NonCommercial purposes only; and
|
162 |
+
|
163 |
+
b. produce, reproduce, and Share Adapted Material for
|
164 |
+
NonCommercial purposes only.
|
165 |
+
|
166 |
+
2. Exceptions and Limitations. For the avoidance of doubt, where
|
167 |
+
Exceptions and Limitations apply to Your use, this Public
|
168 |
+
License does not apply, and You do not need to comply with
|
169 |
+
its terms and conditions.
|
170 |
+
|
171 |
+
3. Term. The term of this Public License is specified in Section
|
172 |
+
6(a).
|
173 |
+
|
174 |
+
4. Media and formats; technical modifications allowed. The
|
175 |
+
Licensor authorizes You to exercise the Licensed Rights in
|
176 |
+
all media and formats whether now known or hereafter created,
|
177 |
+
and to make technical modifications necessary to do so. The
|
178 |
+
Licensor waives and/or agrees not to assert any right or
|
179 |
+
authority to forbid You from making technical modifications
|
180 |
+
necessary to exercise the Licensed Rights, including
|
181 |
+
technical modifications necessary to circumvent Effective
|
182 |
+
Technological Measures. For purposes of this Public License,
|
183 |
+
simply making modifications authorized by this Section 2(a)
|
184 |
+
(4) never produces Adapted Material.
|
185 |
+
|
186 |
+
5. Downstream recipients.
|
187 |
+
|
188 |
+
a. Offer from the Licensor -- Licensed Material. Every
|
189 |
+
recipient of the Licensed Material automatically
|
190 |
+
receives an offer from the Licensor to exercise the
|
191 |
+
Licensed Rights under the terms and conditions of this
|
192 |
+
Public License.
|
193 |
+
|
194 |
+
b. Additional offer from the Licensor -- Adapted Material.
|
195 |
+
Every recipient of Adapted Material from You
|
196 |
+
automatically receives an offer from the Licensor to
|
197 |
+
exercise the Licensed Rights in the Adapted Material
|
198 |
+
under the conditions of the Adapter's License You apply.
|
199 |
+
|
200 |
+
c. No downstream restrictions. You may not offer or impose
|
201 |
+
any additional or different terms or conditions on, or
|
202 |
+
apply any Effective Technological Measures to, the
|
203 |
+
Licensed Material if doing so restricts exercise of the
|
204 |
+
Licensed Rights by any recipient of the Licensed
|
205 |
+
Material.
|
206 |
+
|
207 |
+
6. No endorsement. Nothing in this Public License constitutes or
|
208 |
+
may be construed as permission to assert or imply that You
|
209 |
+
are, or that Your use of the Licensed Material is, connected
|
210 |
+
with, or sponsored, endorsed, or granted official status by,
|
211 |
+
the Licensor or others designated to receive attribution as
|
212 |
+
provided in Section 3(a)(1)(A)(i).
|
213 |
+
|
214 |
+
b. Other rights.
|
215 |
+
|
216 |
+
1. Moral rights, such as the right of integrity, are not
|
217 |
+
licensed under this Public License, nor are publicity,
|
218 |
+
privacy, and/or other similar personality rights; however, to
|
219 |
+
the extent possible, the Licensor waives and/or agrees not to
|
220 |
+
assert any such rights held by the Licensor to the limited
|
221 |
+
extent necessary to allow You to exercise the Licensed
|
222 |
+
Rights, but not otherwise.
|
223 |
+
|
224 |
+
2. Patent and trademark rights are not licensed under this
|
225 |
+
Public License.
|
226 |
+
|
227 |
+
3. To the extent possible, the Licensor waives any right to
|
228 |
+
collect royalties from You for the exercise of the Licensed
|
229 |
+
Rights, whether directly or through a collecting society
|
230 |
+
under any voluntary or waivable statutory or compulsory
|
231 |
+
licensing scheme. In all other cases the Licensor expressly
|
232 |
+
reserves any right to collect such royalties, including when
|
233 |
+
the Licensed Material is used other than for NonCommercial
|
234 |
+
purposes.
|
235 |
+
|
236 |
+
|
237 |
+
Section 3 -- License Conditions.
|
238 |
+
|
239 |
+
Your exercise of the Licensed Rights is expressly made subject to the
|
240 |
+
following conditions.
|
241 |
+
|
242 |
+
a. Attribution.
|
243 |
+
|
244 |
+
1. If You Share the Licensed Material (including in modified
|
245 |
+
form), You must:
|
246 |
+
|
247 |
+
a. retain the following if it is supplied by the Licensor
|
248 |
+
with the Licensed Material:
|
249 |
+
|
250 |
+
i. identification of the creator(s) of the Licensed
|
251 |
+
Material and any others designated to receive
|
252 |
+
attribution, in any reasonable manner requested by
|
253 |
+
the Licensor (including by pseudonym if
|
254 |
+
designated);
|
255 |
+
|
256 |
+
ii. a copyright notice;
|
257 |
+
|
258 |
+
iii. a notice that refers to this Public License;
|
259 |
+
|
260 |
+
iv. a notice that refers to the disclaimer of
|
261 |
+
warranties;
|
262 |
+
|
263 |
+
v. a URI or hyperlink to the Licensed Material to the
|
264 |
+
extent reasonably practicable;
|
265 |
+
|
266 |
+
b. indicate if You modified the Licensed Material and
|
267 |
+
retain an indication of any previous modifications; and
|
268 |
+
|
269 |
+
c. indicate the Licensed Material is licensed under this
|
270 |
+
Public License, and include the text of, or the URI or
|
271 |
+
hyperlink to, this Public License.
|
272 |
+
|
273 |
+
2. You may satisfy the conditions in Section 3(a)(1) in any
|
274 |
+
reasonable manner based on the medium, means, and context in
|
275 |
+
which You Share the Licensed Material. For example, it may be
|
276 |
+
reasonable to satisfy the conditions by providing a URI or
|
277 |
+
hyperlink to a resource that includes the required
|
278 |
+
information.
|
279 |
+
3. If requested by the Licensor, You must remove any of the
|
280 |
+
information required by Section 3(a)(1)(A) to the extent
|
281 |
+
reasonably practicable.
|
282 |
+
|
283 |
+
b. ShareAlike.
|
284 |
+
|
285 |
+
In addition to the conditions in Section 3(a), if You Share
|
286 |
+
Adapted Material You produce, the following conditions also apply.
|
287 |
+
|
288 |
+
1. The Adapter's License You apply must be a Creative Commons
|
289 |
+
license with the same License Elements, this version or
|
290 |
+
later, or a BY-NC-SA Compatible License.
|
291 |
+
|
292 |
+
2. You must include the text of, or the URI or hyperlink to, the
|
293 |
+
Adapter's License You apply. You may satisfy this condition
|
294 |
+
in any reasonable manner based on the medium, means, and
|
295 |
+
context in which You Share Adapted Material.
|
296 |
+
|
297 |
+
3. You may not offer or impose any additional or different terms
|
298 |
+
or conditions on, or apply any Effective Technological
|
299 |
+
Measures to, Adapted Material that restrict exercise of the
|
300 |
+
rights granted under the Adapter's License You apply.
|
301 |
+
|
302 |
+
|
303 |
+
Section 4 -- Sui Generis Database Rights.
|
304 |
+
|
305 |
+
Where the Licensed Rights include Sui Generis Database Rights that
|
306 |
+
apply to Your use of the Licensed Material:
|
307 |
+
|
308 |
+
a. for the avoidance of doubt, Section 2(a)(1) grants You the right
|
309 |
+
to extract, reuse, reproduce, and Share all or a substantial
|
310 |
+
portion of the contents of the database for NonCommercial purposes
|
311 |
+
only;
|
312 |
+
|
313 |
+
b. if You include all or a substantial portion of the database
|
314 |
+
contents in a database in which You have Sui Generis Database
|
315 |
+
Rights, then the database in which You have Sui Generis Database
|
316 |
+
Rights (but not its individual contents) is Adapted Material,
|
317 |
+
including for purposes of Section 3(b); and
|
318 |
+
|
319 |
+
c. You must comply with the conditions in Section 3(a) if You Share
|
320 |
+
all or a substantial portion of the contents of the database.
|
321 |
+
|
322 |
+
For the avoidance of doubt, this Section 4 supplements and does not
|
323 |
+
replace Your obligations under this Public License where the Licensed
|
324 |
+
Rights include other Copyright and Similar Rights.
|
325 |
+
|
326 |
+
|
327 |
+
Section 5 -- Disclaimer of Warranties and Limitation of Liability.
|
328 |
+
|
329 |
+
a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
|
330 |
+
EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
|
331 |
+
AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
|
332 |
+
ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
|
333 |
+
IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
|
334 |
+
WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
|
335 |
+
PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
|
336 |
+
ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
|
337 |
+
KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
|
338 |
+
ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
|
339 |
+
|
340 |
+
b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
|
341 |
+
TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
|
342 |
+
NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
|
343 |
+
INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
|
344 |
+
COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
|
345 |
+
USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
|
346 |
+
ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
|
347 |
+
DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
|
348 |
+
IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
|
349 |
+
|
350 |
+
c. The disclaimer of warranties and limitation of liability provided
|
351 |
+
above shall be interpreted in a manner that, to the extent
|
352 |
+
possible, most closely approximates an absolute disclaimer and
|
353 |
+
waiver of all liability.
|
354 |
+
|
355 |
+
|
356 |
+
Section 6 -- Term and Termination.
|
357 |
+
|
358 |
+
a. This Public License applies for the term of the Copyright and
|
359 |
+
Similar Rights licensed here. However, if You fail to comply with
|
360 |
+
this Public License, then Your rights under this Public License
|
361 |
+
terminate automatically.
|
362 |
+
|
363 |
+
b. Where Your right to use the Licensed Material has terminated under
|
364 |
+
Section 6(a), it reinstates:
|
365 |
+
|
366 |
+
1. automatically as of the date the violation is cured, provided
|
367 |
+
it is cured within 30 days of Your discovery of the
|
368 |
+
violation; or
|
369 |
+
|
370 |
+
2. upon express reinstatement by the Licensor.
|
371 |
+
|
372 |
+
For the avoidance of doubt, this Section 6(b) does not affect any
|
373 |
+
right the Licensor may have to seek remedies for Your violations
|
374 |
+
of this Public License.
|
375 |
+
|
376 |
+
c. For the avoidance of doubt, the Licensor may also offer the
|
377 |
+
Licensed Material under separate terms or conditions or stop
|
378 |
+
distributing the Licensed Material at any time; however, doing so
|
379 |
+
will not terminate this Public License.
|
380 |
+
|
381 |
+
d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
|
382 |
+
License.
|
383 |
+
|
384 |
+
|
385 |
+
Section 7 -- Other Terms and Conditions.
|
386 |
+
|
387 |
+
a. The Licensor shall not be bound by any additional or different
|
388 |
+
terms or conditions communicated by You unless expressly agreed.
|
389 |
+
|
390 |
+
b. Any arrangements, understandings, or agreements regarding the
|
391 |
+
Licensed Material not stated herein are separate from and
|
392 |
+
independent of the terms and conditions of this Public License.
|
393 |
+
|
394 |
+
|
395 |
+
Section 8 -- Interpretation.
|
396 |
+
|
397 |
+
a. For the avoidance of doubt, this Public License does not, and
|
398 |
+
shall not be interpreted to, reduce, limit, restrict, or impose
|
399 |
+
conditions on any use of the Licensed Material that could lawfully
|
400 |
+
be made without permission under this Public License.
|
401 |
+
|
402 |
+
b. To the extent possible, if any provision of this Public License is
|
403 |
+
deemed unenforceable, it shall be automatically reformed to the
|
404 |
+
minimum extent necessary to make it enforceable. If the provision
|
405 |
+
cannot be reformed, it shall be severed from this Public License
|
406 |
+
without affecting the enforceability of the remaining terms and
|
407 |
+
conditions.
|
408 |
+
|
409 |
+
c. No term or condition of this Public License will be waived and no
|
410 |
+
failure to comply consented to unless expressly agreed to by the
|
411 |
+
Licensor.
|
412 |
+
|
413 |
+
d. Nothing in this Public License constitutes or may be interpreted
|
414 |
+
as a limitation upon, or waiver of, any privileges and immunities
|
415 |
+
that apply to the Licensor or You, including from the legal
|
416 |
+
processes of any jurisdiction or authority.
|
417 |
+
|
418 |
+
=======================================================================
|
419 |
+
|
420 |
+
Creative Commons is not a party to its public
|
421 |
+
licenses. Notwithstanding, Creative Commons may elect to apply one of
|
422 |
+
its public licenses to material it publishes and in those instances
|
423 |
+
will be considered the “Licensor.” The text of the Creative Commons
|
424 |
+
public licenses is dedicated to the public domain under the CC0 Public
|
425 |
+
Domain Dedication. Except for the limited purpose of indicating that
|
426 |
+
material is shared under a Creative Commons public license or as
|
427 |
+
otherwise permitted by the Creative Commons policies published at
|
428 |
+
creativecommons.org/policies, Creative Commons does not authorize the
|
429 |
+
use of the trademark "Creative Commons" or any other trademark or logo
|
430 |
+
of Creative Commons without its prior written consent including,
|
431 |
+
without limitation, in connection with any unauthorized modifications
|
432 |
+
to any of its public licenses or any other arrangements,
|
433 |
+
understandings, or agreements concerning use of licensed material. For
|
434 |
+
the avoidance of doubt, this paragraph does not form part of the
|
435 |
+
public licenses.
|
436 |
+
|
437 |
+
Creative Commons may be contacted at creativecommons.org.
|
MANIFEST.in
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
include *.py LICENSE README.md
|
2 |
+
recursive-include audioldm2 *.txt *.py *.gz *.npy *.json
|
README.md
CHANGED
@@ -1,13 +1,23 @@
|
|
1 |
---
|
2 |
-
title:
|
3 |
-
emoji:
|
4 |
-
colorFrom:
|
5 |
-
colorTo:
|
6 |
sdk: gradio
|
7 |
-
sdk_version: 3.
|
8 |
app_file: app.py
|
9 |
pinned: false
|
10 |
-
license:
|
|
|
|
|
11 |
---
|
12 |
|
13 |
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
---
|
2 |
+
title: AudioLDM2 Text2Audio Text2Music Generation
|
3 |
+
emoji: 🔊
|
4 |
+
colorFrom: indigo
|
5 |
+
colorTo: red
|
6 |
sdk: gradio
|
7 |
+
sdk_version: 3.27.0
|
8 |
app_file: app.py
|
9 |
pinned: false
|
10 |
+
license: bigscience-openrail-m
|
11 |
+
duplicated_from: haoheliu/audioldm2-text2audio-text2music
|
12 |
+
|
13 |
---
|
14 |
|
15 |
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
16 |
+
|
17 |
+
## Reference
|
18 |
+
Part of the code from this repo is borrowed from the following repos. We would like to thank the authors of them for their contribution.
|
19 |
+
|
20 |
+
> https://github.com/LAION-AI/CLAP
|
21 |
+
> https://github.com/CompVis/stable-diffusion
|
22 |
+
> https://github.com/v-iashin/SpecVQGAN
|
23 |
+
> https://github.com/toshas/torch-fidelity
|
app.py
ADDED
@@ -0,0 +1,361 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from huggingface_hub import hf_hub_download
|
2 |
+
import torch
|
3 |
+
import os
|
4 |
+
|
5 |
+
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
6 |
+
|
7 |
+
import gradio as gr
|
8 |
+
from audioldm2 import text_to_audio, build_model
|
9 |
+
from share_btn import community_icon_html, loading_icon_html, share_js
|
10 |
+
|
11 |
+
model_id = "haoheliu/audioldm2-full"
|
12 |
+
hf_hub_download(repo_id="haoheliu/audioldm2-full", filename="audioldm2-full.pth")
|
13 |
+
|
14 |
+
audioldm = None
|
15 |
+
current_model_name = None
|
16 |
+
|
17 |
+
def text2audio(
|
18 |
+
text,
|
19 |
+
guidance_scale,
|
20 |
+
random_seed,
|
21 |
+
n_candidates,
|
22 |
+
model_name="audioldm2-full",
|
23 |
+
):
|
24 |
+
global audioldm, current_model_name
|
25 |
+
torch.set_float32_matmul_precision("high")
|
26 |
+
|
27 |
+
if audioldm is None or model_name != current_model_name:
|
28 |
+
audioldm = build_model(model_name=model_name)
|
29 |
+
current_model_name = model_name
|
30 |
+
audioldm = torch.compile(audioldm)
|
31 |
+
|
32 |
+
# print(text, length, guidance_scale)
|
33 |
+
waveform = text_to_audio(
|
34 |
+
latent_diffusion=audioldm,
|
35 |
+
text=text,
|
36 |
+
seed=random_seed,
|
37 |
+
duration=10,
|
38 |
+
guidance_scale=guidance_scale,
|
39 |
+
n_candidate_gen_per_text=int(n_candidates),
|
40 |
+
) # [bs, 1, samples]
|
41 |
+
waveform = [
|
42 |
+
gr.make_waveform((16000, wave[0]), bg_image="bg.png") for wave in waveform
|
43 |
+
]
|
44 |
+
# waveform = [(16000, np.random.randn(16000)), (16000, np.random.randn(16000))]
|
45 |
+
if len(waveform) == 1:
|
46 |
+
waveform = waveform[0]
|
47 |
+
return waveform
|
48 |
+
|
49 |
+
css = """
|
50 |
+
a {
|
51 |
+
color: inherit;
|
52 |
+
text-decoration: underline;
|
53 |
+
}
|
54 |
+
.gradio-container {
|
55 |
+
font-family: 'IBM Plex Sans', sans-serif;
|
56 |
+
}
|
57 |
+
.gr-button {
|
58 |
+
color: white;
|
59 |
+
border-color: #000000;
|
60 |
+
background: #000000;
|
61 |
+
}
|
62 |
+
input[type='range'] {
|
63 |
+
accent-color: #000000;
|
64 |
+
}
|
65 |
+
.dark input[type='range'] {
|
66 |
+
accent-color: #dfdfdf;
|
67 |
+
}
|
68 |
+
.container {
|
69 |
+
max-width: 730px;
|
70 |
+
margin: auto;
|
71 |
+
padding-top: 1.5rem;
|
72 |
+
}
|
73 |
+
#gallery {
|
74 |
+
min-height: 22rem;
|
75 |
+
margin-bottom: 15px;
|
76 |
+
margin-left: auto;
|
77 |
+
margin-right: auto;
|
78 |
+
border-bottom-right-radius: .5rem !important;
|
79 |
+
border-bottom-left-radius: .5rem !important;
|
80 |
+
}
|
81 |
+
#gallery>div>.h-full {
|
82 |
+
min-height: 20rem;
|
83 |
+
}
|
84 |
+
.details:hover {
|
85 |
+
text-decoration: underline;
|
86 |
+
}
|
87 |
+
.gr-button {
|
88 |
+
white-space: nowrap;
|
89 |
+
}
|
90 |
+
.gr-button:focus {
|
91 |
+
border-color: rgb(147 197 253 / var(--tw-border-opacity));
|
92 |
+
outline: none;
|
93 |
+
box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);
|
94 |
+
--tw-border-opacity: 1;
|
95 |
+
--tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);
|
96 |
+
--tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);
|
97 |
+
--tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));
|
98 |
+
--tw-ring-opacity: .5;
|
99 |
+
}
|
100 |
+
#advanced-btn {
|
101 |
+
font-size: .7rem !important;
|
102 |
+
line-height: 19px;
|
103 |
+
margin-top: 12px;
|
104 |
+
margin-bottom: 12px;
|
105 |
+
padding: 2px 8px;
|
106 |
+
border-radius: 14px !important;
|
107 |
+
}
|
108 |
+
#advanced-options {
|
109 |
+
margin-bottom: 20px;
|
110 |
+
}
|
111 |
+
.footer {
|
112 |
+
margin-bottom: 45px;
|
113 |
+
margin-top: 35px;
|
114 |
+
text-align: center;
|
115 |
+
border-bottom: 1px solid #e5e5e5;
|
116 |
+
}
|
117 |
+
.footer>p {
|
118 |
+
font-size: .8rem;
|
119 |
+
display: inline-block;
|
120 |
+
padding: 0 10px;
|
121 |
+
transform: translateY(10px);
|
122 |
+
background: white;
|
123 |
+
}
|
124 |
+
.dark .footer {
|
125 |
+
border-color: #303030;
|
126 |
+
}
|
127 |
+
.dark .footer>p {
|
128 |
+
background: #0b0f19;
|
129 |
+
}
|
130 |
+
.acknowledgments h4{
|
131 |
+
margin: 1.25em 0 .25em 0;
|
132 |
+
font-weight: bold;
|
133 |
+
font-size: 115%;
|
134 |
+
}
|
135 |
+
#container-advanced-btns{
|
136 |
+
display: flex;
|
137 |
+
flex-wrap: wrap;
|
138 |
+
justify-content: space-between;
|
139 |
+
align-items: center;
|
140 |
+
}
|
141 |
+
.animate-spin {
|
142 |
+
animation: spin 1s linear infinite;
|
143 |
+
}
|
144 |
+
@keyframes spin {
|
145 |
+
from {
|
146 |
+
transform: rotate(0deg);
|
147 |
+
}
|
148 |
+
to {
|
149 |
+
transform: rotate(360deg);
|
150 |
+
}
|
151 |
+
}
|
152 |
+
#share-btn-container {
|
153 |
+
display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
|
154 |
+
margin-top: 10px;
|
155 |
+
margin-left: auto;
|
156 |
+
}
|
157 |
+
#share-btn {
|
158 |
+
all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;right:0;
|
159 |
+
}
|
160 |
+
#share-btn * {
|
161 |
+
all: unset;
|
162 |
+
}
|
163 |
+
#share-btn-container div:nth-child(-n+2){
|
164 |
+
width: auto !important;
|
165 |
+
min-height: 0px !important;
|
166 |
+
}
|
167 |
+
#share-btn-container .wrap {
|
168 |
+
display: none !important;
|
169 |
+
}
|
170 |
+
.gr-form{
|
171 |
+
flex: 1 1 50%; border-top-right-radius: 0; border-bottom-right-radius: 0;
|
172 |
+
}
|
173 |
+
#prompt-container{
|
174 |
+
gap: 0;
|
175 |
+
}
|
176 |
+
#generated_id{
|
177 |
+
min-height: 700px
|
178 |
+
}
|
179 |
+
#setting_id{
|
180 |
+
margin-bottom: 12px;
|
181 |
+
text-align: center;
|
182 |
+
font-weight: 900;
|
183 |
+
}
|
184 |
+
"""
|
185 |
+
iface = gr.Blocks(css=css)
|
186 |
+
|
187 |
+
with iface:
|
188 |
+
gr.HTML(
|
189 |
+
"""
|
190 |
+
<div style="text-align: center; max-width: 700px; margin: 0 auto;">
|
191 |
+
<div
|
192 |
+
style="
|
193 |
+
display: inline-flex;
|
194 |
+
align-items: center;
|
195 |
+
gap: 0.8rem;
|
196 |
+
font-size: 1.75rem;
|
197 |
+
"
|
198 |
+
>
|
199 |
+
<h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;">
|
200 |
+
AudioLDM 2: A General Framework for Audio, Music, and Speech Generation
|
201 |
+
</h1>
|
202 |
+
</div>
|
203 |
+
<p style="margin-bottom: 10px; font-size: 94%">
|
204 |
+
<a href="https://arxiv.org/abs/2301.12503">[Paper]</a> <a href="https://audioldm.github.io/">[Project page]</a>
|
205 |
+
</p>
|
206 |
+
</div>
|
207 |
+
"""
|
208 |
+
)
|
209 |
+
gr.HTML(
|
210 |
+
"""
|
211 |
+
<h1 style="font-weight: 900; margin-bottom: 7px;">
|
212 |
+
AudioLDM 2: A General Framework for Audio, Music, and Speech Generation
|
213 |
+
</h1>
|
214 |
+
<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
|
215 |
+
<br/>
|
216 |
+
<a href="https://huggingface.co/spaces/haoheliu/audioldm2-text2audio-text2music?duplicate=true">
|
217 |
+
<img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
|
218 |
+
<p/>
|
219 |
+
"""
|
220 |
+
)
|
221 |
+
with gr.Group():
|
222 |
+
with gr.Box():
|
223 |
+
############# Input
|
224 |
+
textbox = gr.Textbox(
|
225 |
+
value="A forest of wind chimes singing a soothing melody in the breeze.",
|
226 |
+
max_lines=1,
|
227 |
+
label="Input your text here. Your text is important for the audio quality. Please ensure it is descriptive by using more adjectives.",
|
228 |
+
elem_id="prompt-in",
|
229 |
+
)
|
230 |
+
|
231 |
+
with gr.Accordion("Click to modify detailed configurations", open=False):
|
232 |
+
seed = gr.Number(
|
233 |
+
value=45,
|
234 |
+
label="Change this value (any integer number) will lead to a different generation result.",
|
235 |
+
)
|
236 |
+
# duration = gr.Slider(
|
237 |
+
# 10, 10, value=10, step=2.5, label="Duration (seconds)"
|
238 |
+
# )
|
239 |
+
guidance_scale = gr.Slider(
|
240 |
+
0,
|
241 |
+
6,
|
242 |
+
value=3.5,
|
243 |
+
step=0.5,
|
244 |
+
label="Guidance scale (Large => better quality and relavancy to text; Small => better diversity)",
|
245 |
+
)
|
246 |
+
n_candidates = gr.Slider(
|
247 |
+
1,
|
248 |
+
3,
|
249 |
+
value=3,
|
250 |
+
step=1,
|
251 |
+
label="Automatic quality control. This number control the number of candidates (e.g., generate three audios and choose the best to show you). A Larger value usually lead to better quality with heavier computation",
|
252 |
+
)
|
253 |
+
# model_name = gr.Dropdown(
|
254 |
+
# ["audioldm-m-text-ft", "audioldm-s-text-ft", "audioldm-m-full","audioldm-s-full-v2", "audioldm-s-full", "audioldm-l-full"], value="audioldm-m-full", label="Choose the model to use. audioldm-m-text-ft and audioldm-s-text-ft are recommanded. -s- means small, -m- means medium and -l- means large",
|
255 |
+
# )
|
256 |
+
############# Output
|
257 |
+
# outputs=gr.Audio(label="Output", type="numpy")
|
258 |
+
outputs = gr.Video(label="Output", elem_id="output-video")
|
259 |
+
|
260 |
+
# with gr.Group(elem_id="container-advanced-btns"):
|
261 |
+
# # advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
|
262 |
+
# with gr.Group(elem_id="share-btn-container"):
|
263 |
+
# community_icon = gr.HTML(community_icon_html, visible=False)
|
264 |
+
# loading_icon = gr.HTML(loading_icon_html, visible=False)
|
265 |
+
# share_button = gr.Button("Share to community", elem_id="share-btn", visible=False)
|
266 |
+
# outputs=[gr.Audio(label="Output", type="numpy"), gr.Audio(label="Output", type="numpy")]
|
267 |
+
btn = gr.Button("Submit").style(full_width=True)
|
268 |
+
|
269 |
+
with gr.Group(elem_id="share-btn-container", visible=False):
|
270 |
+
community_icon = gr.HTML(community_icon_html)
|
271 |
+
loading_icon = gr.HTML(loading_icon_html)
|
272 |
+
share_button = gr.Button("Share to community", elem_id="share-btn")
|
273 |
+
|
274 |
+
# btn.click(text2audio, inputs=[
|
275 |
+
# textbox, duration, guidance_scale, seed, n_candidates, model_name], outputs=[outputs])
|
276 |
+
btn.click(
|
277 |
+
text2audio,
|
278 |
+
inputs=[textbox, guidance_scale, seed, n_candidates],
|
279 |
+
outputs=[outputs],
|
280 |
+
)
|
281 |
+
|
282 |
+
share_button.click(None, [], [], _js=share_js)
|
283 |
+
gr.HTML(
|
284 |
+
"""
|
285 |
+
<div class="footer" style="text-align: center; max-width: 700px; margin: 0 auto;">
|
286 |
+
<p>Follow the latest update of AudioLDM on our<a href="https://github.com/haoheliu/AudioLDM" style="text-decoration: underline;" target="_blank"> Github repo</a>
|
287 |
+
</p>
|
288 |
+
<br>
|
289 |
+
<p>Model by <a href="https://twitter.com/LiuHaohe" style="text-decoration: underline;" target="_blank">Haohe Liu</a></p>
|
290 |
+
<br>
|
291 |
+
</div>
|
292 |
+
"""
|
293 |
+
)
|
294 |
+
gr.Examples(
|
295 |
+
[
|
296 |
+
[
|
297 |
+
"An excited crowd cheering at a sports game.",
|
298 |
+
3.5,
|
299 |
+
45,
|
300 |
+
3,
|
301 |
+
"audioldm2-full",
|
302 |
+
],
|
303 |
+
[
|
304 |
+
"A cat is meowing for attention.",
|
305 |
+
3.5,
|
306 |
+
45,
|
307 |
+
3,
|
308 |
+
"audioldm2-full",
|
309 |
+
],
|
310 |
+
[
|
311 |
+
"Birds singing sweetly in a blooming garden.",
|
312 |
+
3.5,
|
313 |
+
45,
|
314 |
+
3,
|
315 |
+
"audioldm2-full",
|
316 |
+
],
|
317 |
+
[
|
318 |
+
"A modern synthesizer creating futuristic soundscapes.",
|
319 |
+
3.5,
|
320 |
+
45,
|
321 |
+
3,
|
322 |
+
"audioldm2-full",
|
323 |
+
],
|
324 |
+
[
|
325 |
+
"The vibrant beat of Brazilian samba drums.",
|
326 |
+
3.5,
|
327 |
+
45,
|
328 |
+
3,
|
329 |
+
"audioldm2-full",
|
330 |
+
],
|
331 |
+
],
|
332 |
+
fn=text2audio,
|
333 |
+
# inputs=[textbox, duration, guidance_scale, seed, n_candidates, model_name],
|
334 |
+
inputs=[textbox, guidance_scale, seed, n_candidates],
|
335 |
+
outputs=[outputs],
|
336 |
+
cache_examples=True,
|
337 |
+
)
|
338 |
+
gr.HTML(
|
339 |
+
"""
|
340 |
+
<div class="acknowledgements">
|
341 |
+
<p>Essential Tricks for Enhancing the Quality of Your Generated Audio</p>
|
342 |
+
<p>1. Try to use more adjectives to describe your sound. For example: "A man is speaking clearly and slowly in a large room" is better than "A man is speaking". This can make sure AudioLDM understands what you want.</p>
|
343 |
+
<p>2. Try to use different random seeds, which can affect the generation quality significantly sometimes.</p>
|
344 |
+
<p>3. It's better to use general terms like 'man' or 'woman' instead of specific names for individuals or abstract objects that humans may not be familiar with, such as 'mummy'.</p>
|
345 |
+
</div>
|
346 |
+
"""
|
347 |
+
)
|
348 |
+
|
349 |
+
with gr.Accordion("Additional information", open=False):
|
350 |
+
gr.HTML(
|
351 |
+
"""
|
352 |
+
<div class="acknowledgments">
|
353 |
+
<p> We build the model with data from <a href="http://research.google.com/audioset/">AudioSet</a>, <a href="https://freesound.org/">Freesound</a> and <a href="https://sound-effects.bbcrewind.co.uk/">BBC Sound Effect library</a>. We share this demo based on the <a href="https://assets.publishing.service.gov.uk/government/uploads/system/uploads/attachment_data/file/375954/Research.pdf">UK copyright exception</a> of data for academic research. </p>
|
354 |
+
</div>
|
355 |
+
"""
|
356 |
+
)
|
357 |
+
# <p>This demo is strictly for research demo purpose only. For commercial use please <a href="[email protected]">contact us</a>.</p>
|
358 |
+
|
359 |
+
iface.queue(concurrency_count=3)
|
360 |
+
# iface.launch(debug=True)
|
361 |
+
iface.launch(debug=True, share=True)
|
audioldm2/__init__.py
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
from .utils import seed_everything, save_wave, get_time, get_duration, read_list
|
2 |
+
from .pipeline import *
|
audioldm2/audiomae_gen/__init__.py
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
from .sequence_input import Sequence2AudioMAE
|
audioldm2/audiomae_gen/sequence_input.py
ADDED
@@ -0,0 +1,429 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
from audioldm2.latent_diffusion.util import (
|
4 |
+
instantiate_from_config,
|
5 |
+
)
|
6 |
+
|
7 |
+
# from latent_diffusion.modules.encoders.modules import CLAPAudioEmbeddingClassifierFreev2
|
8 |
+
from transformers import GPT2Config, GPT2Model
|
9 |
+
import torch.optim.lr_scheduler as lr_scheduler
|
10 |
+
|
11 |
+
class Sequence2AudioMAE(nn.Module):
|
12 |
+
def __init__(
|
13 |
+
self,
|
14 |
+
base_learning_rate,
|
15 |
+
sequence_gen_length,
|
16 |
+
sequence_input_key,
|
17 |
+
sequence_input_embed_dim,
|
18 |
+
cond_stage_config,
|
19 |
+
optimizer_type="AdamW",
|
20 |
+
use_warmup=True,
|
21 |
+
use_ar_gen_loss=False,
|
22 |
+
use_audiomae_linear=False,
|
23 |
+
target_tokens_mask_ratio=0.0,
|
24 |
+
random_mask_ratio=False,
|
25 |
+
**kwargs
|
26 |
+
):
|
27 |
+
super().__init__()
|
28 |
+
assert use_audiomae_linear == False
|
29 |
+
self.random_mask_ratio = random_mask_ratio
|
30 |
+
self.learning_rate = base_learning_rate
|
31 |
+
self.cond_stage_config = cond_stage_config
|
32 |
+
self.use_audiomae_linear = use_audiomae_linear
|
33 |
+
self.optimizer_type = optimizer_type
|
34 |
+
self.use_warmup = use_warmup
|
35 |
+
self.use_ar_gen_loss = use_ar_gen_loss
|
36 |
+
# Even though the LDM can be conditioned on mutliple pooling rate
|
37 |
+
# Our model always predict the higest pooling rate
|
38 |
+
|
39 |
+
# self.time_pool = max(self.cond_stage_config["crossattn_audiomae_pooled"]["params"]["time_pooling_factors"])
|
40 |
+
# self.freq_pool = max(self.cond_stage_config["crossattn_audiomae_pooled"]["params"]["freq_pooling_factors"])
|
41 |
+
# self.mae_token_num = int(512/(self.time_pool*self.freq_pool))
|
42 |
+
|
43 |
+
self.mae_token_num = sequence_gen_length
|
44 |
+
self.sequence_input_key = sequence_input_key
|
45 |
+
self.sequence_input_embed_dim = sequence_input_embed_dim
|
46 |
+
self.target_tokens_mask_ratio = target_tokens_mask_ratio
|
47 |
+
|
48 |
+
self.start_of_sequence_tokens = nn.Embedding(32, 768)
|
49 |
+
self.end_of_sequence_tokens = nn.Embedding(32, 768)
|
50 |
+
|
51 |
+
self.input_sequence_embed_linear = nn.ModuleList([])
|
52 |
+
self.initial_learning_rate = None
|
53 |
+
|
54 |
+
for dim in self.sequence_input_embed_dim:
|
55 |
+
self.input_sequence_embed_linear.append(nn.Linear(dim, 768))
|
56 |
+
|
57 |
+
self.cond_stage_models = nn.ModuleList([])
|
58 |
+
self.instantiate_cond_stage(cond_stage_config)
|
59 |
+
self.initialize_param_check_toolkit()
|
60 |
+
|
61 |
+
# configuration = GPT2Config(n_layer=1) # TODO
|
62 |
+
# self.model=GPT2Model(configuration)
|
63 |
+
###################
|
64 |
+
# self.model=nn.Linear(768,768, bias=False) # TODO change the model
|
65 |
+
# with torch.no_grad():
|
66 |
+
# self.model.weight.copy_(torch.eye(768))
|
67 |
+
###################
|
68 |
+
self.model = GPT2Model(GPT2Config.from_pretrained("gpt2"))
|
69 |
+
###################
|
70 |
+
# self.model = nn.LSTM(input_size=768, hidden_size=768, num_layers=1,bias=False) # TODO
|
71 |
+
|
72 |
+
# self.loss_fn = nn.MSELoss()
|
73 |
+
self.loss_fn = nn.L1Loss()
|
74 |
+
|
75 |
+
self.logger_save_dir = None
|
76 |
+
self.logger_exp_name = None
|
77 |
+
self.logger_exp_group_name = None
|
78 |
+
self.logger_version = None
|
79 |
+
|
80 |
+
def set_log_dir(self, save_dir, exp_group_name, exp_name):
|
81 |
+
self.logger_save_dir = save_dir
|
82 |
+
self.logger_exp_group_name = exp_group_name
|
83 |
+
self.logger_exp_name = exp_name
|
84 |
+
|
85 |
+
def cfg_uncond(self, batch_size):
|
86 |
+
unconditional_conditioning = {}
|
87 |
+
for key in self.cond_stage_model_metadata:
|
88 |
+
model_idx = self.cond_stage_model_metadata[key]["model_idx"]
|
89 |
+
unconditional_conditioning[key] = self.cond_stage_models[
|
90 |
+
model_idx
|
91 |
+
].get_unconditional_condition(batch_size)
|
92 |
+
assert (
|
93 |
+
"crossattn_audiomae_pooled" in unconditional_conditioning.keys()
|
94 |
+
), "The module is not initialized with AudioMAE"
|
95 |
+
unconditional_conditioning[
|
96 |
+
"crossattn_clap_to_audiomae_feature"
|
97 |
+
] = unconditional_conditioning["crossattn_audiomae_pooled"]
|
98 |
+
return unconditional_conditioning
|
99 |
+
|
100 |
+
def configure_optimizers(self):
|
101 |
+
lr = float(self.learning_rate)
|
102 |
+
# params = list(self.model.parameters()) + list(self.input_sequence_embed_linear.parameters())
|
103 |
+
params = list(self.parameters())
|
104 |
+
|
105 |
+
# opt = torch.optim.Adam(params, lr=lr, betas=(0.9, 0.98), eps=1e-9)
|
106 |
+
opt = eval(self.optimizer_type)(params, lr=lr)
|
107 |
+
scheduler = lr_scheduler.StepLR(opt, step_size=10, gamma=0.8)
|
108 |
+
return [opt], [scheduler]
|
109 |
+
|
110 |
+
def add_sos_eos_tokens(self, _id, sequence, attn_mask):
|
111 |
+
batchsize = sequence.size(0)
|
112 |
+
|
113 |
+
new_attn_mask_step = torch.ones((batchsize, 1)).to(sequence.device)
|
114 |
+
key_id = torch.tensor([_id]).to(sequence.device)
|
115 |
+
|
116 |
+
# Add two more steps to attn mask
|
117 |
+
new_attn_mask = torch.cat(
|
118 |
+
[new_attn_mask_step, attn_mask, new_attn_mask_step], dim=1
|
119 |
+
)
|
120 |
+
|
121 |
+
# Add two more tokens in the sequence
|
122 |
+
sos_token = self.start_of_sequence_tokens(key_id).expand(batchsize, 1, -1)
|
123 |
+
eos_token = self.end_of_sequence_tokens(key_id).expand(batchsize, 1, -1)
|
124 |
+
new_sequence = torch.cat([sos_token, sequence, eos_token], dim=1)
|
125 |
+
return new_sequence, new_attn_mask
|
126 |
+
|
127 |
+
def truncate_sequence_and_mask(self, sequence, mask, max_len=512):
|
128 |
+
if sequence.size(1) > max_len:
|
129 |
+
print(
|
130 |
+
"The input sequence length to GPT-2 model is too long:",
|
131 |
+
sequence.size(1),
|
132 |
+
)
|
133 |
+
return sequence[:, :max_len], mask[:, :max_len]
|
134 |
+
else:
|
135 |
+
return sequence, mask
|
136 |
+
|
137 |
+
def get_input_sequence_and_mask(self, cond_dict):
|
138 |
+
input_embeds = None
|
139 |
+
input_embeds_attn_mask = None
|
140 |
+
for _id, sequence_key in enumerate(self.sequence_input_key):
|
141 |
+
assert sequence_key in cond_dict.keys(), (
|
142 |
+
"Invalid sequence key %s" % sequence_key
|
143 |
+
)
|
144 |
+
cond_embed = cond_dict[sequence_key]
|
145 |
+
if isinstance(cond_embed, list):
|
146 |
+
assert (
|
147 |
+
len(cond_embed) == 2
|
148 |
+
), "The crossattn returned list should have length 2, including embed and attn_mask"
|
149 |
+
item_input_embeds, item_attn_mask = cond_embed
|
150 |
+
|
151 |
+
item_input_embeds = self.input_sequence_embed_linear[_id](
|
152 |
+
item_input_embeds
|
153 |
+
)
|
154 |
+
|
155 |
+
item_input_embeds, item_attn_mask = self.add_sos_eos_tokens(
|
156 |
+
_id, item_input_embeds, item_attn_mask
|
157 |
+
)
|
158 |
+
|
159 |
+
if input_embeds is None and input_embeds_attn_mask is None:
|
160 |
+
input_embeds, input_embeds_attn_mask = (
|
161 |
+
item_input_embeds,
|
162 |
+
item_attn_mask,
|
163 |
+
)
|
164 |
+
else:
|
165 |
+
input_embeds = torch.cat(
|
166 |
+
[input_embeds, item_input_embeds], dim=1
|
167 |
+
) # The 1-st dimension is time steps
|
168 |
+
input_embeds_attn_mask = torch.cat(
|
169 |
+
[input_embeds_attn_mask, item_attn_mask], dim=1
|
170 |
+
) # The 1-st dimension is time steps
|
171 |
+
else:
|
172 |
+
assert isinstance(cond_embed, torch.Tensor)
|
173 |
+
cond_embed = self.input_sequence_embed_linear[_id](cond_embed)
|
174 |
+
attn_mask = torch.ones((cond_embed.size(0), cond_embed.size(1))).to(
|
175 |
+
cond_embed.device
|
176 |
+
)
|
177 |
+
|
178 |
+
item_input_embeds, item_attn_mask = self.add_sos_eos_tokens(
|
179 |
+
_id, cond_embed, attn_mask
|
180 |
+
)
|
181 |
+
|
182 |
+
if input_embeds is None and input_embeds_attn_mask is None:
|
183 |
+
input_embeds, input_embeds_attn_mask = (
|
184 |
+
item_input_embeds,
|
185 |
+
item_attn_mask,
|
186 |
+
)
|
187 |
+
else:
|
188 |
+
input_embeds, input_embeds_attn_mask = torch.cat(
|
189 |
+
[input_embeds, item_input_embeds], dim=1
|
190 |
+
), torch.cat([input_embeds_attn_mask, item_attn_mask], dim=1)
|
191 |
+
|
192 |
+
assert input_embeds is not None and input_embeds_attn_mask is not None
|
193 |
+
|
194 |
+
input_embeds, input_embeds_attn_mask = self.truncate_sequence_and_mask(
|
195 |
+
input_embeds, input_embeds_attn_mask, int(1024 - self.mae_token_num)
|
196 |
+
)
|
197 |
+
cond_sequence_end_time_idx = input_embeds.size(
|
198 |
+
1
|
199 |
+
) # The index that we start to collect the output embeds
|
200 |
+
|
201 |
+
return input_embeds, input_embeds_attn_mask, cond_sequence_end_time_idx
|
202 |
+
|
203 |
+
def warmup_step(self):
|
204 |
+
if self.initial_learning_rate is None:
|
205 |
+
self.initial_learning_rate = float(self.learning_rate)
|
206 |
+
|
207 |
+
# Only the first parameter group
|
208 |
+
if self.global_step <= 1000:
|
209 |
+
if self.global_step == 0:
|
210 |
+
print(
|
211 |
+
"Warming up learning rate start with %s"
|
212 |
+
% self.initial_learning_rate
|
213 |
+
)
|
214 |
+
self.trainer.optimizers[0].param_groups[0]["lr"] = (
|
215 |
+
self.global_step / 1000
|
216 |
+
) * self.initial_learning_rate
|
217 |
+
else:
|
218 |
+
# TODO set learning rate here
|
219 |
+
self.trainer.optimizers[0].param_groups[0][
|
220 |
+
"lr"
|
221 |
+
] = self.initial_learning_rate
|
222 |
+
|
223 |
+
def mask_target_sequence(self, target_embeds, target_embeds_attn_mask):
|
224 |
+
time_seq_mask = None
|
225 |
+
if self.target_tokens_mask_ratio > 1e-4:
|
226 |
+
batchsize, time_seq_len, embed_dim = target_embeds.size()
|
227 |
+
_, time_seq_len = target_embeds_attn_mask.size()
|
228 |
+
# Generate random mask
|
229 |
+
if self.random_mask_ratio:
|
230 |
+
mask_ratio = torch.rand(1).item() * self.target_tokens_mask_ratio
|
231 |
+
else:
|
232 |
+
mask_ratio = self.target_tokens_mask_ratio
|
233 |
+
|
234 |
+
time_seq_mask = (torch.rand((batchsize, time_seq_len)) > mask_ratio).to(
|
235 |
+
target_embeds.device
|
236 |
+
)
|
237 |
+
# Mask the target embedding
|
238 |
+
target_embeds = target_embeds * time_seq_mask.unsqueeze(-1)
|
239 |
+
target_embeds_attn_mask = target_embeds_attn_mask * time_seq_mask
|
240 |
+
return target_embeds, target_embeds_attn_mask, time_seq_mask
|
241 |
+
|
242 |
+
def generate_partial(self, batch, cond_dict=None, no_grad=False):
|
243 |
+
if cond_dict is None:
|
244 |
+
cond_dict = self.get_input(batch)
|
245 |
+
|
246 |
+
print("Generate partially prompted audio with in-context learning")
|
247 |
+
# self.model.train()
|
248 |
+
# assert self.model.training==True
|
249 |
+
|
250 |
+
target_embeds, target_embeds_attn_mask = (
|
251 |
+
cond_dict["crossattn_audiomae_pooled"][0],
|
252 |
+
cond_dict["crossattn_audiomae_pooled"][1],
|
253 |
+
)
|
254 |
+
|
255 |
+
target_time_steps = target_embeds.size(1)
|
256 |
+
|
257 |
+
(
|
258 |
+
input_embeds,
|
259 |
+
input_embeds_attn_mask,
|
260 |
+
cond_sequence_end_time_idx,
|
261 |
+
) = self.get_input_sequence_and_mask(cond_dict)
|
262 |
+
|
263 |
+
model_input = torch.cat(
|
264 |
+
[input_embeds, target_embeds[:, : target_time_steps // 4, :]], dim=1
|
265 |
+
)
|
266 |
+
model_input_mask = torch.cat(
|
267 |
+
[
|
268 |
+
input_embeds_attn_mask,
|
269 |
+
target_embeds_attn_mask[:, : target_time_steps // 4],
|
270 |
+
],
|
271 |
+
dim=1,
|
272 |
+
)
|
273 |
+
|
274 |
+
steps = self.mae_token_num
|
275 |
+
|
276 |
+
for _ in range(3 * steps // 4):
|
277 |
+
output = self.model(
|
278 |
+
inputs_embeds=model_input, attention_mask=model_input_mask
|
279 |
+
)["last_hidden_state"]
|
280 |
+
# Update the model input
|
281 |
+
model_input = torch.cat([model_input, output[:, -1:, :]], dim=1)
|
282 |
+
# Update the attention mask
|
283 |
+
attention_mask_new_step = torch.ones((model_input_mask.size(0), 1)).to(
|
284 |
+
model_input.device
|
285 |
+
)
|
286 |
+
model_input_mask = torch.cat(
|
287 |
+
[model_input_mask, attention_mask_new_step], dim=1
|
288 |
+
)
|
289 |
+
|
290 |
+
output = model_input[:, cond_sequence_end_time_idx:]
|
291 |
+
|
292 |
+
return output, cond_dict
|
293 |
+
|
294 |
+
def generate(self, batch, cond_dict=None, no_grad=False):
|
295 |
+
if cond_dict is None:
|
296 |
+
cond_dict = self.get_input(batch)
|
297 |
+
|
298 |
+
# self.model.train()
|
299 |
+
# print("!!!!!!!!!!!!!train")
|
300 |
+
|
301 |
+
(
|
302 |
+
input_embeds,
|
303 |
+
input_embeds_attn_mask,
|
304 |
+
cond_sequence_end_time_idx,
|
305 |
+
) = self.get_input_sequence_and_mask(cond_dict)
|
306 |
+
model_input = input_embeds
|
307 |
+
model_input_mask = input_embeds_attn_mask
|
308 |
+
|
309 |
+
steps = self.mae_token_num
|
310 |
+
|
311 |
+
for _ in range(steps):
|
312 |
+
output = self.model(
|
313 |
+
inputs_embeds=model_input, attention_mask=model_input_mask
|
314 |
+
)["last_hidden_state"]
|
315 |
+
# Update the model input
|
316 |
+
model_input = torch.cat([model_input, output[:, -1:, :]], dim=1)
|
317 |
+
# Update the attention mask
|
318 |
+
attention_mask_new_step = torch.ones((model_input_mask.size(0), 1)).to(
|
319 |
+
model_input.device
|
320 |
+
)
|
321 |
+
model_input_mask = torch.cat(
|
322 |
+
[model_input_mask, attention_mask_new_step], dim=1
|
323 |
+
)
|
324 |
+
|
325 |
+
return model_input[:, cond_sequence_end_time_idx:], cond_dict
|
326 |
+
|
327 |
+
def get_input_item(self, batch, k):
|
328 |
+
fname, text, waveform, stft, fbank = (
|
329 |
+
batch["fname"],
|
330 |
+
batch["text"],
|
331 |
+
batch["waveform"],
|
332 |
+
batch["stft"],
|
333 |
+
batch["log_mel_spec"],
|
334 |
+
)
|
335 |
+
ret = {}
|
336 |
+
|
337 |
+
ret["fbank"] = (
|
338 |
+
fbank.unsqueeze(1).to(memory_format=torch.contiguous_format).float()
|
339 |
+
)
|
340 |
+
ret["stft"] = stft.to(memory_format=torch.contiguous_format).float()
|
341 |
+
# ret["clip_label"] = clip_label.to(memory_format=torch.contiguous_format).float()
|
342 |
+
ret["waveform"] = waveform.to(memory_format=torch.contiguous_format).float()
|
343 |
+
ret["text"] = list(text)
|
344 |
+
ret["fname"] = fname
|
345 |
+
|
346 |
+
for key in batch.keys():
|
347 |
+
if key not in ret.keys():
|
348 |
+
ret[key] = batch[key]
|
349 |
+
|
350 |
+
return ret[k]
|
351 |
+
|
352 |
+
def get_input(self, batch):
|
353 |
+
cond_dict = {}
|
354 |
+
if len(self.cond_stage_model_metadata.keys()) > 0:
|
355 |
+
unconditional_cfg = False
|
356 |
+
|
357 |
+
for cond_model_key in self.cond_stage_model_metadata.keys():
|
358 |
+
cond_stage_key = self.cond_stage_model_metadata[cond_model_key][
|
359 |
+
"cond_stage_key"
|
360 |
+
]
|
361 |
+
|
362 |
+
# if(not self.training):
|
363 |
+
# if(isinstance(self.cond_stage_models[self.cond_stage_model_metadata[cond_model_key]["model_idx"]], CLAPAudioEmbeddingClassifierFreev2)):
|
364 |
+
# assert cond_stage_key == "text" # CLAP model should use text for evaluation
|
365 |
+
|
366 |
+
# The original data for conditioning
|
367 |
+
xc = self.get_input_item(batch, cond_stage_key)
|
368 |
+
if type(xc) == torch.Tensor:
|
369 |
+
xc = xc.to(self.device)
|
370 |
+
|
371 |
+
c = self.get_learned_conditioning(
|
372 |
+
xc, key=cond_model_key, unconditional_cfg=unconditional_cfg
|
373 |
+
)
|
374 |
+
cond_dict[cond_model_key] = c
|
375 |
+
|
376 |
+
return cond_dict
|
377 |
+
|
378 |
+
def instantiate_cond_stage(self, config):
|
379 |
+
self.cond_stage_model_metadata = {}
|
380 |
+
|
381 |
+
for i, cond_model_key in enumerate(config.keys()):
|
382 |
+
model = instantiate_from_config(config[cond_model_key])
|
383 |
+
self.cond_stage_models.append(model)
|
384 |
+
self.cond_stage_model_metadata[cond_model_key] = {
|
385 |
+
"model_idx": i,
|
386 |
+
"cond_stage_key": config[cond_model_key]["cond_stage_key"],
|
387 |
+
"conditioning_key": config[cond_model_key]["conditioning_key"],
|
388 |
+
}
|
389 |
+
|
390 |
+
def get_learned_conditioning(self, c, key, unconditional_cfg):
|
391 |
+
assert key in self.cond_stage_model_metadata.keys()
|
392 |
+
|
393 |
+
# Classifier-free guidance
|
394 |
+
if not unconditional_cfg:
|
395 |
+
c = self.cond_stage_models[
|
396 |
+
self.cond_stage_model_metadata[key]["model_idx"]
|
397 |
+
](c)
|
398 |
+
else:
|
399 |
+
if isinstance(c, torch.Tensor):
|
400 |
+
batchsize = c.size(0)
|
401 |
+
elif isinstance(c, list):
|
402 |
+
batchsize = len(c)
|
403 |
+
else:
|
404 |
+
raise NotImplementedError()
|
405 |
+
c = self.cond_stage_models[
|
406 |
+
self.cond_stage_model_metadata[key]["model_idx"]
|
407 |
+
].get_unconditional_condition(batchsize)
|
408 |
+
|
409 |
+
return c
|
410 |
+
|
411 |
+
def initialize_param_check_toolkit(self):
|
412 |
+
self.tracked_steps = 0
|
413 |
+
self.param_dict = {}
|
414 |
+
|
415 |
+
def statistic_require_grad_tensor_number(self, module, name=None):
|
416 |
+
requires_grad_num = 0
|
417 |
+
total_num = 0
|
418 |
+
require_grad_tensor = None
|
419 |
+
for p in module.parameters():
|
420 |
+
if p.requires_grad:
|
421 |
+
requires_grad_num += 1
|
422 |
+
if require_grad_tensor is None:
|
423 |
+
require_grad_tensor = p
|
424 |
+
total_num += 1
|
425 |
+
print(
|
426 |
+
"Module: [%s] have %s trainable parameters out of %s total parameters (%.2f)"
|
427 |
+
% (name, requires_grad_num, total_num, requires_grad_num / total_num)
|
428 |
+
)
|
429 |
+
return require_grad_tensor
|
audioldm2/audiomae_gen/utils.py
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch.nn as nn
|
2 |
+
|
3 |
+
|
4 |
+
class Prenet(nn.Module):
|
5 |
+
def __init__(self, in_dim, sizes=[256, 128], dropout_rate=0.5):
|
6 |
+
super(Prenet, self).__init__()
|
7 |
+
in_sizes = [in_dim] + sizes[:-1]
|
8 |
+
self.layers = nn.ModuleList(
|
9 |
+
[
|
10 |
+
nn.Linear(in_size, out_size)
|
11 |
+
for (in_size, out_size) in zip(in_sizes, sizes)
|
12 |
+
]
|
13 |
+
)
|
14 |
+
self.relu = nn.ReLU()
|
15 |
+
self.dropout = nn.Dropout(dropout_rate)
|
16 |
+
|
17 |
+
def forward(self, inputs):
|
18 |
+
for linear in self.layers:
|
19 |
+
inputs = self.dropout(self.relu(linear(inputs)))
|
20 |
+
return inputs
|
21 |
+
|
22 |
+
|
23 |
+
if __name__ == "__main__":
|
24 |
+
model = Prenet(in_dim=128, sizes=[256, 256, 128])
|
25 |
+
import ipdb
|
26 |
+
|
27 |
+
ipdb.set_trace()
|
audioldm2/clap/__init__.py
ADDED
File without changes
|
audioldm2/clap/open_clip/__init__.py
ADDED
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from .factory import (
|
2 |
+
list_models,
|
3 |
+
create_model,
|
4 |
+
create_model_and_transforms,
|
5 |
+
add_model_config,
|
6 |
+
)
|
7 |
+
from .loss import ClipLoss, gather_features, LPLoss, lp_gather_features, LPMetrics
|
8 |
+
from .model import (
|
9 |
+
CLAP,
|
10 |
+
CLAPTextCfg,
|
11 |
+
CLAPVisionCfg,
|
12 |
+
CLAPAudioCfp,
|
13 |
+
convert_weights_to_fp16,
|
14 |
+
trace_model,
|
15 |
+
)
|
16 |
+
from .openai import load_openai_model, list_openai_models
|
17 |
+
from .pretrained import (
|
18 |
+
list_pretrained,
|
19 |
+
list_pretrained_tag_models,
|
20 |
+
list_pretrained_model_tags,
|
21 |
+
get_pretrained_url,
|
22 |
+
download_pretrained,
|
23 |
+
)
|
24 |
+
from .tokenizer import SimpleTokenizer, tokenize
|
25 |
+
from .transform import image_transform
|
audioldm2/clap/open_clip/bpe_simple_vocab_16e6.txt.gz
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
|
3 |
+
size 1356917
|
audioldm2/clap/open_clip/factory.py
ADDED
@@ -0,0 +1,276 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
import logging
|
3 |
+
import os
|
4 |
+
import re
|
5 |
+
from copy import deepcopy
|
6 |
+
from pathlib import Path
|
7 |
+
|
8 |
+
import torch
|
9 |
+
|
10 |
+
from .model import CLAP, convert_weights_to_fp16
|
11 |
+
from .openai import load_openai_model
|
12 |
+
from .pretrained import get_pretrained_url, download_pretrained
|
13 |
+
from .transform import image_transform
|
14 |
+
|
15 |
+
_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
|
16 |
+
_MODEL_CONFIGS = {} # directory (model_name: config) of model architecture configs
|
17 |
+
|
18 |
+
|
19 |
+
def _natural_key(string_):
|
20 |
+
return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", string_.lower())]
|
21 |
+
|
22 |
+
|
23 |
+
def _rescan_model_configs():
|
24 |
+
global _MODEL_CONFIGS
|
25 |
+
|
26 |
+
config_ext = (".json",)
|
27 |
+
config_files = []
|
28 |
+
for config_path in _MODEL_CONFIG_PATHS:
|
29 |
+
if config_path.is_file() and config_path.suffix in config_ext:
|
30 |
+
config_files.append(config_path)
|
31 |
+
elif config_path.is_dir():
|
32 |
+
for ext in config_ext:
|
33 |
+
config_files.extend(config_path.glob(f"*{ext}"))
|
34 |
+
|
35 |
+
for cf in config_files:
|
36 |
+
if os.path.basename(cf)[0] == ".":
|
37 |
+
continue # Ignore hidden files
|
38 |
+
|
39 |
+
with open(cf, "r") as f:
|
40 |
+
model_cfg = json.load(f)
|
41 |
+
if all(a in model_cfg for a in ("embed_dim", "audio_cfg", "text_cfg")):
|
42 |
+
_MODEL_CONFIGS[cf.stem] = model_cfg
|
43 |
+
|
44 |
+
_MODEL_CONFIGS = {
|
45 |
+
k: v
|
46 |
+
for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))
|
47 |
+
}
|
48 |
+
|
49 |
+
|
50 |
+
_rescan_model_configs() # initial populate of model config registry
|
51 |
+
|
52 |
+
|
53 |
+
def load_state_dict(checkpoint_path: str, map_location="cpu", skip_params=True):
|
54 |
+
checkpoint = torch.load(checkpoint_path, map_location=map_location)
|
55 |
+
if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
|
56 |
+
state_dict = checkpoint["state_dict"]
|
57 |
+
else:
|
58 |
+
state_dict = checkpoint
|
59 |
+
if skip_params:
|
60 |
+
if next(iter(state_dict.items()))[0].startswith("module"):
|
61 |
+
state_dict = {k[7:]: v for k, v in state_dict.items()}
|
62 |
+
# for k in state_dict:
|
63 |
+
# if k.startswith('transformer'):
|
64 |
+
# v = state_dict.pop(k)
|
65 |
+
# state_dict['text_branch.' + k[12:]] = v
|
66 |
+
return state_dict
|
67 |
+
|
68 |
+
|
69 |
+
def create_model(
|
70 |
+
amodel_name: str,
|
71 |
+
tmodel_name: str,
|
72 |
+
pretrained: str = "",
|
73 |
+
precision: str = "fp32",
|
74 |
+
device: torch.device = torch.device("cpu"),
|
75 |
+
jit: bool = False,
|
76 |
+
force_quick_gelu: bool = False,
|
77 |
+
openai_model_cache_dir: str = os.path.expanduser("~/.cache/clip"),
|
78 |
+
skip_params=True,
|
79 |
+
pretrained_audio: str = "",
|
80 |
+
pretrained_text: str = "",
|
81 |
+
enable_fusion: bool = False,
|
82 |
+
fusion_type: str = "None"
|
83 |
+
# pretrained_image: bool = False,
|
84 |
+
):
|
85 |
+
amodel_name = amodel_name.replace(
|
86 |
+
"/", "-"
|
87 |
+
) # for callers using old naming with / in ViT names
|
88 |
+
pretrained_orig = pretrained
|
89 |
+
pretrained = pretrained.lower()
|
90 |
+
if pretrained == "openai":
|
91 |
+
if amodel_name in _MODEL_CONFIGS:
|
92 |
+
logging.info(f"Loading {amodel_name} model config.")
|
93 |
+
model_cfg = deepcopy(_MODEL_CONFIGS[amodel_name])
|
94 |
+
else:
|
95 |
+
logging.error(
|
96 |
+
f"Model config for {amodel_name} not found; available models {list_models()}."
|
97 |
+
)
|
98 |
+
raise RuntimeError(f"Model config for {amodel_name} not found.")
|
99 |
+
|
100 |
+
logging.info(f"Loading pretrained ViT-B-16 text encoder from OpenAI.")
|
101 |
+
# Hard Code in model name
|
102 |
+
model_cfg["text_cfg"]["model_type"] = tmodel_name
|
103 |
+
model = load_openai_model(
|
104 |
+
"ViT-B-16",
|
105 |
+
model_cfg,
|
106 |
+
device=device,
|
107 |
+
jit=jit,
|
108 |
+
cache_dir=openai_model_cache_dir,
|
109 |
+
enable_fusion=enable_fusion,
|
110 |
+
fusion_type=fusion_type,
|
111 |
+
)
|
112 |
+
# See https://discuss.pytorch.org/t/valueerror-attemting-to-unscale-fp16-gradients/81372
|
113 |
+
if precision == "amp" or precision == "fp32":
|
114 |
+
model = model.float()
|
115 |
+
else:
|
116 |
+
if amodel_name in _MODEL_CONFIGS:
|
117 |
+
logging.info(f"Loading {amodel_name} model config.")
|
118 |
+
model_cfg = deepcopy(_MODEL_CONFIGS[amodel_name])
|
119 |
+
else:
|
120 |
+
logging.error(
|
121 |
+
f"Model config for {amodel_name} not found; available models {list_models()}."
|
122 |
+
)
|
123 |
+
raise RuntimeError(f"Model config for {amodel_name} not found.")
|
124 |
+
|
125 |
+
if force_quick_gelu:
|
126 |
+
# override for use of QuickGELU on non-OpenAI transformer models
|
127 |
+
model_cfg["quick_gelu"] = True
|
128 |
+
|
129 |
+
# if pretrained_image:
|
130 |
+
# if 'timm_amodel_name' in model_cfg.get('vision_cfg', {}):
|
131 |
+
# # pretrained weight loading for timm models set via vision_cfg
|
132 |
+
# model_cfg['vision_cfg']['timm_model_pretrained'] = True
|
133 |
+
# else:
|
134 |
+
# assert False, 'pretrained image towers currently only supported for timm models'
|
135 |
+
model_cfg["text_cfg"]["model_type"] = tmodel_name
|
136 |
+
model_cfg["enable_fusion"] = enable_fusion
|
137 |
+
model_cfg["fusion_type"] = fusion_type
|
138 |
+
model = CLAP(**model_cfg)
|
139 |
+
|
140 |
+
if pretrained:
|
141 |
+
checkpoint_path = ""
|
142 |
+
url = get_pretrained_url(amodel_name, pretrained)
|
143 |
+
if url:
|
144 |
+
checkpoint_path = download_pretrained(url, root=openai_model_cache_dir)
|
145 |
+
elif os.path.exists(pretrained_orig):
|
146 |
+
checkpoint_path = pretrained_orig
|
147 |
+
if checkpoint_path:
|
148 |
+
logging.info(
|
149 |
+
f"Loading pretrained {amodel_name}-{tmodel_name} weights ({pretrained})."
|
150 |
+
)
|
151 |
+
ckpt = load_state_dict(checkpoint_path, skip_params=True)
|
152 |
+
model.load_state_dict(ckpt)
|
153 |
+
param_names = [n for n, p in model.named_parameters()]
|
154 |
+
# for n in param_names:
|
155 |
+
# print(n, "\t", "Loaded" if n in ckpt else "Unloaded")
|
156 |
+
else:
|
157 |
+
logging.warning(
|
158 |
+
f"Pretrained weights ({pretrained}) not found for model {amodel_name}."
|
159 |
+
)
|
160 |
+
raise RuntimeError(
|
161 |
+
f"Pretrained weights ({pretrained}) not found for model {amodel_name}."
|
162 |
+
)
|
163 |
+
|
164 |
+
if pretrained_audio:
|
165 |
+
if amodel_name.startswith("PANN"):
|
166 |
+
if "Cnn14_mAP" in pretrained_audio: # official checkpoint
|
167 |
+
audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
|
168 |
+
audio_ckpt = audio_ckpt["model"]
|
169 |
+
keys = list(audio_ckpt.keys())
|
170 |
+
for key in keys:
|
171 |
+
if (
|
172 |
+
"spectrogram_extractor" not in key
|
173 |
+
and "logmel_extractor" not in key
|
174 |
+
):
|
175 |
+
v = audio_ckpt.pop(key)
|
176 |
+
audio_ckpt["audio_branch." + key] = v
|
177 |
+
elif os.path.basename(pretrained_audio).startswith(
|
178 |
+
"PANN"
|
179 |
+
): # checkpoint trained via HTSAT codebase
|
180 |
+
audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
|
181 |
+
audio_ckpt = audio_ckpt["state_dict"]
|
182 |
+
keys = list(audio_ckpt.keys())
|
183 |
+
for key in keys:
|
184 |
+
if key.startswith("sed_model"):
|
185 |
+
v = audio_ckpt.pop(key)
|
186 |
+
audio_ckpt["audio_branch." + key[10:]] = v
|
187 |
+
elif os.path.basename(pretrained_audio).startswith(
|
188 |
+
"finetuned"
|
189 |
+
): # checkpoint trained via linear probe codebase
|
190 |
+
audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
|
191 |
+
else:
|
192 |
+
raise ValueError("Unknown audio checkpoint")
|
193 |
+
elif amodel_name.startswith("HTSAT"):
|
194 |
+
if "HTSAT_AudioSet_Saved" in pretrained_audio: # official checkpoint
|
195 |
+
audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
|
196 |
+
audio_ckpt = audio_ckpt["state_dict"]
|
197 |
+
keys = list(audio_ckpt.keys())
|
198 |
+
for key in keys:
|
199 |
+
if key.startswith("sed_model") and (
|
200 |
+
"spectrogram_extractor" not in key
|
201 |
+
and "logmel_extractor" not in key
|
202 |
+
):
|
203 |
+
v = audio_ckpt.pop(key)
|
204 |
+
audio_ckpt["audio_branch." + key[10:]] = v
|
205 |
+
elif os.path.basename(pretrained_audio).startswith(
|
206 |
+
"HTSAT"
|
207 |
+
): # checkpoint trained via HTSAT codebase
|
208 |
+
audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
|
209 |
+
audio_ckpt = audio_ckpt["state_dict"]
|
210 |
+
keys = list(audio_ckpt.keys())
|
211 |
+
for key in keys:
|
212 |
+
if key.startswith("sed_model"):
|
213 |
+
v = audio_ckpt.pop(key)
|
214 |
+
audio_ckpt["audio_branch." + key[10:]] = v
|
215 |
+
elif os.path.basename(pretrained_audio).startswith(
|
216 |
+
"finetuned"
|
217 |
+
): # checkpoint trained via linear probe codebase
|
218 |
+
audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
|
219 |
+
else:
|
220 |
+
raise ValueError("Unknown audio checkpoint")
|
221 |
+
else:
|
222 |
+
raise f"this audio encoder pretrained checkpoint is not support"
|
223 |
+
|
224 |
+
model.load_state_dict(audio_ckpt, strict=False)
|
225 |
+
logging.info(
|
226 |
+
f"Loading pretrained {amodel_name} weights ({pretrained_audio})."
|
227 |
+
)
|
228 |
+
param_names = [n for n, p in model.named_parameters()]
|
229 |
+
for n in param_names:
|
230 |
+
print(n, "\t", "Loaded" if n in audio_ckpt else "Unloaded")
|
231 |
+
|
232 |
+
model.to(device=device)
|
233 |
+
if precision == "fp16":
|
234 |
+
assert device.type != "cpu"
|
235 |
+
convert_weights_to_fp16(model)
|
236 |
+
|
237 |
+
if jit:
|
238 |
+
model = torch.jit.script(model)
|
239 |
+
|
240 |
+
return model, model_cfg
|
241 |
+
|
242 |
+
|
243 |
+
def create_model_and_transforms(
|
244 |
+
model_name: str,
|
245 |
+
pretrained: str = "",
|
246 |
+
precision: str = "fp32",
|
247 |
+
device: torch.device = torch.device("cpu"),
|
248 |
+
jit: bool = False,
|
249 |
+
force_quick_gelu: bool = False,
|
250 |
+
# pretrained_image: bool = False,
|
251 |
+
):
|
252 |
+
model = create_model(
|
253 |
+
model_name,
|
254 |
+
pretrained,
|
255 |
+
precision,
|
256 |
+
device,
|
257 |
+
jit,
|
258 |
+
force_quick_gelu=force_quick_gelu,
|
259 |
+
# pretrained_image=pretrained_image
|
260 |
+
)
|
261 |
+
preprocess_train = image_transform(model.visual.image_size, is_train=True)
|
262 |
+
preprocess_val = image_transform(model.visual.image_size, is_train=False)
|
263 |
+
return model, preprocess_train, preprocess_val
|
264 |
+
|
265 |
+
|
266 |
+
def list_models():
|
267 |
+
"""enumerate available model architectures based on config files"""
|
268 |
+
return list(_MODEL_CONFIGS.keys())
|
269 |
+
|
270 |
+
|
271 |
+
def add_model_config(path):
|
272 |
+
"""add model config path or file and update registry"""
|
273 |
+
if not isinstance(path, Path):
|
274 |
+
path = Path(path)
|
275 |
+
_MODEL_CONFIG_PATHS.append(path)
|
276 |
+
_rescan_model_configs()
|
audioldm2/clap/open_clip/feature_fusion.py
ADDED
@@ -0,0 +1,192 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""
|
2 |
+
Feature Fusion for Varible-Length Data Processing
|
3 |
+
AFF/iAFF is referred and modified from https://github.com/YimianDai/open-aff/blob/master/aff_pytorch/aff_net/fusion.py
|
4 |
+
According to the paper: Yimian Dai et al, Attentional Feature Fusion, IEEE Winter Conference on Applications of Computer Vision, WACV 2021
|
5 |
+
"""
|
6 |
+
|
7 |
+
import torch
|
8 |
+
import torch.nn as nn
|
9 |
+
|
10 |
+
|
11 |
+
class DAF(nn.Module):
|
12 |
+
"""
|
13 |
+
直接相加 DirectAddFuse
|
14 |
+
"""
|
15 |
+
|
16 |
+
def __init__(self):
|
17 |
+
super(DAF, self).__init__()
|
18 |
+
|
19 |
+
def forward(self, x, residual):
|
20 |
+
return x + residual
|
21 |
+
|
22 |
+
|
23 |
+
class iAFF(nn.Module):
|
24 |
+
"""
|
25 |
+
多特征融合 iAFF
|
26 |
+
"""
|
27 |
+
|
28 |
+
def __init__(self, channels=64, r=4, type="2D"):
|
29 |
+
super(iAFF, self).__init__()
|
30 |
+
inter_channels = int(channels // r)
|
31 |
+
|
32 |
+
if type == "1D":
|
33 |
+
# 本地注意力
|
34 |
+
self.local_att = nn.Sequential(
|
35 |
+
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
36 |
+
nn.BatchNorm1d(inter_channels),
|
37 |
+
nn.ReLU(inplace=True),
|
38 |
+
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
39 |
+
nn.BatchNorm1d(channels),
|
40 |
+
)
|
41 |
+
|
42 |
+
# 全局注意力
|
43 |
+
self.global_att = nn.Sequential(
|
44 |
+
nn.AdaptiveAvgPool1d(1),
|
45 |
+
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
46 |
+
nn.BatchNorm1d(inter_channels),
|
47 |
+
nn.ReLU(inplace=True),
|
48 |
+
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
49 |
+
nn.BatchNorm1d(channels),
|
50 |
+
)
|
51 |
+
|
52 |
+
# 第二次本地注意力
|
53 |
+
self.local_att2 = nn.Sequential(
|
54 |
+
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
55 |
+
nn.BatchNorm1d(inter_channels),
|
56 |
+
nn.ReLU(inplace=True),
|
57 |
+
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
58 |
+
nn.BatchNorm1d(channels),
|
59 |
+
)
|
60 |
+
# 第二次全局注意力
|
61 |
+
self.global_att2 = nn.Sequential(
|
62 |
+
nn.AdaptiveAvgPool1d(1),
|
63 |
+
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
64 |
+
nn.BatchNorm1d(inter_channels),
|
65 |
+
nn.ReLU(inplace=True),
|
66 |
+
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
67 |
+
nn.BatchNorm1d(channels),
|
68 |
+
)
|
69 |
+
elif type == "2D":
|
70 |
+
# 本地注意力
|
71 |
+
self.local_att = nn.Sequential(
|
72 |
+
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
73 |
+
nn.BatchNorm2d(inter_channels),
|
74 |
+
nn.ReLU(inplace=True),
|
75 |
+
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
76 |
+
nn.BatchNorm2d(channels),
|
77 |
+
)
|
78 |
+
|
79 |
+
# 全局注意力
|
80 |
+
self.global_att = nn.Sequential(
|
81 |
+
nn.AdaptiveAvgPool2d(1),
|
82 |
+
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
83 |
+
nn.BatchNorm2d(inter_channels),
|
84 |
+
nn.ReLU(inplace=True),
|
85 |
+
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
86 |
+
nn.BatchNorm2d(channels),
|
87 |
+
)
|
88 |
+
|
89 |
+
# 第二次本地注意力
|
90 |
+
self.local_att2 = nn.Sequential(
|
91 |
+
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
92 |
+
nn.BatchNorm2d(inter_channels),
|
93 |
+
nn.ReLU(inplace=True),
|
94 |
+
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
95 |
+
nn.BatchNorm2d(channels),
|
96 |
+
)
|
97 |
+
# 第二次全局注意力
|
98 |
+
self.global_att2 = nn.Sequential(
|
99 |
+
nn.AdaptiveAvgPool2d(1),
|
100 |
+
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
101 |
+
nn.BatchNorm2d(inter_channels),
|
102 |
+
nn.ReLU(inplace=True),
|
103 |
+
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
104 |
+
nn.BatchNorm2d(channels),
|
105 |
+
)
|
106 |
+
else:
|
107 |
+
raise f"the type is not supported"
|
108 |
+
|
109 |
+
self.sigmoid = nn.Sigmoid()
|
110 |
+
|
111 |
+
def forward(self, x, residual):
|
112 |
+
flag = False
|
113 |
+
xa = x + residual
|
114 |
+
if xa.size(0) == 1:
|
115 |
+
xa = torch.cat([xa, xa], dim=0)
|
116 |
+
flag = True
|
117 |
+
xl = self.local_att(xa)
|
118 |
+
xg = self.global_att(xa)
|
119 |
+
xlg = xl + xg
|
120 |
+
wei = self.sigmoid(xlg)
|
121 |
+
xi = x * wei + residual * (1 - wei)
|
122 |
+
|
123 |
+
xl2 = self.local_att2(xi)
|
124 |
+
xg2 = self.global_att(xi)
|
125 |
+
xlg2 = xl2 + xg2
|
126 |
+
wei2 = self.sigmoid(xlg2)
|
127 |
+
xo = x * wei2 + residual * (1 - wei2)
|
128 |
+
if flag:
|
129 |
+
xo = xo[0].unsqueeze(0)
|
130 |
+
return xo
|
131 |
+
|
132 |
+
|
133 |
+
class AFF(nn.Module):
|
134 |
+
"""
|
135 |
+
多特征融合 AFF
|
136 |
+
"""
|
137 |
+
|
138 |
+
def __init__(self, channels=64, r=4, type="2D"):
|
139 |
+
super(AFF, self).__init__()
|
140 |
+
inter_channels = int(channels // r)
|
141 |
+
|
142 |
+
if type == "1D":
|
143 |
+
self.local_att = nn.Sequential(
|
144 |
+
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
145 |
+
nn.BatchNorm1d(inter_channels),
|
146 |
+
nn.ReLU(inplace=True),
|
147 |
+
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
148 |
+
nn.BatchNorm1d(channels),
|
149 |
+
)
|
150 |
+
self.global_att = nn.Sequential(
|
151 |
+
nn.AdaptiveAvgPool1d(1),
|
152 |
+
nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
153 |
+
nn.BatchNorm1d(inter_channels),
|
154 |
+
nn.ReLU(inplace=True),
|
155 |
+
nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
156 |
+
nn.BatchNorm1d(channels),
|
157 |
+
)
|
158 |
+
elif type == "2D":
|
159 |
+
self.local_att = nn.Sequential(
|
160 |
+
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
161 |
+
nn.BatchNorm2d(inter_channels),
|
162 |
+
nn.ReLU(inplace=True),
|
163 |
+
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
164 |
+
nn.BatchNorm2d(channels),
|
165 |
+
)
|
166 |
+
self.global_att = nn.Sequential(
|
167 |
+
nn.AdaptiveAvgPool2d(1),
|
168 |
+
nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
|
169 |
+
nn.BatchNorm2d(inter_channels),
|
170 |
+
nn.ReLU(inplace=True),
|
171 |
+
nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
|
172 |
+
nn.BatchNorm2d(channels),
|
173 |
+
)
|
174 |
+
else:
|
175 |
+
raise f"the type is not supported."
|
176 |
+
|
177 |
+
self.sigmoid = nn.Sigmoid()
|
178 |
+
|
179 |
+
def forward(self, x, residual):
|
180 |
+
flag = False
|
181 |
+
xa = x + residual
|
182 |
+
if xa.size(0) == 1:
|
183 |
+
xa = torch.cat([xa, xa], dim=0)
|
184 |
+
flag = True
|
185 |
+
xl = self.local_att(xa)
|
186 |
+
xg = self.global_att(xa)
|
187 |
+
xlg = xl + xg
|
188 |
+
wei = self.sigmoid(xlg)
|
189 |
+
xo = 2 * x * wei + 2 * residual * (1 - wei)
|
190 |
+
if flag:
|
191 |
+
xo = xo[0].unsqueeze(0)
|
192 |
+
return xo
|
audioldm2/clap/open_clip/htsat.py
ADDED
@@ -0,0 +1,1304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Ke Chen
|
2 | |
3 |
+
# HTS-AT: A HIERARCHICAL TOKEN-SEMANTIC AUDIO TRANSFORMER FOR SOUND CLASSIFICATION AND DETECTION
|
4 |
+
# Some layers designed on the model
|
5 |
+
# below codes are based and referred from https://github.com/microsoft/Swin-Transformer
|
6 |
+
# Swin Transformer for Computer Vision: https://arxiv.org/pdf/2103.14030.pdf
|
7 |
+
|
8 |
+
import torch
|
9 |
+
import torch.nn as nn
|
10 |
+
from itertools import repeat
|
11 |
+
import collections.abc
|
12 |
+
import math
|
13 |
+
import warnings
|
14 |
+
|
15 |
+
from torch.nn.init import _calculate_fan_in_and_fan_out
|
16 |
+
import torch.utils.checkpoint as checkpoint
|
17 |
+
|
18 |
+
import random
|
19 |
+
|
20 |
+
from torchlibrosa.stft import Spectrogram, LogmelFilterBank
|
21 |
+
from torchlibrosa.augmentation import SpecAugmentation
|
22 |
+
|
23 |
+
from itertools import repeat
|
24 |
+
from .utils import do_mixup, interpolate
|
25 |
+
|
26 |
+
from .feature_fusion import iAFF, AFF, DAF
|
27 |
+
|
28 |
+
|
29 |
+
# from PyTorch internals
|
30 |
+
def _ntuple(n):
|
31 |
+
def parse(x):
|
32 |
+
if isinstance(x, collections.abc.Iterable):
|
33 |
+
return x
|
34 |
+
return tuple(repeat(x, n))
|
35 |
+
|
36 |
+
return parse
|
37 |
+
|
38 |
+
|
39 |
+
to_1tuple = _ntuple(1)
|
40 |
+
to_2tuple = _ntuple(2)
|
41 |
+
to_3tuple = _ntuple(3)
|
42 |
+
to_4tuple = _ntuple(4)
|
43 |
+
to_ntuple = _ntuple
|
44 |
+
|
45 |
+
|
46 |
+
def drop_path(x, drop_prob: float = 0.0, training: bool = False):
|
47 |
+
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
48 |
+
This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
|
49 |
+
the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
|
50 |
+
See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
|
51 |
+
changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
|
52 |
+
'survival rate' as the argument.
|
53 |
+
"""
|
54 |
+
if drop_prob == 0.0 or not training:
|
55 |
+
return x
|
56 |
+
keep_prob = 1 - drop_prob
|
57 |
+
shape = (x.shape[0],) + (1,) * (
|
58 |
+
x.ndim - 1
|
59 |
+
) # work with diff dim tensors, not just 2D ConvNets
|
60 |
+
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
|
61 |
+
random_tensor.floor_() # binarize
|
62 |
+
output = x.div(keep_prob) * random_tensor
|
63 |
+
return output
|
64 |
+
|
65 |
+
|
66 |
+
class DropPath(nn.Module):
|
67 |
+
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
|
68 |
+
|
69 |
+
def __init__(self, drop_prob=None):
|
70 |
+
super(DropPath, self).__init__()
|
71 |
+
self.drop_prob = drop_prob
|
72 |
+
|
73 |
+
def forward(self, x):
|
74 |
+
return drop_path(x, self.drop_prob, self.training)
|
75 |
+
|
76 |
+
|
77 |
+
class PatchEmbed(nn.Module):
|
78 |
+
"""2D Image to Patch Embedding"""
|
79 |
+
|
80 |
+
def __init__(
|
81 |
+
self,
|
82 |
+
img_size=224,
|
83 |
+
patch_size=16,
|
84 |
+
in_chans=3,
|
85 |
+
embed_dim=768,
|
86 |
+
norm_layer=None,
|
87 |
+
flatten=True,
|
88 |
+
patch_stride=16,
|
89 |
+
enable_fusion=False,
|
90 |
+
fusion_type="None",
|
91 |
+
):
|
92 |
+
super().__init__()
|
93 |
+
img_size = to_2tuple(img_size)
|
94 |
+
patch_size = to_2tuple(patch_size)
|
95 |
+
patch_stride = to_2tuple(patch_stride)
|
96 |
+
self.img_size = img_size
|
97 |
+
self.patch_size = patch_size
|
98 |
+
self.patch_stride = patch_stride
|
99 |
+
self.grid_size = (
|
100 |
+
img_size[0] // patch_stride[0],
|
101 |
+
img_size[1] // patch_stride[1],
|
102 |
+
)
|
103 |
+
self.num_patches = self.grid_size[0] * self.grid_size[1]
|
104 |
+
self.flatten = flatten
|
105 |
+
self.in_chans = in_chans
|
106 |
+
self.embed_dim = embed_dim
|
107 |
+
|
108 |
+
self.enable_fusion = enable_fusion
|
109 |
+
self.fusion_type = fusion_type
|
110 |
+
|
111 |
+
padding = (
|
112 |
+
(patch_size[0] - patch_stride[0]) // 2,
|
113 |
+
(patch_size[1] - patch_stride[1]) // 2,
|
114 |
+
)
|
115 |
+
|
116 |
+
if (self.enable_fusion) and (self.fusion_type == "channel_map"):
|
117 |
+
self.proj = nn.Conv2d(
|
118 |
+
in_chans * 4,
|
119 |
+
embed_dim,
|
120 |
+
kernel_size=patch_size,
|
121 |
+
stride=patch_stride,
|
122 |
+
padding=padding,
|
123 |
+
)
|
124 |
+
else:
|
125 |
+
self.proj = nn.Conv2d(
|
126 |
+
in_chans,
|
127 |
+
embed_dim,
|
128 |
+
kernel_size=patch_size,
|
129 |
+
stride=patch_stride,
|
130 |
+
padding=padding,
|
131 |
+
)
|
132 |
+
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
133 |
+
|
134 |
+
if (self.enable_fusion) and (
|
135 |
+
self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d"]
|
136 |
+
):
|
137 |
+
self.mel_conv2d = nn.Conv2d(
|
138 |
+
in_chans,
|
139 |
+
embed_dim,
|
140 |
+
kernel_size=(patch_size[0], patch_size[1] * 3),
|
141 |
+
stride=(patch_stride[0], patch_stride[1] * 3),
|
142 |
+
padding=padding,
|
143 |
+
)
|
144 |
+
if self.fusion_type == "daf_2d":
|
145 |
+
self.fusion_model = DAF()
|
146 |
+
elif self.fusion_type == "aff_2d":
|
147 |
+
self.fusion_model = AFF(channels=embed_dim, type="2D")
|
148 |
+
elif self.fusion_type == "iaff_2d":
|
149 |
+
self.fusion_model = iAFF(channels=embed_dim, type="2D")
|
150 |
+
|
151 |
+
def forward(self, x, longer_idx=None):
|
152 |
+
if (self.enable_fusion) and (
|
153 |
+
self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d"]
|
154 |
+
):
|
155 |
+
global_x = x[:, 0:1, :, :]
|
156 |
+
|
157 |
+
# global processing
|
158 |
+
B, C, H, W = global_x.shape
|
159 |
+
assert (
|
160 |
+
H == self.img_size[0] and W == self.img_size[1]
|
161 |
+
), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
|
162 |
+
global_x = self.proj(global_x)
|
163 |
+
TW = global_x.size(-1)
|
164 |
+
if len(longer_idx) > 0:
|
165 |
+
# local processing
|
166 |
+
local_x = x[longer_idx, 1:, :, :].contiguous()
|
167 |
+
B, C, H, W = local_x.shape
|
168 |
+
local_x = local_x.view(B * C, 1, H, W)
|
169 |
+
local_x = self.mel_conv2d(local_x)
|
170 |
+
local_x = local_x.view(
|
171 |
+
B, C, local_x.size(1), local_x.size(2), local_x.size(3)
|
172 |
+
)
|
173 |
+
local_x = local_x.permute((0, 2, 3, 1, 4)).contiguous().flatten(3)
|
174 |
+
TB, TC, TH, _ = local_x.size()
|
175 |
+
if local_x.size(-1) < TW:
|
176 |
+
local_x = torch.cat(
|
177 |
+
[
|
178 |
+
local_x,
|
179 |
+
torch.zeros(
|
180 |
+
(TB, TC, TH, TW - local_x.size(-1)),
|
181 |
+
device=global_x.device,
|
182 |
+
),
|
183 |
+
],
|
184 |
+
dim=-1,
|
185 |
+
)
|
186 |
+
else:
|
187 |
+
local_x = local_x[:, :, :, :TW]
|
188 |
+
|
189 |
+
global_x[longer_idx] = self.fusion_model(global_x[longer_idx], local_x)
|
190 |
+
x = global_x
|
191 |
+
else:
|
192 |
+
B, C, H, W = x.shape
|
193 |
+
assert (
|
194 |
+
H == self.img_size[0] and W == self.img_size[1]
|
195 |
+
), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
|
196 |
+
x = self.proj(x)
|
197 |
+
|
198 |
+
if self.flatten:
|
199 |
+
x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
|
200 |
+
x = self.norm(x)
|
201 |
+
return x
|
202 |
+
|
203 |
+
|
204 |
+
class Mlp(nn.Module):
|
205 |
+
"""MLP as used in Vision Transformer, MLP-Mixer and related networks"""
|
206 |
+
|
207 |
+
def __init__(
|
208 |
+
self,
|
209 |
+
in_features,
|
210 |
+
hidden_features=None,
|
211 |
+
out_features=None,
|
212 |
+
act_layer=nn.GELU,
|
213 |
+
drop=0.0,
|
214 |
+
):
|
215 |
+
super().__init__()
|
216 |
+
out_features = out_features or in_features
|
217 |
+
hidden_features = hidden_features or in_features
|
218 |
+
self.fc1 = nn.Linear(in_features, hidden_features)
|
219 |
+
self.act = act_layer()
|
220 |
+
self.fc2 = nn.Linear(hidden_features, out_features)
|
221 |
+
self.drop = nn.Dropout(drop)
|
222 |
+
|
223 |
+
def forward(self, x):
|
224 |
+
x = self.fc1(x)
|
225 |
+
x = self.act(x)
|
226 |
+
x = self.drop(x)
|
227 |
+
x = self.fc2(x)
|
228 |
+
x = self.drop(x)
|
229 |
+
return x
|
230 |
+
|
231 |
+
|
232 |
+
def _no_grad_trunc_normal_(tensor, mean, std, a, b):
|
233 |
+
# Cut & paste from PyTorch official master until it's in a few official releases - RW
|
234 |
+
# Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
|
235 |
+
def norm_cdf(x):
|
236 |
+
# Computes standard normal cumulative distribution function
|
237 |
+
return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
|
238 |
+
|
239 |
+
if (mean < a - 2 * std) or (mean > b + 2 * std):
|
240 |
+
warnings.warn(
|
241 |
+
"mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
|
242 |
+
"The distribution of values may be incorrect.",
|
243 |
+
stacklevel=2,
|
244 |
+
)
|
245 |
+
|
246 |
+
with torch.no_grad():
|
247 |
+
# Values are generated by using a truncated uniform distribution and
|
248 |
+
# then using the inverse CDF for the normal distribution.
|
249 |
+
# Get upper and lower cdf values
|
250 |
+
l = norm_cdf((a - mean) / std)
|
251 |
+
u = norm_cdf((b - mean) / std)
|
252 |
+
|
253 |
+
# Uniformly fill tensor with values from [l, u], then translate to
|
254 |
+
# [2l-1, 2u-1].
|
255 |
+
tensor.uniform_(2 * l - 1, 2 * u - 1)
|
256 |
+
|
257 |
+
# Use inverse cdf transform for normal distribution to get truncated
|
258 |
+
# standard normal
|
259 |
+
tensor.erfinv_()
|
260 |
+
|
261 |
+
# Transform to proper mean, std
|
262 |
+
tensor.mul_(std * math.sqrt(2.0))
|
263 |
+
tensor.add_(mean)
|
264 |
+
|
265 |
+
# Clamp to ensure it's in the proper range
|
266 |
+
tensor.clamp_(min=a, max=b)
|
267 |
+
return tensor
|
268 |
+
|
269 |
+
|
270 |
+
def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
|
271 |
+
# type: (Tensor, float, float, float, float) -> Tensor
|
272 |
+
r"""Fills the input Tensor with values drawn from a truncated
|
273 |
+
normal distribution. The values are effectively drawn from the
|
274 |
+
normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
|
275 |
+
with values outside :math:`[a, b]` redrawn until they are within
|
276 |
+
the bounds. The method used for generating the random values works
|
277 |
+
best when :math:`a \leq \text{mean} \leq b`.
|
278 |
+
Args:
|
279 |
+
tensor: an n-dimensional `torch.Tensor`
|
280 |
+
mean: the mean of the normal distribution
|
281 |
+
std: the standard deviation of the normal distribution
|
282 |
+
a: the minimum cutoff value
|
283 |
+
b: the maximum cutoff value
|
284 |
+
Examples:
|
285 |
+
>>> w = torch.empty(3, 5)
|
286 |
+
>>> nn.init.trunc_normal_(w)
|
287 |
+
"""
|
288 |
+
return _no_grad_trunc_normal_(tensor, mean, std, a, b)
|
289 |
+
|
290 |
+
|
291 |
+
def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):
|
292 |
+
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
|
293 |
+
if mode == "fan_in":
|
294 |
+
denom = fan_in
|
295 |
+
elif mode == "fan_out":
|
296 |
+
denom = fan_out
|
297 |
+
elif mode == "fan_avg":
|
298 |
+
denom = (fan_in + fan_out) / 2
|
299 |
+
|
300 |
+
variance = scale / denom
|
301 |
+
|
302 |
+
if distribution == "truncated_normal":
|
303 |
+
# constant is stddev of standard normal truncated to (-2, 2)
|
304 |
+
trunc_normal_(tensor, std=math.sqrt(variance) / 0.87962566103423978)
|
305 |
+
elif distribution == "normal":
|
306 |
+
tensor.normal_(std=math.sqrt(variance))
|
307 |
+
elif distribution == "uniform":
|
308 |
+
bound = math.sqrt(3 * variance)
|
309 |
+
tensor.uniform_(-bound, bound)
|
310 |
+
else:
|
311 |
+
raise ValueError(f"invalid distribution {distribution}")
|
312 |
+
|
313 |
+
|
314 |
+
def lecun_normal_(tensor):
|
315 |
+
variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")
|
316 |
+
|
317 |
+
|
318 |
+
def window_partition(x, window_size):
|
319 |
+
"""
|
320 |
+
Args:
|
321 |
+
x: (B, H, W, C)
|
322 |
+
window_size (int): window size
|
323 |
+
Returns:
|
324 |
+
windows: (num_windows*B, window_size, window_size, C)
|
325 |
+
"""
|
326 |
+
B, H, W, C = x.shape
|
327 |
+
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
|
328 |
+
windows = (
|
329 |
+
x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
|
330 |
+
)
|
331 |
+
return windows
|
332 |
+
|
333 |
+
|
334 |
+
def window_reverse(windows, window_size, H, W):
|
335 |
+
"""
|
336 |
+
Args:
|
337 |
+
windows: (num_windows*B, window_size, window_size, C)
|
338 |
+
window_size (int): Window size
|
339 |
+
H (int): Height of image
|
340 |
+
W (int): Width of image
|
341 |
+
Returns:
|
342 |
+
x: (B, H, W, C)
|
343 |
+
"""
|
344 |
+
B = int(windows.shape[0] / (H * W / window_size / window_size))
|
345 |
+
x = windows.view(
|
346 |
+
B, H // window_size, W // window_size, window_size, window_size, -1
|
347 |
+
)
|
348 |
+
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
|
349 |
+
return x
|
350 |
+
|
351 |
+
|
352 |
+
class WindowAttention(nn.Module):
|
353 |
+
r"""Window based multi-head self attention (W-MSA) module with relative position bias.
|
354 |
+
It supports both of shifted and non-shifted window.
|
355 |
+
Args:
|
356 |
+
dim (int): Number of input channels.
|
357 |
+
window_size (tuple[int]): The height and width of the window.
|
358 |
+
num_heads (int): Number of attention heads.
|
359 |
+
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
|
360 |
+
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
|
361 |
+
attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
|
362 |
+
proj_drop (float, optional): Dropout ratio of output. Default: 0.0
|
363 |
+
"""
|
364 |
+
|
365 |
+
def __init__(
|
366 |
+
self,
|
367 |
+
dim,
|
368 |
+
window_size,
|
369 |
+
num_heads,
|
370 |
+
qkv_bias=True,
|
371 |
+
qk_scale=None,
|
372 |
+
attn_drop=0.0,
|
373 |
+
proj_drop=0.0,
|
374 |
+
):
|
375 |
+
super().__init__()
|
376 |
+
self.dim = dim
|
377 |
+
self.window_size = window_size # Wh, Ww
|
378 |
+
self.num_heads = num_heads
|
379 |
+
head_dim = dim // num_heads
|
380 |
+
self.scale = qk_scale or head_dim**-0.5
|
381 |
+
|
382 |
+
# define a parameter table of relative position bias
|
383 |
+
self.relative_position_bias_table = nn.Parameter(
|
384 |
+
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
|
385 |
+
) # 2*Wh-1 * 2*Ww-1, nH
|
386 |
+
|
387 |
+
# get pair-wise relative position index for each token inside the window
|
388 |
+
coords_h = torch.arange(self.window_size[0])
|
389 |
+
coords_w = torch.arange(self.window_size[1])
|
390 |
+
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
|
391 |
+
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
|
392 |
+
relative_coords = (
|
393 |
+
coords_flatten[:, :, None] - coords_flatten[:, None, :]
|
394 |
+
) # 2, Wh*Ww, Wh*Ww
|
395 |
+
relative_coords = relative_coords.permute(
|
396 |
+
1, 2, 0
|
397 |
+
).contiguous() # Wh*Ww, Wh*Ww, 2
|
398 |
+
relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
|
399 |
+
relative_coords[:, :, 1] += self.window_size[1] - 1
|
400 |
+
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
|
401 |
+
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
|
402 |
+
self.register_buffer("relative_position_index", relative_position_index)
|
403 |
+
|
404 |
+
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
405 |
+
self.attn_drop = nn.Dropout(attn_drop)
|
406 |
+
self.proj = nn.Linear(dim, dim)
|
407 |
+
self.proj_drop = nn.Dropout(proj_drop)
|
408 |
+
|
409 |
+
trunc_normal_(self.relative_position_bias_table, std=0.02)
|
410 |
+
self.softmax = nn.Softmax(dim=-1)
|
411 |
+
|
412 |
+
def forward(self, x, mask=None):
|
413 |
+
"""
|
414 |
+
Args:
|
415 |
+
x: input features with shape of (num_windows*B, N, C)
|
416 |
+
mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
|
417 |
+
"""
|
418 |
+
B_, N, C = x.shape
|
419 |
+
qkv = (
|
420 |
+
self.qkv(x)
|
421 |
+
.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
|
422 |
+
.permute(2, 0, 3, 1, 4)
|
423 |
+
)
|
424 |
+
q, k, v = (
|
425 |
+
qkv[0],
|
426 |
+
qkv[1],
|
427 |
+
qkv[2],
|
428 |
+
) # make torchscript happy (cannot use tensor as tuple)
|
429 |
+
|
430 |
+
q = q * self.scale
|
431 |
+
attn = q @ k.transpose(-2, -1)
|
432 |
+
|
433 |
+
relative_position_bias = self.relative_position_bias_table[
|
434 |
+
self.relative_position_index.view(-1)
|
435 |
+
].view(
|
436 |
+
self.window_size[0] * self.window_size[1],
|
437 |
+
self.window_size[0] * self.window_size[1],
|
438 |
+
-1,
|
439 |
+
) # Wh*Ww,Wh*Ww,nH
|
440 |
+
relative_position_bias = relative_position_bias.permute(
|
441 |
+
2, 0, 1
|
442 |
+
).contiguous() # nH, Wh*Ww, Wh*Ww
|
443 |
+
attn = attn + relative_position_bias.unsqueeze(0)
|
444 |
+
|
445 |
+
if mask is not None:
|
446 |
+
nW = mask.shape[0]
|
447 |
+
attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(
|
448 |
+
1
|
449 |
+
).unsqueeze(0)
|
450 |
+
attn = attn.view(-1, self.num_heads, N, N)
|
451 |
+
attn = self.softmax(attn)
|
452 |
+
else:
|
453 |
+
attn = self.softmax(attn)
|
454 |
+
|
455 |
+
attn = self.attn_drop(attn)
|
456 |
+
|
457 |
+
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
|
458 |
+
x = self.proj(x)
|
459 |
+
x = self.proj_drop(x)
|
460 |
+
return x, attn
|
461 |
+
|
462 |
+
def extra_repr(self):
|
463 |
+
return f"dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}"
|
464 |
+
|
465 |
+
|
466 |
+
# We use the model based on Swintransformer Block, therefore we can use the swin-transformer pretrained model
|
467 |
+
class SwinTransformerBlock(nn.Module):
|
468 |
+
r"""Swin Transformer Block.
|
469 |
+
Args:
|
470 |
+
dim (int): Number of input channels.
|
471 |
+
input_resolution (tuple[int]): Input resulotion.
|
472 |
+
num_heads (int): Number of attention heads.
|
473 |
+
window_size (int): Window size.
|
474 |
+
shift_size (int): Shift size for SW-MSA.
|
475 |
+
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
|
476 |
+
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
|
477 |
+
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
|
478 |
+
drop (float, optional): Dropout rate. Default: 0.0
|
479 |
+
attn_drop (float, optional): Attention dropout rate. Default: 0.0
|
480 |
+
drop_path (float, optional): Stochastic depth rate. Default: 0.0
|
481 |
+
act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
|
482 |
+
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
|
483 |
+
"""
|
484 |
+
|
485 |
+
def __init__(
|
486 |
+
self,
|
487 |
+
dim,
|
488 |
+
input_resolution,
|
489 |
+
num_heads,
|
490 |
+
window_size=7,
|
491 |
+
shift_size=0,
|
492 |
+
mlp_ratio=4.0,
|
493 |
+
qkv_bias=True,
|
494 |
+
qk_scale=None,
|
495 |
+
drop=0.0,
|
496 |
+
attn_drop=0.0,
|
497 |
+
drop_path=0.0,
|
498 |
+
act_layer=nn.GELU,
|
499 |
+
norm_layer=nn.LayerNorm,
|
500 |
+
norm_before_mlp="ln",
|
501 |
+
):
|
502 |
+
super().__init__()
|
503 |
+
self.dim = dim
|
504 |
+
self.input_resolution = input_resolution
|
505 |
+
self.num_heads = num_heads
|
506 |
+
self.window_size = window_size
|
507 |
+
self.shift_size = shift_size
|
508 |
+
self.mlp_ratio = mlp_ratio
|
509 |
+
self.norm_before_mlp = norm_before_mlp
|
510 |
+
if min(self.input_resolution) <= self.window_size:
|
511 |
+
# if window size is larger than input resolution, we don't partition windows
|
512 |
+
self.shift_size = 0
|
513 |
+
self.window_size = min(self.input_resolution)
|
514 |
+
assert (
|
515 |
+
0 <= self.shift_size < self.window_size
|
516 |
+
), "shift_size must in 0-window_size"
|
517 |
+
|
518 |
+
self.norm1 = norm_layer(dim)
|
519 |
+
self.attn = WindowAttention(
|
520 |
+
dim,
|
521 |
+
window_size=to_2tuple(self.window_size),
|
522 |
+
num_heads=num_heads,
|
523 |
+
qkv_bias=qkv_bias,
|
524 |
+
qk_scale=qk_scale,
|
525 |
+
attn_drop=attn_drop,
|
526 |
+
proj_drop=drop,
|
527 |
+
)
|
528 |
+
|
529 |
+
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
|
530 |
+
if self.norm_before_mlp == "ln":
|
531 |
+
self.norm2 = nn.LayerNorm(dim)
|
532 |
+
elif self.norm_before_mlp == "bn":
|
533 |
+
self.norm2 = lambda x: nn.BatchNorm1d(dim)(x.transpose(1, 2)).transpose(
|
534 |
+
1, 2
|
535 |
+
)
|
536 |
+
else:
|
537 |
+
raise NotImplementedError
|
538 |
+
mlp_hidden_dim = int(dim * mlp_ratio)
|
539 |
+
self.mlp = Mlp(
|
540 |
+
in_features=dim,
|
541 |
+
hidden_features=mlp_hidden_dim,
|
542 |
+
act_layer=act_layer,
|
543 |
+
drop=drop,
|
544 |
+
)
|
545 |
+
|
546 |
+
if self.shift_size > 0:
|
547 |
+
# calculate attention mask for SW-MSA
|
548 |
+
H, W = self.input_resolution
|
549 |
+
img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1
|
550 |
+
h_slices = (
|
551 |
+
slice(0, -self.window_size),
|
552 |
+
slice(-self.window_size, -self.shift_size),
|
553 |
+
slice(-self.shift_size, None),
|
554 |
+
)
|
555 |
+
w_slices = (
|
556 |
+
slice(0, -self.window_size),
|
557 |
+
slice(-self.window_size, -self.shift_size),
|
558 |
+
slice(-self.shift_size, None),
|
559 |
+
)
|
560 |
+
cnt = 0
|
561 |
+
for h in h_slices:
|
562 |
+
for w in w_slices:
|
563 |
+
img_mask[:, h, w, :] = cnt
|
564 |
+
cnt += 1
|
565 |
+
|
566 |
+
mask_windows = window_partition(
|
567 |
+
img_mask, self.window_size
|
568 |
+
) # nW, window_size, window_size, 1
|
569 |
+
mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
|
570 |
+
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
|
571 |
+
attn_mask = attn_mask.masked_fill(
|
572 |
+
attn_mask != 0, float(-100.0)
|
573 |
+
).masked_fill(attn_mask == 0, float(0.0))
|
574 |
+
else:
|
575 |
+
attn_mask = None
|
576 |
+
|
577 |
+
self.register_buffer("attn_mask", attn_mask)
|
578 |
+
|
579 |
+
def forward(self, x):
|
580 |
+
# pdb.set_trace()
|
581 |
+
H, W = self.input_resolution
|
582 |
+
# print("H: ", H)
|
583 |
+
# print("W: ", W)
|
584 |
+
# pdb.set_trace()
|
585 |
+
B, L, C = x.shape
|
586 |
+
# assert L == H * W, "input feature has wrong size"
|
587 |
+
|
588 |
+
shortcut = x
|
589 |
+
x = self.norm1(x)
|
590 |
+
x = x.view(B, H, W, C)
|
591 |
+
|
592 |
+
# cyclic shift
|
593 |
+
if self.shift_size > 0:
|
594 |
+
shifted_x = torch.roll(
|
595 |
+
x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)
|
596 |
+
)
|
597 |
+
else:
|
598 |
+
shifted_x = x
|
599 |
+
|
600 |
+
# partition windows
|
601 |
+
x_windows = window_partition(
|
602 |
+
shifted_x, self.window_size
|
603 |
+
) # nW*B, window_size, window_size, C
|
604 |
+
x_windows = x_windows.view(
|
605 |
+
-1, self.window_size * self.window_size, C
|
606 |
+
) # nW*B, window_size*window_size, C
|
607 |
+
|
608 |
+
# W-MSA/SW-MSA
|
609 |
+
attn_windows, attn = self.attn(
|
610 |
+
x_windows, mask=self.attn_mask
|
611 |
+
) # nW*B, window_size*window_size, C
|
612 |
+
|
613 |
+
# merge windows
|
614 |
+
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
|
615 |
+
shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C
|
616 |
+
|
617 |
+
# reverse cyclic shift
|
618 |
+
if self.shift_size > 0:
|
619 |
+
x = torch.roll(
|
620 |
+
shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)
|
621 |
+
)
|
622 |
+
else:
|
623 |
+
x = shifted_x
|
624 |
+
x = x.view(B, H * W, C)
|
625 |
+
|
626 |
+
# FFN
|
627 |
+
x = shortcut + self.drop_path(x)
|
628 |
+
x = x + self.drop_path(self.mlp(self.norm2(x)))
|
629 |
+
|
630 |
+
return x, attn
|
631 |
+
|
632 |
+
def extra_repr(self):
|
633 |
+
return (
|
634 |
+
f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, "
|
635 |
+
f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
|
636 |
+
)
|
637 |
+
|
638 |
+
|
639 |
+
class PatchMerging(nn.Module):
|
640 |
+
r"""Patch Merging Layer.
|
641 |
+
Args:
|
642 |
+
input_resolution (tuple[int]): Resolution of input feature.
|
643 |
+
dim (int): Number of input channels.
|
644 |
+
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
|
645 |
+
"""
|
646 |
+
|
647 |
+
def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
|
648 |
+
super().__init__()
|
649 |
+
self.input_resolution = input_resolution
|
650 |
+
self.dim = dim
|
651 |
+
self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
|
652 |
+
self.norm = norm_layer(4 * dim)
|
653 |
+
|
654 |
+
def forward(self, x):
|
655 |
+
"""
|
656 |
+
x: B, H*W, C
|
657 |
+
"""
|
658 |
+
H, W = self.input_resolution
|
659 |
+
B, L, C = x.shape
|
660 |
+
assert L == H * W, "input feature has wrong size"
|
661 |
+
assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
|
662 |
+
|
663 |
+
x = x.view(B, H, W, C)
|
664 |
+
|
665 |
+
x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C
|
666 |
+
x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C
|
667 |
+
x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C
|
668 |
+
x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C
|
669 |
+
x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C
|
670 |
+
x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C
|
671 |
+
|
672 |
+
x = self.norm(x)
|
673 |
+
x = self.reduction(x)
|
674 |
+
|
675 |
+
return x
|
676 |
+
|
677 |
+
def extra_repr(self):
|
678 |
+
return f"input_resolution={self.input_resolution}, dim={self.dim}"
|
679 |
+
|
680 |
+
|
681 |
+
class BasicLayer(nn.Module):
|
682 |
+
"""A basic Swin Transformer layer for one stage.
|
683 |
+
Args:
|
684 |
+
dim (int): Number of input channels.
|
685 |
+
input_resolution (tuple[int]): Input resolution.
|
686 |
+
depth (int): Number of blocks.
|
687 |
+
num_heads (int): Number of attention heads.
|
688 |
+
window_size (int): Local window size.
|
689 |
+
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
|
690 |
+
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
|
691 |
+
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
|
692 |
+
drop (float, optional): Dropout rate. Default: 0.0
|
693 |
+
attn_drop (float, optional): Attention dropout rate. Default: 0.0
|
694 |
+
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
|
695 |
+
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
|
696 |
+
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
|
697 |
+
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
|
698 |
+
"""
|
699 |
+
|
700 |
+
def __init__(
|
701 |
+
self,
|
702 |
+
dim,
|
703 |
+
input_resolution,
|
704 |
+
depth,
|
705 |
+
num_heads,
|
706 |
+
window_size,
|
707 |
+
mlp_ratio=4.0,
|
708 |
+
qkv_bias=True,
|
709 |
+
qk_scale=None,
|
710 |
+
drop=0.0,
|
711 |
+
attn_drop=0.0,
|
712 |
+
drop_path=0.0,
|
713 |
+
norm_layer=nn.LayerNorm,
|
714 |
+
downsample=None,
|
715 |
+
use_checkpoint=False,
|
716 |
+
norm_before_mlp="ln",
|
717 |
+
):
|
718 |
+
super().__init__()
|
719 |
+
self.dim = dim
|
720 |
+
self.input_resolution = input_resolution
|
721 |
+
self.depth = depth
|
722 |
+
self.use_checkpoint = use_checkpoint
|
723 |
+
|
724 |
+
# build blocks
|
725 |
+
self.blocks = nn.ModuleList(
|
726 |
+
[
|
727 |
+
SwinTransformerBlock(
|
728 |
+
dim=dim,
|
729 |
+
input_resolution=input_resolution,
|
730 |
+
num_heads=num_heads,
|
731 |
+
window_size=window_size,
|
732 |
+
shift_size=0 if (i % 2 == 0) else window_size // 2,
|
733 |
+
mlp_ratio=mlp_ratio,
|
734 |
+
qkv_bias=qkv_bias,
|
735 |
+
qk_scale=qk_scale,
|
736 |
+
drop=drop,
|
737 |
+
attn_drop=attn_drop,
|
738 |
+
drop_path=drop_path[i]
|
739 |
+
if isinstance(drop_path, list)
|
740 |
+
else drop_path,
|
741 |
+
norm_layer=norm_layer,
|
742 |
+
norm_before_mlp=norm_before_mlp,
|
743 |
+
)
|
744 |
+
for i in range(depth)
|
745 |
+
]
|
746 |
+
)
|
747 |
+
|
748 |
+
# patch merging layer
|
749 |
+
if downsample is not None:
|
750 |
+
self.downsample = downsample(
|
751 |
+
input_resolution, dim=dim, norm_layer=norm_layer
|
752 |
+
)
|
753 |
+
else:
|
754 |
+
self.downsample = None
|
755 |
+
|
756 |
+
def forward(self, x):
|
757 |
+
attns = []
|
758 |
+
for blk in self.blocks:
|
759 |
+
if self.use_checkpoint:
|
760 |
+
x = checkpoint.checkpoint(blk, x)
|
761 |
+
else:
|
762 |
+
x, attn = blk(x)
|
763 |
+
if not self.training:
|
764 |
+
attns.append(attn.unsqueeze(0))
|
765 |
+
if self.downsample is not None:
|
766 |
+
x = self.downsample(x)
|
767 |
+
if not self.training:
|
768 |
+
attn = torch.cat(attns, dim=0)
|
769 |
+
attn = torch.mean(attn, dim=0)
|
770 |
+
return x, attn
|
771 |
+
|
772 |
+
def extra_repr(self):
|
773 |
+
return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
|
774 |
+
|
775 |
+
|
776 |
+
# The Core of HTSAT
|
777 |
+
class HTSAT_Swin_Transformer(nn.Module):
|
778 |
+
r"""HTSAT based on the Swin Transformer
|
779 |
+
Args:
|
780 |
+
spec_size (int | tuple(int)): Input Spectrogram size. Default 256
|
781 |
+
patch_size (int | tuple(int)): Patch size. Default: 4
|
782 |
+
path_stride (iot | tuple(int)): Patch Stride for Frequency and Time Axis. Default: 4
|
783 |
+
in_chans (int): Number of input image channels. Default: 1 (mono)
|
784 |
+
num_classes (int): Number of classes for classification head. Default: 527
|
785 |
+
embed_dim (int): Patch embedding dimension. Default: 96
|
786 |
+
depths (tuple(int)): Depth of each HTSAT-Swin Transformer layer.
|
787 |
+
num_heads (tuple(int)): Number of attention heads in different layers.
|
788 |
+
window_size (int): Window size. Default: 8
|
789 |
+
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
|
790 |
+
qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
|
791 |
+
qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
|
792 |
+
drop_rate (float): Dropout rate. Default: 0
|
793 |
+
attn_drop_rate (float): Attention dropout rate. Default: 0
|
794 |
+
drop_path_rate (float): Stochastic depth rate. Default: 0.1
|
795 |
+
norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
|
796 |
+
ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
|
797 |
+
patch_norm (bool): If True, add normalization after patch embedding. Default: True
|
798 |
+
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
|
799 |
+
config (module): The configuration Module from config.py
|
800 |
+
"""
|
801 |
+
|
802 |
+
def __init__(
|
803 |
+
self,
|
804 |
+
spec_size=256,
|
805 |
+
patch_size=4,
|
806 |
+
patch_stride=(4, 4),
|
807 |
+
in_chans=1,
|
808 |
+
num_classes=527,
|
809 |
+
embed_dim=96,
|
810 |
+
depths=[2, 2, 6, 2],
|
811 |
+
num_heads=[4, 8, 16, 32],
|
812 |
+
window_size=8,
|
813 |
+
mlp_ratio=4.0,
|
814 |
+
qkv_bias=True,
|
815 |
+
qk_scale=None,
|
816 |
+
drop_rate=0.0,
|
817 |
+
attn_drop_rate=0.0,
|
818 |
+
drop_path_rate=0.1,
|
819 |
+
norm_layer=nn.LayerNorm,
|
820 |
+
ape=False,
|
821 |
+
patch_norm=True,
|
822 |
+
use_checkpoint=False,
|
823 |
+
norm_before_mlp="ln",
|
824 |
+
config=None,
|
825 |
+
enable_fusion=False,
|
826 |
+
fusion_type="None",
|
827 |
+
**kwargs,
|
828 |
+
):
|
829 |
+
super(HTSAT_Swin_Transformer, self).__init__()
|
830 |
+
|
831 |
+
self.config = config
|
832 |
+
self.spec_size = spec_size
|
833 |
+
self.patch_stride = patch_stride
|
834 |
+
self.patch_size = patch_size
|
835 |
+
self.window_size = window_size
|
836 |
+
self.embed_dim = embed_dim
|
837 |
+
self.depths = depths
|
838 |
+
self.ape = ape
|
839 |
+
self.in_chans = in_chans
|
840 |
+
self.num_classes = num_classes
|
841 |
+
self.num_heads = num_heads
|
842 |
+
self.num_layers = len(self.depths)
|
843 |
+
self.num_features = int(self.embed_dim * 2 ** (self.num_layers - 1))
|
844 |
+
|
845 |
+
self.drop_rate = drop_rate
|
846 |
+
self.attn_drop_rate = attn_drop_rate
|
847 |
+
self.drop_path_rate = drop_path_rate
|
848 |
+
|
849 |
+
self.qkv_bias = qkv_bias
|
850 |
+
self.qk_scale = None
|
851 |
+
|
852 |
+
self.patch_norm = patch_norm
|
853 |
+
self.norm_layer = norm_layer if self.patch_norm else None
|
854 |
+
self.norm_before_mlp = norm_before_mlp
|
855 |
+
self.mlp_ratio = mlp_ratio
|
856 |
+
|
857 |
+
self.use_checkpoint = use_checkpoint
|
858 |
+
|
859 |
+
self.enable_fusion = enable_fusion
|
860 |
+
self.fusion_type = fusion_type
|
861 |
+
|
862 |
+
# process mel-spec ; used only once
|
863 |
+
self.freq_ratio = self.spec_size // self.config.mel_bins
|
864 |
+
window = "hann"
|
865 |
+
center = True
|
866 |
+
pad_mode = "reflect"
|
867 |
+
ref = 1.0
|
868 |
+
amin = 1e-10
|
869 |
+
top_db = None
|
870 |
+
self.interpolate_ratio = 32 # Downsampled ratio
|
871 |
+
# Spectrogram extractor
|
872 |
+
self.spectrogram_extractor = Spectrogram(
|
873 |
+
n_fft=config.window_size,
|
874 |
+
hop_length=config.hop_size,
|
875 |
+
win_length=config.window_size,
|
876 |
+
window=window,
|
877 |
+
center=center,
|
878 |
+
pad_mode=pad_mode,
|
879 |
+
freeze_parameters=True,
|
880 |
+
)
|
881 |
+
# Logmel feature extractor
|
882 |
+
self.logmel_extractor = LogmelFilterBank(
|
883 |
+
sr=config.sample_rate,
|
884 |
+
n_fft=config.window_size,
|
885 |
+
n_mels=config.mel_bins,
|
886 |
+
fmin=config.fmin,
|
887 |
+
fmax=config.fmax,
|
888 |
+
ref=ref,
|
889 |
+
amin=amin,
|
890 |
+
top_db=top_db,
|
891 |
+
freeze_parameters=True,
|
892 |
+
)
|
893 |
+
# Spec augmenter
|
894 |
+
self.spec_augmenter = SpecAugmentation(
|
895 |
+
time_drop_width=64,
|
896 |
+
time_stripes_num=2,
|
897 |
+
freq_drop_width=8,
|
898 |
+
freq_stripes_num=2,
|
899 |
+
) # 2 2
|
900 |
+
self.bn0 = nn.BatchNorm2d(self.config.mel_bins)
|
901 |
+
|
902 |
+
# split spctrogram into non-overlapping patches
|
903 |
+
self.patch_embed = PatchEmbed(
|
904 |
+
img_size=self.spec_size,
|
905 |
+
patch_size=self.patch_size,
|
906 |
+
in_chans=self.in_chans,
|
907 |
+
embed_dim=self.embed_dim,
|
908 |
+
norm_layer=self.norm_layer,
|
909 |
+
patch_stride=patch_stride,
|
910 |
+
enable_fusion=self.enable_fusion,
|
911 |
+
fusion_type=self.fusion_type,
|
912 |
+
)
|
913 |
+
|
914 |
+
num_patches = self.patch_embed.num_patches
|
915 |
+
patches_resolution = self.patch_embed.grid_size
|
916 |
+
self.patches_resolution = patches_resolution
|
917 |
+
|
918 |
+
# absolute position embedding
|
919 |
+
if self.ape:
|
920 |
+
self.absolute_pos_embed = nn.Parameter(
|
921 |
+
torch.zeros(1, num_patches, self.embed_dim)
|
922 |
+
)
|
923 |
+
trunc_normal_(self.absolute_pos_embed, std=0.02)
|
924 |
+
|
925 |
+
self.pos_drop = nn.Dropout(p=self.drop_rate)
|
926 |
+
|
927 |
+
# stochastic depth
|
928 |
+
dpr = [
|
929 |
+
x.item() for x in torch.linspace(0, self.drop_path_rate, sum(self.depths))
|
930 |
+
] # stochastic depth decay rule
|
931 |
+
|
932 |
+
# build layers
|
933 |
+
self.layers = nn.ModuleList()
|
934 |
+
for i_layer in range(self.num_layers):
|
935 |
+
layer = BasicLayer(
|
936 |
+
dim=int(self.embed_dim * 2**i_layer),
|
937 |
+
input_resolution=(
|
938 |
+
patches_resolution[0] // (2**i_layer),
|
939 |
+
patches_resolution[1] // (2**i_layer),
|
940 |
+
),
|
941 |
+
depth=self.depths[i_layer],
|
942 |
+
num_heads=self.num_heads[i_layer],
|
943 |
+
window_size=self.window_size,
|
944 |
+
mlp_ratio=self.mlp_ratio,
|
945 |
+
qkv_bias=self.qkv_bias,
|
946 |
+
qk_scale=self.qk_scale,
|
947 |
+
drop=self.drop_rate,
|
948 |
+
attn_drop=self.attn_drop_rate,
|
949 |
+
drop_path=dpr[
|
950 |
+
sum(self.depths[:i_layer]) : sum(self.depths[: i_layer + 1])
|
951 |
+
],
|
952 |
+
norm_layer=self.norm_layer,
|
953 |
+
downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
|
954 |
+
use_checkpoint=use_checkpoint,
|
955 |
+
norm_before_mlp=self.norm_before_mlp,
|
956 |
+
)
|
957 |
+
self.layers.append(layer)
|
958 |
+
|
959 |
+
self.norm = self.norm_layer(self.num_features)
|
960 |
+
self.avgpool = nn.AdaptiveAvgPool1d(1)
|
961 |
+
self.maxpool = nn.AdaptiveMaxPool1d(1)
|
962 |
+
|
963 |
+
SF = (
|
964 |
+
self.spec_size
|
965 |
+
// (2 ** (len(self.depths) - 1))
|
966 |
+
// self.patch_stride[0]
|
967 |
+
// self.freq_ratio
|
968 |
+
)
|
969 |
+
self.tscam_conv = nn.Conv2d(
|
970 |
+
in_channels=self.num_features,
|
971 |
+
out_channels=self.num_classes,
|
972 |
+
kernel_size=(SF, 3),
|
973 |
+
padding=(0, 1),
|
974 |
+
)
|
975 |
+
self.head = nn.Linear(num_classes, num_classes)
|
976 |
+
|
977 |
+
if (self.enable_fusion) and (
|
978 |
+
self.fusion_type in ["daf_1d", "aff_1d", "iaff_1d"]
|
979 |
+
):
|
980 |
+
self.mel_conv1d = nn.Sequential(
|
981 |
+
nn.Conv1d(64, 64, kernel_size=5, stride=3, padding=2),
|
982 |
+
nn.BatchNorm1d(64),
|
983 |
+
)
|
984 |
+
if self.fusion_type == "daf_1d":
|
985 |
+
self.fusion_model = DAF()
|
986 |
+
elif self.fusion_type == "aff_1d":
|
987 |
+
self.fusion_model = AFF(channels=64, type="1D")
|
988 |
+
elif self.fusion_type == "iaff_1d":
|
989 |
+
self.fusion_model = iAFF(channels=64, type="1D")
|
990 |
+
|
991 |
+
self.apply(self._init_weights)
|
992 |
+
|
993 |
+
def _init_weights(self, m):
|
994 |
+
if isinstance(m, nn.Linear):
|
995 |
+
trunc_normal_(m.weight, std=0.02)
|
996 |
+
if isinstance(m, nn.Linear) and m.bias is not None:
|
997 |
+
nn.init.constant_(m.bias, 0)
|
998 |
+
elif isinstance(m, nn.LayerNorm):
|
999 |
+
nn.init.constant_(m.bias, 0)
|
1000 |
+
nn.init.constant_(m.weight, 1.0)
|
1001 |
+
|
1002 |
+
@torch.jit.ignore
|
1003 |
+
def no_weight_decay(self):
|
1004 |
+
return {"absolute_pos_embed"}
|
1005 |
+
|
1006 |
+
@torch.jit.ignore
|
1007 |
+
def no_weight_decay_keywords(self):
|
1008 |
+
return {"relative_position_bias_table"}
|
1009 |
+
|
1010 |
+
def forward_features(self, x, longer_idx=None):
|
1011 |
+
# A deprecated optimization for using a hierarchical output from different blocks
|
1012 |
+
|
1013 |
+
frames_num = x.shape[2]
|
1014 |
+
x = self.patch_embed(x, longer_idx=longer_idx)
|
1015 |
+
if self.ape:
|
1016 |
+
x = x + self.absolute_pos_embed
|
1017 |
+
x = self.pos_drop(x)
|
1018 |
+
for i, layer in enumerate(self.layers):
|
1019 |
+
x, attn = layer(x)
|
1020 |
+
# for x
|
1021 |
+
x = self.norm(x)
|
1022 |
+
B, N, C = x.shape
|
1023 |
+
SF = frames_num // (2 ** (len(self.depths) - 1)) // self.patch_stride[0]
|
1024 |
+
ST = frames_num // (2 ** (len(self.depths) - 1)) // self.patch_stride[1]
|
1025 |
+
x = x.permute(0, 2, 1).contiguous().reshape(B, C, SF, ST)
|
1026 |
+
B, C, F, T = x.shape
|
1027 |
+
# group 2D CNN
|
1028 |
+
c_freq_bin = F // self.freq_ratio
|
1029 |
+
x = x.reshape(B, C, F // c_freq_bin, c_freq_bin, T)
|
1030 |
+
x = x.permute(0, 1, 3, 2, 4).contiguous().reshape(B, C, c_freq_bin, -1)
|
1031 |
+
# get latent_output
|
1032 |
+
fine_grained_latent_output = torch.mean(x, dim=2)
|
1033 |
+
fine_grained_latent_output = interpolate(
|
1034 |
+
fine_grained_latent_output.permute(0, 2, 1).contiguous(),
|
1035 |
+
8 * self.patch_stride[1],
|
1036 |
+
)
|
1037 |
+
|
1038 |
+
latent_output = self.avgpool(torch.flatten(x, 2))
|
1039 |
+
latent_output = torch.flatten(latent_output, 1)
|
1040 |
+
|
1041 |
+
# display the attention map, if needed
|
1042 |
+
|
1043 |
+
x = self.tscam_conv(x)
|
1044 |
+
x = torch.flatten(x, 2) # B, C, T
|
1045 |
+
|
1046 |
+
fpx = interpolate(
|
1047 |
+
torch.sigmoid(x).permute(0, 2, 1).contiguous(), 8 * self.patch_stride[1]
|
1048 |
+
)
|
1049 |
+
|
1050 |
+
x = self.avgpool(x)
|
1051 |
+
x = torch.flatten(x, 1)
|
1052 |
+
|
1053 |
+
output_dict = {
|
1054 |
+
"framewise_output": fpx, # already sigmoided
|
1055 |
+
"clipwise_output": torch.sigmoid(x),
|
1056 |
+
"fine_grained_embedding": fine_grained_latent_output,
|
1057 |
+
"embedding": latent_output,
|
1058 |
+
}
|
1059 |
+
|
1060 |
+
return output_dict
|
1061 |
+
|
1062 |
+
def crop_wav(self, x, crop_size, spe_pos=None):
|
1063 |
+
time_steps = x.shape[2]
|
1064 |
+
tx = torch.zeros(x.shape[0], x.shape[1], crop_size, x.shape[3]).to(x.device)
|
1065 |
+
for i in range(len(x)):
|
1066 |
+
if spe_pos is None:
|
1067 |
+
crop_pos = random.randint(0, time_steps - crop_size - 1)
|
1068 |
+
else:
|
1069 |
+
crop_pos = spe_pos
|
1070 |
+
tx[i][0] = x[i, 0, crop_pos : crop_pos + crop_size, :]
|
1071 |
+
return tx
|
1072 |
+
|
1073 |
+
# Reshape the wavform to a img size, if you want to use the pretrained swin transformer model
|
1074 |
+
def reshape_wav2img(self, x):
|
1075 |
+
B, C, T, F = x.shape
|
1076 |
+
target_T = int(self.spec_size * self.freq_ratio)
|
1077 |
+
target_F = self.spec_size // self.freq_ratio
|
1078 |
+
assert (
|
1079 |
+
T <= target_T and F <= target_F
|
1080 |
+
), "the wav size should less than or equal to the swin input size"
|
1081 |
+
# to avoid bicubic zero error
|
1082 |
+
if T < target_T:
|
1083 |
+
x = nn.functional.interpolate(
|
1084 |
+
x, (target_T, x.shape[3]), mode="bicubic", align_corners=True
|
1085 |
+
)
|
1086 |
+
if F < target_F:
|
1087 |
+
x = nn.functional.interpolate(
|
1088 |
+
x, (x.shape[2], target_F), mode="bicubic", align_corners=True
|
1089 |
+
)
|
1090 |
+
x = x.permute(0, 1, 3, 2).contiguous()
|
1091 |
+
x = x.reshape(
|
1092 |
+
x.shape[0],
|
1093 |
+
x.shape[1],
|
1094 |
+
x.shape[2],
|
1095 |
+
self.freq_ratio,
|
1096 |
+
x.shape[3] // self.freq_ratio,
|
1097 |
+
)
|
1098 |
+
# print(x.shape)
|
1099 |
+
x = x.permute(0, 1, 3, 2, 4).contiguous()
|
1100 |
+
x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3], x.shape[4])
|
1101 |
+
return x
|
1102 |
+
|
1103 |
+
# Repeat the wavform to a img size, if you want to use the pretrained swin transformer model
|
1104 |
+
def repeat_wat2img(self, x, cur_pos):
|
1105 |
+
B, C, T, F = x.shape
|
1106 |
+
target_T = int(self.spec_size * self.freq_ratio)
|
1107 |
+
target_F = self.spec_size // self.freq_ratio
|
1108 |
+
assert (
|
1109 |
+
T <= target_T and F <= target_F
|
1110 |
+
), "the wav size should less than or equal to the swin input size"
|
1111 |
+
# to avoid bicubic zero error
|
1112 |
+
if T < target_T:
|
1113 |
+
x = nn.functional.interpolate(
|
1114 |
+
x, (target_T, x.shape[3]), mode="bicubic", align_corners=True
|
1115 |
+
)
|
1116 |
+
if F < target_F:
|
1117 |
+
x = nn.functional.interpolate(
|
1118 |
+
x, (x.shape[2], target_F), mode="bicubic", align_corners=True
|
1119 |
+
)
|
1120 |
+
x = x.permute(0, 1, 3, 2).contiguous() # B C F T
|
1121 |
+
x = x[:, :, :, cur_pos : cur_pos + self.spec_size]
|
1122 |
+
x = x.repeat(repeats=(1, 1, 4, 1))
|
1123 |
+
return x
|
1124 |
+
|
1125 |
+
def forward(
|
1126 |
+
self, x: torch.Tensor, mixup_lambda=None, infer_mode=False, device=None
|
1127 |
+
): # out_feat_keys: List[str] = None):
|
1128 |
+
if self.enable_fusion and x["longer"].sum() == 0:
|
1129 |
+
# if no audio is longer than 10s, then randomly select one audio to be longer
|
1130 |
+
x["longer"][torch.randint(0, x["longer"].shape[0], (1,))] = True
|
1131 |
+
|
1132 |
+
if not self.enable_fusion:
|
1133 |
+
x = x["waveform"].to(device=device, non_blocking=True)
|
1134 |
+
x = self.spectrogram_extractor(x) # (batch_size, 1, time_steps, freq_bins)
|
1135 |
+
x = self.logmel_extractor(x) # (batch_size, 1, time_steps, mel_bins)
|
1136 |
+
x = x.transpose(1, 3)
|
1137 |
+
x = self.bn0(x)
|
1138 |
+
x = x.transpose(1, 3)
|
1139 |
+
if self.training:
|
1140 |
+
x = self.spec_augmenter(x)
|
1141 |
+
|
1142 |
+
if self.training and mixup_lambda is not None:
|
1143 |
+
x = do_mixup(x, mixup_lambda)
|
1144 |
+
|
1145 |
+
x = self.reshape_wav2img(x)
|
1146 |
+
output_dict = self.forward_features(x)
|
1147 |
+
else:
|
1148 |
+
longer_list = x["longer"].to(device=device, non_blocking=True)
|
1149 |
+
x = x["mel_fusion"].to(device=device, non_blocking=True)
|
1150 |
+
x = x.transpose(1, 3)
|
1151 |
+
x = self.bn0(x)
|
1152 |
+
x = x.transpose(1, 3)
|
1153 |
+
longer_list_idx = torch.where(longer_list)[0]
|
1154 |
+
if self.fusion_type in ["daf_1d", "aff_1d", "iaff_1d"]:
|
1155 |
+
new_x = x[:, 0:1, :, :].clone().contiguous()
|
1156 |
+
if len(longer_list_idx) > 0:
|
1157 |
+
# local processing
|
1158 |
+
fusion_x_local = x[longer_list_idx, 1:, :, :].clone().contiguous()
|
1159 |
+
FB, FC, FT, FF = fusion_x_local.size()
|
1160 |
+
fusion_x_local = fusion_x_local.view(FB * FC, FT, FF)
|
1161 |
+
fusion_x_local = torch.permute(
|
1162 |
+
fusion_x_local, (0, 2, 1)
|
1163 |
+
).contiguous()
|
1164 |
+
fusion_x_local = self.mel_conv1d(fusion_x_local)
|
1165 |
+
fusion_x_local = fusion_x_local.view(
|
1166 |
+
FB, FC, FF, fusion_x_local.size(-1)
|
1167 |
+
)
|
1168 |
+
fusion_x_local = (
|
1169 |
+
torch.permute(fusion_x_local, (0, 2, 1, 3))
|
1170 |
+
.contiguous()
|
1171 |
+
.flatten(2)
|
1172 |
+
)
|
1173 |
+
if fusion_x_local.size(-1) < FT:
|
1174 |
+
fusion_x_local = torch.cat(
|
1175 |
+
[
|
1176 |
+
fusion_x_local,
|
1177 |
+
torch.zeros(
|
1178 |
+
(FB, FF, FT - fusion_x_local.size(-1)),
|
1179 |
+
device=device,
|
1180 |
+
),
|
1181 |
+
],
|
1182 |
+
dim=-1,
|
1183 |
+
)
|
1184 |
+
else:
|
1185 |
+
fusion_x_local = fusion_x_local[:, :, :FT]
|
1186 |
+
# 1D fusion
|
1187 |
+
new_x = new_x.squeeze(1).permute((0, 2, 1)).contiguous()
|
1188 |
+
new_x[longer_list_idx] = self.fusion_model(
|
1189 |
+
new_x[longer_list_idx], fusion_x_local
|
1190 |
+
)
|
1191 |
+
x = new_x.permute((0, 2, 1)).contiguous()[:, None, :, :]
|
1192 |
+
else:
|
1193 |
+
x = new_x
|
1194 |
+
|
1195 |
+
elif self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d", "channel_map"]:
|
1196 |
+
x = x # no change
|
1197 |
+
|
1198 |
+
if self.training:
|
1199 |
+
x = self.spec_augmenter(x)
|
1200 |
+
if self.training and mixup_lambda is not None:
|
1201 |
+
x = do_mixup(x, mixup_lambda)
|
1202 |
+
|
1203 |
+
x = self.reshape_wav2img(x)
|
1204 |
+
output_dict = self.forward_features(x, longer_idx=longer_list_idx)
|
1205 |
+
|
1206 |
+
# if infer_mode:
|
1207 |
+
# # in infer mode. we need to handle different length audio input
|
1208 |
+
# frame_num = x.shape[2]
|
1209 |
+
# target_T = int(self.spec_size * self.freq_ratio)
|
1210 |
+
# repeat_ratio = math.floor(target_T / frame_num)
|
1211 |
+
# x = x.repeat(repeats=(1,1,repeat_ratio,1))
|
1212 |
+
# x = self.reshape_wav2img(x)
|
1213 |
+
# output_dict = self.forward_features(x)
|
1214 |
+
# else:
|
1215 |
+
# if x.shape[2] > self.freq_ratio * self.spec_size:
|
1216 |
+
# if self.training:
|
1217 |
+
# x = self.crop_wav(x, crop_size=self.freq_ratio * self.spec_size)
|
1218 |
+
# x = self.reshape_wav2img(x)
|
1219 |
+
# output_dict = self.forward_features(x)
|
1220 |
+
# else:
|
1221 |
+
# # Change: Hard code here
|
1222 |
+
# overlap_size = (x.shape[2] - 1) // 4
|
1223 |
+
# output_dicts = []
|
1224 |
+
# crop_size = (x.shape[2] - 1) // 2
|
1225 |
+
# for cur_pos in range(0, x.shape[2] - crop_size - 1, overlap_size):
|
1226 |
+
# tx = self.crop_wav(x, crop_size = crop_size, spe_pos = cur_pos)
|
1227 |
+
# tx = self.reshape_wav2img(tx)
|
1228 |
+
# output_dicts.append(self.forward_features(tx))
|
1229 |
+
# clipwise_output = torch.zeros_like(output_dicts[0]["clipwise_output"]).float().to(x.device)
|
1230 |
+
# framewise_output = torch.zeros_like(output_dicts[0]["framewise_output"]).float().to(x.device)
|
1231 |
+
# for d in output_dicts:
|
1232 |
+
# clipwise_output += d["clipwise_output"]
|
1233 |
+
# framewise_output += d["framewise_output"]
|
1234 |
+
# clipwise_output = clipwise_output / len(output_dicts)
|
1235 |
+
# framewise_output = framewise_output / len(output_dicts)
|
1236 |
+
# output_dict = {
|
1237 |
+
# 'framewise_output': framewise_output,
|
1238 |
+
# 'clipwise_output': clipwise_output
|
1239 |
+
# }
|
1240 |
+
# else: # this part is typically used, and most easy one
|
1241 |
+
# x = self.reshape_wav2img(x)
|
1242 |
+
# output_dict = self.forward_features(x)
|
1243 |
+
# x = self.head(x)
|
1244 |
+
|
1245 |
+
# We process the data in the dataloader part, in that here we only consider the input_T < fixed_T
|
1246 |
+
|
1247 |
+
return output_dict
|
1248 |
+
|
1249 |
+
|
1250 |
+
def create_htsat_model(audio_cfg, enable_fusion=False, fusion_type="None"):
|
1251 |
+
try:
|
1252 |
+
assert audio_cfg.model_name in [
|
1253 |
+
"tiny",
|
1254 |
+
"base",
|
1255 |
+
"large",
|
1256 |
+
], "model name for HTS-AT is wrong!"
|
1257 |
+
if audio_cfg.model_name == "tiny":
|
1258 |
+
model = HTSAT_Swin_Transformer(
|
1259 |
+
spec_size=256,
|
1260 |
+
patch_size=4,
|
1261 |
+
patch_stride=(4, 4),
|
1262 |
+
num_classes=audio_cfg.class_num,
|
1263 |
+
embed_dim=96,
|
1264 |
+
depths=[2, 2, 6, 2],
|
1265 |
+
num_heads=[4, 8, 16, 32],
|
1266 |
+
window_size=8,
|
1267 |
+
config=audio_cfg,
|
1268 |
+
enable_fusion=enable_fusion,
|
1269 |
+
fusion_type=fusion_type,
|
1270 |
+
)
|
1271 |
+
elif audio_cfg.model_name == "base":
|
1272 |
+
model = HTSAT_Swin_Transformer(
|
1273 |
+
spec_size=256,
|
1274 |
+
patch_size=4,
|
1275 |
+
patch_stride=(4, 4),
|
1276 |
+
num_classes=audio_cfg.class_num,
|
1277 |
+
embed_dim=128,
|
1278 |
+
depths=[2, 2, 12, 2],
|
1279 |
+
num_heads=[4, 8, 16, 32],
|
1280 |
+
window_size=8,
|
1281 |
+
config=audio_cfg,
|
1282 |
+
enable_fusion=enable_fusion,
|
1283 |
+
fusion_type=fusion_type,
|
1284 |
+
)
|
1285 |
+
elif audio_cfg.model_name == "large":
|
1286 |
+
model = HTSAT_Swin_Transformer(
|
1287 |
+
spec_size=256,
|
1288 |
+
patch_size=4,
|
1289 |
+
patch_stride=(4, 4),
|
1290 |
+
num_classes=audio_cfg.class_num,
|
1291 |
+
embed_dim=256,
|
1292 |
+
depths=[2, 2, 12, 2],
|
1293 |
+
num_heads=[4, 8, 16, 32],
|
1294 |
+
window_size=8,
|
1295 |
+
config=audio_cfg,
|
1296 |
+
enable_fusion=enable_fusion,
|
1297 |
+
fusion_type=fusion_type,
|
1298 |
+
)
|
1299 |
+
|
1300 |
+
return model
|
1301 |
+
except:
|
1302 |
+
raise RuntimeError(
|
1303 |
+
f"Import Model for {audio_cfg.model_name} not found, or the audio cfg parameters are not enough."
|
1304 |
+
)
|
audioldm2/clap/open_clip/loss.py
ADDED
@@ -0,0 +1,397 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.distributed.nn
|
3 |
+
from torch import distributed as dist, nn as nn
|
4 |
+
from torch.nn import functional as F
|
5 |
+
import numpy as np
|
6 |
+
from sklearn.metrics import average_precision_score, roc_auc_score, accuracy_score
|
7 |
+
|
8 |
+
try:
|
9 |
+
import horovod.torch as hvd
|
10 |
+
except ImportError:
|
11 |
+
hvd = None
|
12 |
+
|
13 |
+
|
14 |
+
def gather_features(
|
15 |
+
audio_features,
|
16 |
+
text_features,
|
17 |
+
audio_features_mlp=None,
|
18 |
+
text_features_mlp=None,
|
19 |
+
local_loss=False,
|
20 |
+
gather_with_grad=False,
|
21 |
+
rank=0,
|
22 |
+
world_size=1,
|
23 |
+
use_horovod=False,
|
24 |
+
mlp_loss=False,
|
25 |
+
):
|
26 |
+
if use_horovod:
|
27 |
+
assert hvd is not None, "Please install horovod"
|
28 |
+
if gather_with_grad:
|
29 |
+
all_audio_features = hvd.allgather(audio_features)
|
30 |
+
all_text_features = hvd.allgather(text_features)
|
31 |
+
if mlp_loss:
|
32 |
+
all_audio_features_mlp = hvd.allgather(audio_features_mlp)
|
33 |
+
all_text_features_mlp = hvd.allgather(text_features_mlp)
|
34 |
+
else:
|
35 |
+
with torch.no_grad():
|
36 |
+
all_audio_features = hvd.allgather(audio_features)
|
37 |
+
all_text_features = hvd.allgather(text_features)
|
38 |
+
if mlp_loss:
|
39 |
+
all_audio_features_mlp = hvd.allgather(audio_features_mlp)
|
40 |
+
all_text_features_mlp = hvd.allgather(text_features_mlp)
|
41 |
+
if not local_loss:
|
42 |
+
# ensure grads for local rank when all_* features don't have a gradient
|
43 |
+
gathered_audio_features = list(
|
44 |
+
all_audio_features.chunk(world_size, dim=0)
|
45 |
+
)
|
46 |
+
gathered_text_features = list(
|
47 |
+
all_text_features.chunk(world_size, dim=0)
|
48 |
+
)
|
49 |
+
gathered_audio_features[rank] = audio_features
|
50 |
+
gathered_text_features[rank] = text_features
|
51 |
+
all_audio_features = torch.cat(gathered_audio_features, dim=0)
|
52 |
+
all_text_features = torch.cat(gathered_text_features, dim=0)
|
53 |
+
if mlp_loss:
|
54 |
+
gathered_audio_features_mlp = list(
|
55 |
+
all_audio_features_mlp.chunk(world_size, dim=0)
|
56 |
+
)
|
57 |
+
gathered_text_features_mlp = list(
|
58 |
+
all_text_features_mlp.chunk(world_size, dim=0)
|
59 |
+
)
|
60 |
+
gathered_audio_features_mlp[rank] = audio_features_mlp
|
61 |
+
gathered_text_features_mlp[rank] = text_features_mlp
|
62 |
+
all_audio_features_mlp = torch.cat(
|
63 |
+
gathered_audio_features_mlp, dim=0
|
64 |
+
)
|
65 |
+
all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
|
66 |
+
else:
|
67 |
+
# We gather tensors from all gpus
|
68 |
+
if gather_with_grad:
|
69 |
+
all_audio_features = torch.cat(
|
70 |
+
torch.distributed.nn.all_gather(audio_features), dim=0
|
71 |
+
)
|
72 |
+
all_text_features = torch.cat(
|
73 |
+
torch.distributed.nn.all_gather(text_features), dim=0
|
74 |
+
)
|
75 |
+
if mlp_loss:
|
76 |
+
all_audio_features_mlp = torch.cat(
|
77 |
+
torch.distributed.nn.all_gather(audio_features_mlp), dim=0
|
78 |
+
)
|
79 |
+
all_text_features_mlp = torch.cat(
|
80 |
+
torch.distributed.nn.all_gather(text_features_mlp), dim=0
|
81 |
+
)
|
82 |
+
else:
|
83 |
+
gathered_audio_features = [
|
84 |
+
torch.zeros_like(audio_features) for _ in range(world_size)
|
85 |
+
]
|
86 |
+
gathered_text_features = [
|
87 |
+
torch.zeros_like(text_features) for _ in range(world_size)
|
88 |
+
]
|
89 |
+
dist.all_gather(gathered_audio_features, audio_features)
|
90 |
+
dist.all_gather(gathered_text_features, text_features)
|
91 |
+
if mlp_loss:
|
92 |
+
gathered_audio_features_mlp = [
|
93 |
+
torch.zeros_like(audio_features_mlp) for _ in range(world_size)
|
94 |
+
]
|
95 |
+
gathered_text_features_mlp = [
|
96 |
+
torch.zeros_like(text_features_mlp) for _ in range(world_size)
|
97 |
+
]
|
98 |
+
dist.all_gather(gathered_audio_features_mlp, audio_features_mlp)
|
99 |
+
dist.all_gather(gathered_text_features_mlp, text_features_mlp)
|
100 |
+
if not local_loss:
|
101 |
+
# ensure grads for local rank when all_* features don't have a gradient
|
102 |
+
gathered_audio_features[rank] = audio_features
|
103 |
+
gathered_text_features[rank] = text_features
|
104 |
+
if mlp_loss:
|
105 |
+
gathered_audio_features_mlp[rank] = audio_features_mlp
|
106 |
+
gathered_text_features_mlp[rank] = text_features_mlp
|
107 |
+
|
108 |
+
all_audio_features = torch.cat(gathered_audio_features, dim=0)
|
109 |
+
all_text_features = torch.cat(gathered_text_features, dim=0)
|
110 |
+
if mlp_loss:
|
111 |
+
all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0)
|
112 |
+
all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
|
113 |
+
if mlp_loss:
|
114 |
+
return (
|
115 |
+
all_audio_features,
|
116 |
+
all_text_features,
|
117 |
+
all_audio_features_mlp,
|
118 |
+
all_text_features_mlp,
|
119 |
+
)
|
120 |
+
else:
|
121 |
+
return all_audio_features, all_text_features
|
122 |
+
|
123 |
+
|
124 |
+
class ClipLoss(nn.Module):
|
125 |
+
def __init__(
|
126 |
+
self,
|
127 |
+
local_loss=False,
|
128 |
+
gather_with_grad=False,
|
129 |
+
cache_labels=False,
|
130 |
+
rank=0,
|
131 |
+
world_size=1,
|
132 |
+
use_horovod=False,
|
133 |
+
mlp_loss=False,
|
134 |
+
weight_loss_kappa=0,
|
135 |
+
):
|
136 |
+
super().__init__()
|
137 |
+
self.local_loss = local_loss
|
138 |
+
self.gather_with_grad = gather_with_grad
|
139 |
+
self.cache_labels = cache_labels
|
140 |
+
self.rank = rank
|
141 |
+
self.world_size = world_size
|
142 |
+
self.use_horovod = use_horovod
|
143 |
+
self.mlp_loss = mlp_loss
|
144 |
+
self.weighted_loss = bool(weight_loss_kappa != 0)
|
145 |
+
self.weight_loss_kappa = weight_loss_kappa
|
146 |
+
# cache state
|
147 |
+
self.prev_num_logits = 0
|
148 |
+
self.labels = {}
|
149 |
+
|
150 |
+
def forward(
|
151 |
+
self,
|
152 |
+
audio_features,
|
153 |
+
text_features,
|
154 |
+
logit_scale_a,
|
155 |
+
logit_scale_t=None,
|
156 |
+
audio_features_mlp=None,
|
157 |
+
text_features_mlp=None,
|
158 |
+
):
|
159 |
+
device = audio_features.device
|
160 |
+
if self.mlp_loss:
|
161 |
+
if self.world_size > 1:
|
162 |
+
(
|
163 |
+
all_audio_features,
|
164 |
+
all_text_features,
|
165 |
+
all_audio_features_mlp,
|
166 |
+
all_text_features_mlp,
|
167 |
+
) = gather_features(
|
168 |
+
audio_features=audio_features,
|
169 |
+
text_features=text_features,
|
170 |
+
audio_features_mlp=audio_features_mlp,
|
171 |
+
text_features_mlp=text_features_mlp,
|
172 |
+
local_loss=self.local_loss,
|
173 |
+
gather_with_grad=self.gather_with_grad,
|
174 |
+
rank=self.rank,
|
175 |
+
world_size=self.world_size,
|
176 |
+
use_horovod=self.use_horovod,
|
177 |
+
mlp_loss=self.mlp_loss,
|
178 |
+
)
|
179 |
+
if self.local_loss:
|
180 |
+
a_logits_per_audio = (
|
181 |
+
logit_scale_a * audio_features @ all_text_features_mlp.T
|
182 |
+
)
|
183 |
+
a_logits_per_text = (
|
184 |
+
logit_scale_a * text_features_mlp @ all_audio_features.T
|
185 |
+
)
|
186 |
+
t_logits_per_audio = (
|
187 |
+
logit_scale_t * audio_features_mlp @ all_text_features.T
|
188 |
+
)
|
189 |
+
t_logits_per_text = (
|
190 |
+
logit_scale_t * text_features @ all_audio_features_mlp.T
|
191 |
+
)
|
192 |
+
else:
|
193 |
+
a_logits_per_audio = (
|
194 |
+
logit_scale_a * all_audio_features @ all_text_features_mlp.T
|
195 |
+
)
|
196 |
+
a_logits_per_text = a_logits_per_audio.T
|
197 |
+
t_logits_per_audio = (
|
198 |
+
logit_scale_t * all_audio_features_mlp @ all_text_features.T
|
199 |
+
)
|
200 |
+
t_logits_per_text = t_logits_per_audio.T
|
201 |
+
else:
|
202 |
+
a_logits_per_audio = (
|
203 |
+
logit_scale_a * audio_features @ text_features_mlp.T
|
204 |
+
)
|
205 |
+
a_logits_per_text = logit_scale_a * text_features_mlp @ audio_features.T
|
206 |
+
t_logits_per_audio = (
|
207 |
+
logit_scale_t * audio_features_mlp @ text_features.T
|
208 |
+
)
|
209 |
+
t_logits_per_text = logit_scale_t * text_features @ audio_features_mlp.T
|
210 |
+
|
211 |
+
# calculated ground-truth and cache if enabled
|
212 |
+
num_logits = a_logits_per_audio.shape[0]
|
213 |
+
if self.prev_num_logits != num_logits or device not in self.labels:
|
214 |
+
labels = torch.arange(num_logits, device=device, dtype=torch.long)
|
215 |
+
if self.world_size > 1 and self.local_loss:
|
216 |
+
labels = labels + num_logits * self.rank
|
217 |
+
if self.cache_labels:
|
218 |
+
self.labels[device] = labels
|
219 |
+
self.prev_num_logits = num_logits
|
220 |
+
else:
|
221 |
+
labels = self.labels[device]
|
222 |
+
|
223 |
+
if not self.weighted_loss:
|
224 |
+
total_loss = (
|
225 |
+
F.cross_entropy(a_logits_per_audio, labels)
|
226 |
+
+ F.cross_entropy(a_logits_per_text, labels)
|
227 |
+
+ F.cross_entropy(t_logits_per_audio, labels)
|
228 |
+
+ F.cross_entropy(t_logits_per_text, labels)
|
229 |
+
) / 4
|
230 |
+
else:
|
231 |
+
audio_weight = (audio_features @ audio_features.T).detach()
|
232 |
+
audio_weight = (
|
233 |
+
torch.exp(
|
234 |
+
torch.sum(audio_weight, axis=1)
|
235 |
+
/ (self.weight_loss_kappa * len(audio_weight))
|
236 |
+
)
|
237 |
+
).detach()
|
238 |
+
text_weight = (text_features @ text_features.T).detach()
|
239 |
+
text_weight = (
|
240 |
+
torch.exp(
|
241 |
+
torch.sum(text_weight, axis=1)
|
242 |
+
/ (self.weight_loss_kappa * len(text_features))
|
243 |
+
)
|
244 |
+
).detach()
|
245 |
+
total_loss = (
|
246 |
+
F.cross_entropy(a_logits_per_audio, labels, weight=audio_weight)
|
247 |
+
+ F.cross_entropy(a_logits_per_text, labels, weight=audio_weight)
|
248 |
+
+ F.cross_entropy(t_logits_per_audio, labels, weight=text_weight)
|
249 |
+
+ F.cross_entropy(t_logits_per_text, labels, weight=text_weight)
|
250 |
+
) / 4
|
251 |
+
else:
|
252 |
+
if self.world_size > 1:
|
253 |
+
all_audio_features, all_text_features = gather_features(
|
254 |
+
audio_features=audio_features,
|
255 |
+
text_features=text_features,
|
256 |
+
local_loss=self.local_loss,
|
257 |
+
gather_with_grad=self.gather_with_grad,
|
258 |
+
rank=self.rank,
|
259 |
+
world_size=self.world_size,
|
260 |
+
use_horovod=self.use_horovod,
|
261 |
+
mlp_loss=self.mlp_loss,
|
262 |
+
)
|
263 |
+
|
264 |
+
if self.local_loss:
|
265 |
+
logits_per_audio = (
|
266 |
+
logit_scale_a * audio_features @ all_text_features.T
|
267 |
+
)
|
268 |
+
logits_per_text = (
|
269 |
+
logit_scale_a * text_features @ all_audio_features.T
|
270 |
+
)
|
271 |
+
else:
|
272 |
+
logits_per_audio = (
|
273 |
+
logit_scale_a * all_audio_features @ all_text_features.T
|
274 |
+
)
|
275 |
+
logits_per_text = logits_per_audio.T
|
276 |
+
else:
|
277 |
+
logits_per_audio = logit_scale_a * audio_features @ text_features.T
|
278 |
+
logits_per_text = logit_scale_a * text_features @ audio_features.T
|
279 |
+
|
280 |
+
# calculated ground-truth and cache if enabled
|
281 |
+
num_logits = logits_per_audio.shape[0]
|
282 |
+
if self.prev_num_logits != num_logits or device not in self.labels:
|
283 |
+
labels = torch.arange(num_logits, device=device, dtype=torch.long)
|
284 |
+
if self.world_size > 1 and self.local_loss:
|
285 |
+
labels = labels + num_logits * self.rank
|
286 |
+
if self.cache_labels:
|
287 |
+
self.labels[device] = labels
|
288 |
+
self.prev_num_logits = num_logits
|
289 |
+
else:
|
290 |
+
labels = self.labels[device]
|
291 |
+
if not self.weighted_loss:
|
292 |
+
total_loss = (
|
293 |
+
F.cross_entropy(logits_per_audio, labels)
|
294 |
+
+ F.cross_entropy(logits_per_text, labels)
|
295 |
+
) / 2
|
296 |
+
else:
|
297 |
+
audio_weight = (all_audio_features @ all_audio_features.T).detach()
|
298 |
+
audio_weight = (
|
299 |
+
torch.exp(
|
300 |
+
torch.sum(audio_weight, axis=1)
|
301 |
+
/ (self.weight_loss_kappa * len(all_audio_features))
|
302 |
+
)
|
303 |
+
).detach()
|
304 |
+
text_weight = (all_text_features @ all_text_features.T).detach()
|
305 |
+
text_weight = (
|
306 |
+
torch.exp(
|
307 |
+
torch.sum(text_weight, axis=1)
|
308 |
+
/ (self.weight_loss_kappa * len(all_text_features))
|
309 |
+
)
|
310 |
+
).detach()
|
311 |
+
total_loss = (
|
312 |
+
F.cross_entropy(logits_per_audio, labels, weight=text_weight)
|
313 |
+
+ F.cross_entropy(logits_per_text, labels, weight=audio_weight)
|
314 |
+
) / 2
|
315 |
+
return total_loss
|
316 |
+
|
317 |
+
|
318 |
+
def lp_gather_features(pred, target, world_size=1, use_horovod=False):
|
319 |
+
if use_horovod:
|
320 |
+
assert hvd is not None, "Please install horovod"
|
321 |
+
with torch.no_grad():
|
322 |
+
all_preds = hvd.allgather(pred)
|
323 |
+
all_targets = hvd.allgath(target)
|
324 |
+
else:
|
325 |
+
gathered_preds = [torch.zeros_like(pred) for _ in range(world_size)]
|
326 |
+
gathered_targets = [torch.zeros_like(target) for _ in range(world_size)]
|
327 |
+
|
328 |
+
dist.all_gather(gathered_preds, pred)
|
329 |
+
dist.all_gather(gathered_targets, target)
|
330 |
+
all_preds = torch.cat(gathered_preds, dim=0)
|
331 |
+
all_targets = torch.cat(gathered_targets, dim=0)
|
332 |
+
|
333 |
+
return all_preds, all_targets
|
334 |
+
|
335 |
+
|
336 |
+
def get_map(pred, target):
|
337 |
+
pred = torch.sigmoid(pred).numpy()
|
338 |
+
target = target.numpy()
|
339 |
+
return np.mean(average_precision_score(target, pred, average=None))
|
340 |
+
|
341 |
+
|
342 |
+
def get_acc(pred, target):
|
343 |
+
pred = torch.argmax(pred, 1).numpy()
|
344 |
+
target = torch.argmax(target, 1).numpy()
|
345 |
+
return accuracy_score(target, pred)
|
346 |
+
|
347 |
+
|
348 |
+
def get_mauc(pred, target):
|
349 |
+
pred = torch.sigmoid(pred).numpy()
|
350 |
+
target = target.numpy()
|
351 |
+
return np.mean(roc_auc_score(target, pred, average=None))
|
352 |
+
|
353 |
+
|
354 |
+
class LPMetrics(object):
|
355 |
+
def __init__(self, metric_names=["map", "acc", "mauc"]):
|
356 |
+
self.metrics = []
|
357 |
+
for name in metric_names:
|
358 |
+
self.metrics.append(self.get_metric(name))
|
359 |
+
self.metric_names = metric_names
|
360 |
+
|
361 |
+
def get_metric(self, name):
|
362 |
+
if name == "map":
|
363 |
+
return get_map
|
364 |
+
elif name == "acc":
|
365 |
+
return get_acc
|
366 |
+
elif name == "mauc":
|
367 |
+
return get_mauc
|
368 |
+
else:
|
369 |
+
raise ValueError(f"the metric should be at least one of [map, acc, mauc]")
|
370 |
+
|
371 |
+
def evaluate_mertics(self, pred, target):
|
372 |
+
metric_dict = {}
|
373 |
+
for i in range(len(self.metric_names)):
|
374 |
+
metric_dict[self.metric_names[i]] = self.metrics[i](pred, target)
|
375 |
+
return metric_dict
|
376 |
+
|
377 |
+
|
378 |
+
def calc_celoss(pred, target):
|
379 |
+
target = torch.argmax(target, 1).long()
|
380 |
+
return nn.CrossEntropyLoss()(pred, target)
|
381 |
+
|
382 |
+
|
383 |
+
class LPLoss(nn.Module):
|
384 |
+
def __init__(self, loss_name):
|
385 |
+
super().__init__()
|
386 |
+
if loss_name == "bce":
|
387 |
+
self.loss_func = nn.BCEWithLogitsLoss()
|
388 |
+
elif loss_name == "ce":
|
389 |
+
self.loss_func = calc_celoss
|
390 |
+
elif loss_name == "mse":
|
391 |
+
self.loss_func = nn.MSELoss()
|
392 |
+
else:
|
393 |
+
raise ValueError(f"the loss func should be at least one of [bce, ce, mse]")
|
394 |
+
|
395 |
+
def forward(self, pred, target):
|
396 |
+
loss = self.loss_func(pred, target)
|
397 |
+
return loss
|
audioldm2/clap/open_clip/model.py
ADDED
@@ -0,0 +1,931 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
""" CLAP Model
|
2 |
+
|
3 |
+
Adapted from CLIP: https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
|
4 |
+
Adapted to the Audio Task.
|
5 |
+
"""
|
6 |
+
|
7 |
+
from collections import OrderedDict
|
8 |
+
from dataclasses import dataclass
|
9 |
+
from typing import Tuple, Union, Callable, Optional
|
10 |
+
|
11 |
+
import numpy as np
|
12 |
+
import torch
|
13 |
+
import torch.nn.functional as F
|
14 |
+
from torch import nn
|
15 |
+
|
16 |
+
import logging
|
17 |
+
from .utils import freeze_batch_norm_2d
|
18 |
+
|
19 |
+
from .pann_model import create_pann_model
|
20 |
+
from .htsat import create_htsat_model
|
21 |
+
from transformers import BertModel, RobertaModel, BartModel, RobertaConfig
|
22 |
+
|
23 |
+
|
24 |
+
class MLPLayers(nn.Module):
|
25 |
+
def __init__(self, units=[512, 512, 512], nonlin=nn.ReLU(), dropout=0.1):
|
26 |
+
super(MLPLayers, self).__init__()
|
27 |
+
self.nonlin = nonlin
|
28 |
+
self.dropout = dropout
|
29 |
+
|
30 |
+
sequence = []
|
31 |
+
for u0, u1 in zip(units[:-1], units[1:]):
|
32 |
+
sequence.append(nn.Linear(u0, u1))
|
33 |
+
sequence.append(self.nonlin)
|
34 |
+
sequence.append(nn.Dropout(self.dropout))
|
35 |
+
sequence = sequence[:-2]
|
36 |
+
|
37 |
+
self.sequential = nn.Sequential(*sequence)
|
38 |
+
|
39 |
+
def forward(self, X):
|
40 |
+
X = self.sequential(X)
|
41 |
+
return X
|
42 |
+
|
43 |
+
|
44 |
+
class Bottleneck(nn.Module):
|
45 |
+
expansion = 4
|
46 |
+
|
47 |
+
def __init__(self, inplanes, planes, stride=1):
|
48 |
+
super().__init__()
|
49 |
+
|
50 |
+
# all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
|
51 |
+
self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
|
52 |
+
self.bn1 = nn.BatchNorm2d(planes)
|
53 |
+
|
54 |
+
self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
|
55 |
+
self.bn2 = nn.BatchNorm2d(planes)
|
56 |
+
|
57 |
+
self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
|
58 |
+
|
59 |
+
self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
|
60 |
+
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
61 |
+
|
62 |
+
self.relu = nn.ReLU(inplace=True)
|
63 |
+
self.downsample = None
|
64 |
+
self.stride = stride
|
65 |
+
|
66 |
+
if stride > 1 or inplanes != planes * Bottleneck.expansion:
|
67 |
+
# downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
|
68 |
+
self.downsample = nn.Sequential(
|
69 |
+
OrderedDict(
|
70 |
+
[
|
71 |
+
("-1", nn.AvgPool2d(stride)),
|
72 |
+
(
|
73 |
+
"0",
|
74 |
+
nn.Conv2d(
|
75 |
+
inplanes,
|
76 |
+
planes * self.expansion,
|
77 |
+
1,
|
78 |
+
stride=1,
|
79 |
+
bias=False,
|
80 |
+
),
|
81 |
+
),
|
82 |
+
("1", nn.BatchNorm2d(planes * self.expansion)),
|
83 |
+
]
|
84 |
+
)
|
85 |
+
)
|
86 |
+
|
87 |
+
def forward(self, x: torch.Tensor):
|
88 |
+
identity = x
|
89 |
+
|
90 |
+
out = self.relu(self.bn1(self.conv1(x)))
|
91 |
+
out = self.relu(self.bn2(self.conv2(out)))
|
92 |
+
out = self.avgpool(out)
|
93 |
+
out = self.bn3(self.conv3(out))
|
94 |
+
|
95 |
+
if self.downsample is not None:
|
96 |
+
identity = self.downsample(x)
|
97 |
+
|
98 |
+
out += identity
|
99 |
+
out = self.relu(out)
|
100 |
+
return out
|
101 |
+
|
102 |
+
|
103 |
+
class AttentionPool2d(nn.Module):
|
104 |
+
def __init__(
|
105 |
+
self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None
|
106 |
+
):
|
107 |
+
super().__init__()
|
108 |
+
self.positional_embedding = nn.Parameter(
|
109 |
+
torch.randn(spacial_dim**2 + 1, embed_dim) / embed_dim**0.5
|
110 |
+
)
|
111 |
+
self.k_proj = nn.Linear(embed_dim, embed_dim)
|
112 |
+
self.q_proj = nn.Linear(embed_dim, embed_dim)
|
113 |
+
self.v_proj = nn.Linear(embed_dim, embed_dim)
|
114 |
+
self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
|
115 |
+
self.num_heads = num_heads
|
116 |
+
|
117 |
+
def forward(self, x):
|
118 |
+
x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(
|
119 |
+
2, 0, 1
|
120 |
+
) # NCHW -> (HW)NC
|
121 |
+
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC
|
122 |
+
x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC
|
123 |
+
x, _ = F.multi_head_attention_forward(
|
124 |
+
query=x,
|
125 |
+
key=x,
|
126 |
+
value=x,
|
127 |
+
embed_dim_to_check=x.shape[-1],
|
128 |
+
num_heads=self.num_heads,
|
129 |
+
q_proj_weight=self.q_proj.weight,
|
130 |
+
k_proj_weight=self.k_proj.weight,
|
131 |
+
v_proj_weight=self.v_proj.weight,
|
132 |
+
in_proj_weight=None,
|
133 |
+
in_proj_bias=torch.cat(
|
134 |
+
[self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]
|
135 |
+
),
|
136 |
+
bias_k=None,
|
137 |
+
bias_v=None,
|
138 |
+
add_zero_attn=False,
|
139 |
+
dropout_p=0,
|
140 |
+
out_proj_weight=self.c_proj.weight,
|
141 |
+
out_proj_bias=self.c_proj.bias,
|
142 |
+
use_separate_proj_weight=True,
|
143 |
+
training=self.training,
|
144 |
+
need_weights=False,
|
145 |
+
)
|
146 |
+
|
147 |
+
return x[0]
|
148 |
+
|
149 |
+
|
150 |
+
class ModifiedResNet(nn.Module):
|
151 |
+
"""
|
152 |
+
A ResNet class that is similar to torchvision's but contains the following changes:
|
153 |
+
- There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
|
154 |
+
- Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
|
155 |
+
- The final pooling layer is a QKV attention instead of an average pool
|
156 |
+
"""
|
157 |
+
|
158 |
+
def __init__(self, layers, output_dim, heads, image_size=224, width=64):
|
159 |
+
super().__init__()
|
160 |
+
self.output_dim = output_dim
|
161 |
+
self.image_size = image_size
|
162 |
+
|
163 |
+
# the 3-layer stem
|
164 |
+
self.conv1 = nn.Conv2d(
|
165 |
+
3, width // 2, kernel_size=3, stride=2, padding=1, bias=False
|
166 |
+
)
|
167 |
+
self.bn1 = nn.BatchNorm2d(width // 2)
|
168 |
+
self.conv2 = nn.Conv2d(
|
169 |
+
width // 2, width // 2, kernel_size=3, padding=1, bias=False
|
170 |
+
)
|
171 |
+
self.bn2 = nn.BatchNorm2d(width // 2)
|
172 |
+
self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
|
173 |
+
self.bn3 = nn.BatchNorm2d(width)
|
174 |
+
self.avgpool = nn.AvgPool2d(2)
|
175 |
+
self.relu = nn.ReLU(inplace=True)
|
176 |
+
|
177 |
+
# residual layers
|
178 |
+
self._inplanes = width # this is a *mutable* variable used during construction
|
179 |
+
self.layer1 = self._make_layer(width, layers[0])
|
180 |
+
self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
|
181 |
+
self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
|
182 |
+
self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
|
183 |
+
|
184 |
+
embed_dim = width * 32 # the ResNet feature dimension
|
185 |
+
self.attnpool = AttentionPool2d(image_size // 32, embed_dim, heads, output_dim)
|
186 |
+
|
187 |
+
self.init_parameters()
|
188 |
+
|
189 |
+
def _make_layer(self, planes, blocks, stride=1):
|
190 |
+
layers = [Bottleneck(self._inplanes, planes, stride)]
|
191 |
+
|
192 |
+
self._inplanes = planes * Bottleneck.expansion
|
193 |
+
for _ in range(1, blocks):
|
194 |
+
layers.append(Bottleneck(self._inplanes, planes))
|
195 |
+
|
196 |
+
return nn.Sequential(*layers)
|
197 |
+
|
198 |
+
def init_parameters(self):
|
199 |
+
if self.attnpool is not None:
|
200 |
+
std = self.attnpool.c_proj.in_features**-0.5
|
201 |
+
nn.init.normal_(self.attnpool.q_proj.weight, std=std)
|
202 |
+
nn.init.normal_(self.attnpool.k_proj.weight, std=std)
|
203 |
+
nn.init.normal_(self.attnpool.v_proj.weight, std=std)
|
204 |
+
nn.init.normal_(self.attnpool.c_proj.weight, std=std)
|
205 |
+
|
206 |
+
for resnet_block in [self.layer1, self.layer2, self.layer3, self.layer4]:
|
207 |
+
for name, param in resnet_block.named_parameters():
|
208 |
+
if name.endswith("bn3.weight"):
|
209 |
+
nn.init.zeros_(param)
|
210 |
+
|
211 |
+
def lock(self, unlocked_groups=0, freeze_bn_stats=False):
|
212 |
+
assert (
|
213 |
+
unlocked_groups == 0
|
214 |
+
), "partial locking not currently supported for this model"
|
215 |
+
for param in self.parameters():
|
216 |
+
param.requires_grad = False
|
217 |
+
if freeze_bn_stats:
|
218 |
+
freeze_batch_norm_2d(self)
|
219 |
+
|
220 |
+
def stem(self, x):
|
221 |
+
for conv, bn in [
|
222 |
+
(self.conv1, self.bn1),
|
223 |
+
(self.conv2, self.bn2),
|
224 |
+
(self.conv3, self.bn3),
|
225 |
+
]:
|
226 |
+
x = self.relu(bn(conv(x)))
|
227 |
+
x = self.avgpool(x)
|
228 |
+
return x
|
229 |
+
|
230 |
+
def forward(self, x):
|
231 |
+
x = self.stem(x)
|
232 |
+
x = self.layer1(x)
|
233 |
+
x = self.layer2(x)
|
234 |
+
x = self.layer3(x)
|
235 |
+
x = self.layer4(x)
|
236 |
+
x = self.attnpool(x)
|
237 |
+
|
238 |
+
return x
|
239 |
+
|
240 |
+
|
241 |
+
class LayerNorm(nn.LayerNorm):
|
242 |
+
"""Subclass torch's LayerNorm to handle fp16."""
|
243 |
+
|
244 |
+
def forward(self, x: torch.Tensor):
|
245 |
+
orig_type = x.dtype
|
246 |
+
x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
|
247 |
+
return x.to(orig_type)
|
248 |
+
|
249 |
+
|
250 |
+
class QuickGELU(nn.Module):
|
251 |
+
# NOTE This is slower than nn.GELU or nn.SiLU and uses more GPU memory
|
252 |
+
def forward(self, x: torch.Tensor):
|
253 |
+
return x * torch.sigmoid(1.702 * x)
|
254 |
+
|
255 |
+
|
256 |
+
class ResidualAttentionBlock(nn.Module):
|
257 |
+
def __init__(self, d_model: int, n_head: int, act_layer: Callable = nn.GELU):
|
258 |
+
super().__init__()
|
259 |
+
|
260 |
+
self.attn = nn.MultiheadAttention(d_model, n_head)
|
261 |
+
self.ln_1 = LayerNorm(d_model)
|
262 |
+
self.mlp = nn.Sequential(
|
263 |
+
OrderedDict(
|
264 |
+
[
|
265 |
+
("c_fc", nn.Linear(d_model, d_model * 4)),
|
266 |
+
("gelu", act_layer()),
|
267 |
+
("c_proj", nn.Linear(d_model * 4, d_model)),
|
268 |
+
]
|
269 |
+
)
|
270 |
+
)
|
271 |
+
self.ln_2 = LayerNorm(d_model)
|
272 |
+
|
273 |
+
def attention(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
|
274 |
+
return self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
|
275 |
+
|
276 |
+
def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
|
277 |
+
x = x + self.attention(self.ln_1(x), attn_mask=attn_mask)
|
278 |
+
x = x + self.mlp(self.ln_2(x))
|
279 |
+
return x
|
280 |
+
|
281 |
+
|
282 |
+
class Transformer(nn.Module):
|
283 |
+
def __init__(
|
284 |
+
self, width: int, layers: int, heads: int, act_layer: Callable = nn.GELU
|
285 |
+
):
|
286 |
+
super().__init__()
|
287 |
+
self.width = width
|
288 |
+
self.layers = layers
|
289 |
+
self.resblocks = nn.ModuleList(
|
290 |
+
[
|
291 |
+
ResidualAttentionBlock(width, heads, act_layer=act_layer)
|
292 |
+
for _ in range(layers)
|
293 |
+
]
|
294 |
+
)
|
295 |
+
|
296 |
+
def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
|
297 |
+
for r in self.resblocks:
|
298 |
+
x = r(x, attn_mask=attn_mask)
|
299 |
+
return x
|
300 |
+
|
301 |
+
|
302 |
+
class VisualTransformer(nn.Module):
|
303 |
+
def __init__(
|
304 |
+
self,
|
305 |
+
image_size: int,
|
306 |
+
patch_size: int,
|
307 |
+
width: int,
|
308 |
+
layers: int,
|
309 |
+
heads: int,
|
310 |
+
output_dim: int,
|
311 |
+
act_layer: Callable = nn.GELU,
|
312 |
+
):
|
313 |
+
super().__init__()
|
314 |
+
self.image_size = image_size
|
315 |
+
self.output_dim = output_dim
|
316 |
+
self.conv1 = nn.Conv2d(
|
317 |
+
in_channels=3,
|
318 |
+
out_channels=width,
|
319 |
+
kernel_size=patch_size,
|
320 |
+
stride=patch_size,
|
321 |
+
bias=False,
|
322 |
+
)
|
323 |
+
|
324 |
+
scale = width**-0.5
|
325 |
+
self.class_embedding = nn.Parameter(scale * torch.randn(width))
|
326 |
+
self.positional_embedding = nn.Parameter(
|
327 |
+
scale * torch.randn((image_size // patch_size) ** 2 + 1, width)
|
328 |
+
)
|
329 |
+
self.ln_pre = LayerNorm(width)
|
330 |
+
|
331 |
+
self.text_branch = Transformer(width, layers, heads, act_layer=act_layer)
|
332 |
+
|
333 |
+
self.ln_post = LayerNorm(width)
|
334 |
+
self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
|
335 |
+
|
336 |
+
def lock(self, unlocked_groups=0, freeze_bn_stats=False):
|
337 |
+
assert (
|
338 |
+
unlocked_groups == 0
|
339 |
+
), "partial locking not currently supported for this model"
|
340 |
+
for param in self.parameters():
|
341 |
+
param.requires_grad = False
|
342 |
+
|
343 |
+
def forward(self, x: torch.Tensor):
|
344 |
+
x = self.conv1(x) # shape = [*, width, grid, grid]
|
345 |
+
x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2]
|
346 |
+
x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
|
347 |
+
x = torch.cat(
|
348 |
+
[
|
349 |
+
self.class_embedding.to(x.dtype)
|
350 |
+
+ torch.zeros(
|
351 |
+
x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device
|
352 |
+
),
|
353 |
+
x,
|
354 |
+
],
|
355 |
+
dim=1,
|
356 |
+
) # shape = [*, grid ** 2 + 1, width]
|
357 |
+
x = x + self.positional_embedding.to(x.dtype)
|
358 |
+
x = self.ln_pre(x)
|
359 |
+
|
360 |
+
x = x.permute(1, 0, 2) # NLD -> LND
|
361 |
+
x = self.text_branch(x)
|
362 |
+
x = x.permute(1, 0, 2) # LND -> NLD
|
363 |
+
|
364 |
+
x = self.ln_post(x[:, 0, :])
|
365 |
+
|
366 |
+
if self.proj is not None:
|
367 |
+
x = x @ self.proj
|
368 |
+
|
369 |
+
return x
|
370 |
+
|
371 |
+
|
372 |
+
@dataclass
|
373 |
+
class CLAPVisionCfg:
|
374 |
+
layers: Union[Tuple[int, int, int, int], int] = 12
|
375 |
+
width: int = 768
|
376 |
+
patch_size: int = 16
|
377 |
+
image_size: Union[Tuple[int, int], int] = 224
|
378 |
+
timm_model_name: str = (
|
379 |
+
None # a valid model name overrides layers, width, patch_size
|
380 |
+
)
|
381 |
+
timm_model_pretrained: bool = (
|
382 |
+
False # use (imagenet) pretrained weights for named model
|
383 |
+
)
|
384 |
+
timm_pool: str = (
|
385 |
+
"avg" # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
|
386 |
+
)
|
387 |
+
timm_proj: str = (
|
388 |
+
"linear" # linear projection for timm model output ('linear', 'mlp', '')
|
389 |
+
)
|
390 |
+
|
391 |
+
|
392 |
+
# Audio Config Class
|
393 |
+
@dataclass
|
394 |
+
class CLAPAudioCfp:
|
395 |
+
model_type: str = "PANN"
|
396 |
+
model_name: str = "Cnn14"
|
397 |
+
sample_rate: int = 48000
|
398 |
+
# Param
|
399 |
+
audio_length: int = 1024
|
400 |
+
window_size: int = 1024
|
401 |
+
hop_size: int = 1024
|
402 |
+
fmin: int = 50
|
403 |
+
fmax: int = 14000
|
404 |
+
class_num: int = 527
|
405 |
+
mel_bins: int = 64
|
406 |
+
clip_samples: int = 480000
|
407 |
+
|
408 |
+
|
409 |
+
@dataclass
|
410 |
+
class CLAPTextCfg:
|
411 |
+
context_length: int
|
412 |
+
vocab_size: int
|
413 |
+
width: int
|
414 |
+
heads: int
|
415 |
+
layers: int
|
416 |
+
model_type: str
|
417 |
+
|
418 |
+
|
419 |
+
class CLAP(nn.Module):
|
420 |
+
def __init__(
|
421 |
+
self,
|
422 |
+
embed_dim: int,
|
423 |
+
audio_cfg: CLAPAudioCfp,
|
424 |
+
text_cfg: CLAPTextCfg,
|
425 |
+
quick_gelu: bool = False,
|
426 |
+
enable_fusion: bool = False,
|
427 |
+
fusion_type: str = "None",
|
428 |
+
joint_embed_shape: int = 512,
|
429 |
+
mlp_act: str = "relu",
|
430 |
+
):
|
431 |
+
super().__init__()
|
432 |
+
if isinstance(audio_cfg, dict):
|
433 |
+
audio_cfg = CLAPAudioCfp(**audio_cfg)
|
434 |
+
if isinstance(text_cfg, dict):
|
435 |
+
text_cfg = CLAPTextCfg(**text_cfg)
|
436 |
+
|
437 |
+
self.audio_cfg = audio_cfg
|
438 |
+
self.text_cfg = text_cfg
|
439 |
+
self.enable_fusion = enable_fusion
|
440 |
+
self.fusion_type = fusion_type
|
441 |
+
self.joint_embed_shape = joint_embed_shape
|
442 |
+
self.mlp_act = mlp_act
|
443 |
+
|
444 |
+
self.context_length = text_cfg.context_length
|
445 |
+
|
446 |
+
# OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
|
447 |
+
# memory efficient in recent PyTorch releases (>= 1.10).
|
448 |
+
# NOTE: timm models always use native GELU regardless of quick_gelu flag.
|
449 |
+
act_layer = QuickGELU if quick_gelu else nn.GELU
|
450 |
+
|
451 |
+
if mlp_act == "relu":
|
452 |
+
mlp_act_layer = nn.ReLU()
|
453 |
+
elif mlp_act == "gelu":
|
454 |
+
mlp_act_layer = nn.GELU()
|
455 |
+
else:
|
456 |
+
raise NotImplementedError
|
457 |
+
|
458 |
+
# audio branch
|
459 |
+
# audio branch parameters
|
460 |
+
if audio_cfg.model_type == "PANN":
|
461 |
+
self.audio_branch = create_pann_model(audio_cfg, enable_fusion, fusion_type)
|
462 |
+
elif audio_cfg.model_type == "HTSAT":
|
463 |
+
self.audio_branch = create_htsat_model(
|
464 |
+
audio_cfg, enable_fusion, fusion_type
|
465 |
+
)
|
466 |
+
else:
|
467 |
+
logging.error(f"Model config for {audio_cfg.model_type} not found")
|
468 |
+
raise RuntimeError(f"Model config for {audio_cfg.model_type} not found.")
|
469 |
+
|
470 |
+
# text branch
|
471 |
+
# text branch parameters
|
472 |
+
if text_cfg.model_type == "transformer":
|
473 |
+
self.text_branch = Transformer(
|
474 |
+
width=text_cfg.width,
|
475 |
+
layers=text_cfg.layers,
|
476 |
+
heads=text_cfg.heads,
|
477 |
+
act_layer=act_layer,
|
478 |
+
)
|
479 |
+
self.vocab_size = text_cfg.vocab_size
|
480 |
+
self.token_embedding = nn.Embedding(text_cfg.vocab_size, text_cfg.width)
|
481 |
+
self.positional_embedding = nn.Parameter(
|
482 |
+
torch.empty(self.context_length, text_cfg.width)
|
483 |
+
)
|
484 |
+
self.ln_final = LayerNorm(text_cfg.width)
|
485 |
+
self.text_transform = MLPLayers(
|
486 |
+
units=[
|
487 |
+
self.joint_embed_shape,
|
488 |
+
self.joint_embed_shape,
|
489 |
+
self.joint_embed_shape,
|
490 |
+
],
|
491 |
+
dropout=0.1,
|
492 |
+
)
|
493 |
+
self.text_projection = nn.Sequential(
|
494 |
+
nn.Linear(text_cfg.width, self.joint_embed_shape),
|
495 |
+
mlp_act_layer,
|
496 |
+
nn.Linear(self.joint_embed_shape, self.joint_embed_shape),
|
497 |
+
)
|
498 |
+
elif text_cfg.model_type == "bert":
|
499 |
+
self.text_branch = BertModel.from_pretrained("bert-base-uncased")
|
500 |
+
self.text_transform = MLPLayers(
|
501 |
+
units=[
|
502 |
+
self.joint_embed_shape,
|
503 |
+
self.joint_embed_shape,
|
504 |
+
self.joint_embed_shape,
|
505 |
+
],
|
506 |
+
dropout=0.1,
|
507 |
+
)
|
508 |
+
self.text_projection = nn.Sequential(
|
509 |
+
nn.Linear(768, self.joint_embed_shape),
|
510 |
+
mlp_act_layer,
|
511 |
+
nn.Linear(self.joint_embed_shape, self.joint_embed_shape),
|
512 |
+
)
|
513 |
+
elif text_cfg.model_type == "roberta":
|
514 |
+
self.text_branch = RobertaModel(
|
515 |
+
RobertaConfig.from_pretrained("roberta-base")
|
516 |
+
)
|
517 |
+
self.text_transform = MLPLayers(
|
518 |
+
units=[
|
519 |
+
self.joint_embed_shape,
|
520 |
+
self.joint_embed_shape,
|
521 |
+
self.joint_embed_shape,
|
522 |
+
],
|
523 |
+
dropout=0.1,
|
524 |
+
)
|
525 |
+
self.text_projection = nn.Sequential(
|
526 |
+
nn.Linear(768, self.joint_embed_shape),
|
527 |
+
mlp_act_layer,
|
528 |
+
nn.Linear(self.joint_embed_shape, self.joint_embed_shape),
|
529 |
+
)
|
530 |
+
elif text_cfg.model_type == "bart":
|
531 |
+
self.text_branch = BartModel.from_pretrained("facebook/bart-base")
|
532 |
+
self.text_transform = MLPLayers(
|
533 |
+
units=[
|
534 |
+
self.joint_embed_shape,
|
535 |
+
self.joint_embed_shape,
|
536 |
+
self.joint_embed_shape,
|
537 |
+
],
|
538 |
+
dropout=0.1,
|
539 |
+
)
|
540 |
+
self.text_projection = nn.Sequential(
|
541 |
+
nn.Linear(768, self.joint_embed_shape),
|
542 |
+
mlp_act_layer,
|
543 |
+
nn.Linear(self.joint_embed_shape, self.joint_embed_shape),
|
544 |
+
)
|
545 |
+
else:
|
546 |
+
logging.error(f"Model config for {text_cfg.model_type} not found")
|
547 |
+
raise RuntimeError(f"Model config for {text_cfg.model_type} not found.")
|
548 |
+
self.text_branch_type = text_cfg.model_type
|
549 |
+
# text branch parameters
|
550 |
+
|
551 |
+
# audio branch parameters
|
552 |
+
self.audio_transform = MLPLayers(
|
553 |
+
units=[
|
554 |
+
self.joint_embed_shape,
|
555 |
+
self.joint_embed_shape,
|
556 |
+
self.joint_embed_shape,
|
557 |
+
],
|
558 |
+
dropout=0.1,
|
559 |
+
)
|
560 |
+
|
561 |
+
# below here is text branch parameters
|
562 |
+
|
563 |
+
# ============================================================================================================
|
564 |
+
self.audio_projection = nn.Sequential(
|
565 |
+
nn.Linear(embed_dim, self.joint_embed_shape),
|
566 |
+
mlp_act_layer,
|
567 |
+
nn.Linear(self.joint_embed_shape, self.joint_embed_shape),
|
568 |
+
)
|
569 |
+
|
570 |
+
self.logit_scale_a = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
|
571 |
+
self.logit_scale_t = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
|
572 |
+
self.register_buffer("attn_mask", self.build_attention_mask(), persistent=False)
|
573 |
+
|
574 |
+
self.init_text_branch_parameters()
|
575 |
+
|
576 |
+
def init_text_branch_parameters(self):
|
577 |
+
if self.text_branch_type == "transformer":
|
578 |
+
nn.init.normal_(self.token_embedding.weight, std=0.02)
|
579 |
+
nn.init.normal_(self.positional_embedding, std=0.01)
|
580 |
+
proj_std = (self.text_branch.width**-0.5) * (
|
581 |
+
(2 * self.text_branch.layers) ** -0.5
|
582 |
+
)
|
583 |
+
attn_std = self.text_branch.width**-0.5
|
584 |
+
fc_std = (2 * self.text_branch.width) ** -0.5
|
585 |
+
for block in self.text_branch.resblocks:
|
586 |
+
nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
|
587 |
+
nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
|
588 |
+
nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
|
589 |
+
nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
|
590 |
+
if self.text_branch_type == "bert" or self.text_branch_type == "roberta":
|
591 |
+
self.text_branch.embeddings.word_embeddings.weight.shape[-1]
|
592 |
+
elif self.text_branch_type == "bart":
|
593 |
+
self.text_branch.shared.weight.shape[-1]
|
594 |
+
else:
|
595 |
+
self.text_branch.width
|
596 |
+
nn.init.constant_(self.logit_scale_a, np.log(1 / 0.07))
|
597 |
+
nn.init.constant_(self.logit_scale_t, np.log(1 / 0.07))
|
598 |
+
|
599 |
+
# deprecated
|
600 |
+
# if hasattr(self.visual, 'init_parameters'):
|
601 |
+
# self.visual.init_parameters()
|
602 |
+
|
603 |
+
# if self.text_projection is not None:
|
604 |
+
# nn.init.normal_(self.text_projection, std=width**-0.5)
|
605 |
+
|
606 |
+
def build_attention_mask(self):
|
607 |
+
# lazily create causal attention mask, with full attention between the vision tokens
|
608 |
+
# pytorch uses additive attention mask; fill with -inf
|
609 |
+
mask = torch.empty(self.context_length, self.context_length)
|
610 |
+
mask.fill_(float("-inf"))
|
611 |
+
mask.triu_(1) # zero out the lower diagonal
|
612 |
+
return mask
|
613 |
+
|
614 |
+
def encode_audio(self, audio, device):
|
615 |
+
return self.audio_branch(
|
616 |
+
audio, mixup_lambda=None, device=device
|
617 |
+
) # mix lambda needs to add
|
618 |
+
|
619 |
+
# def list_of_dict_of_tensor2dict_of_tensor(self, x, device):
|
620 |
+
# tmp = {}
|
621 |
+
# for k in x[0].keys():
|
622 |
+
# tmp[k] = []
|
623 |
+
# for i in range(len(x)):
|
624 |
+
# tmp[k].append(x[i][k][:77])
|
625 |
+
# for k in x[0].keys():
|
626 |
+
# tmp[k] = torch.tensor(tmp[k]).to(device=device, non_blocking=True)
|
627 |
+
# return tmp
|
628 |
+
|
629 |
+
def encode_text(self, text, device):
|
630 |
+
if self.text_branch_type == "transformer":
|
631 |
+
text = text.to(device=device, non_blocking=True)
|
632 |
+
x = self.token_embedding(text) # [batch_size, n_ctx, d_model]
|
633 |
+
|
634 |
+
x = x + self.positional_embedding
|
635 |
+
x = x.permute(1, 0, 2) # NLD -> LND
|
636 |
+
x = self.text_branch(x, attn_mask=self.attn_mask)
|
637 |
+
x = x.permute(1, 0, 2) # LND -> NLD
|
638 |
+
x = self.ln_final(x)
|
639 |
+
|
640 |
+
# x.shape = [batch_size, n_ctx, transformer.width]
|
641 |
+
# take features from the eot embedding (eot_token is the highest number in each sequence)
|
642 |
+
x = self.text_projection(x[torch.arange(x.shape[0]), text.argmax(dim=-1)])
|
643 |
+
elif self.text_branch_type == "bert":
|
644 |
+
# text = self.list_of_dict_of_tensor2dict_of_tensor(text, device)
|
645 |
+
# text = BatchEncoding(text)
|
646 |
+
x = self.text_branch(
|
647 |
+
input_ids=text["input_ids"].to(device=device, non_blocking=True),
|
648 |
+
attention_mask=text["attention_mask"].to(
|
649 |
+
device=device, non_blocking=True
|
650 |
+
),
|
651 |
+
token_type_ids=text["token_type_ids"].to(
|
652 |
+
device=device, non_blocking=True
|
653 |
+
),
|
654 |
+
)["pooler_output"]
|
655 |
+
x = self.text_projection(x)
|
656 |
+
elif self.text_branch_type == "roberta":
|
657 |
+
x = self.text_branch(
|
658 |
+
input_ids=text["input_ids"].to(device=device, non_blocking=True),
|
659 |
+
attention_mask=text["attention_mask"].to(
|
660 |
+
device=device, non_blocking=True
|
661 |
+
),
|
662 |
+
)["pooler_output"]
|
663 |
+
x = self.text_projection(x)
|
664 |
+
elif self.text_branch_type == "bart":
|
665 |
+
x = torch.mean(
|
666 |
+
self.text_branch(
|
667 |
+
input_ids=text["input_ids"].to(device=device, non_blocking=True),
|
668 |
+
attention_mask=text["attention_mask"].to(
|
669 |
+
device=device, non_blocking=True
|
670 |
+
),
|
671 |
+
)["encoder_last_hidden_state"],
|
672 |
+
axis=1,
|
673 |
+
)
|
674 |
+
x = self.text_projection(x)
|
675 |
+
else:
|
676 |
+
logging.error(f"Model type {self.text_branch_type} not found")
|
677 |
+
raise RuntimeError(f"Model type {self.text_branch_type} not found.")
|
678 |
+
return x
|
679 |
+
|
680 |
+
def forward(self, audio, text, device=None):
|
681 |
+
"""Forward audio and text into the CLAP
|
682 |
+
|
683 |
+
Parameters
|
684 |
+
----------
|
685 |
+
audio: torch.Tensor (batch_size, audio_length)
|
686 |
+
the time-domain audio input / the batch of mel_spec and longer list.
|
687 |
+
text: torch.Tensor () // need to add
|
688 |
+
the text token input
|
689 |
+
"""
|
690 |
+
if device is None:
|
691 |
+
if audio is not None:
|
692 |
+
device = audio.device
|
693 |
+
elif text is not None:
|
694 |
+
device = text.device
|
695 |
+
if audio is None and text is None:
|
696 |
+
# a hack to get the logit scale
|
697 |
+
return self.logit_scale_a.exp(), self.logit_scale_t.exp()
|
698 |
+
elif audio is None:
|
699 |
+
return self.encode_text(text, device=device)
|
700 |
+
elif text is None:
|
701 |
+
return self.audio_projection(
|
702 |
+
self.encode_audio(audio, device=device)["embedding"]
|
703 |
+
)
|
704 |
+
audio_features = self.audio_projection(
|
705 |
+
self.encode_audio(audio, device=device)["embedding"]
|
706 |
+
)
|
707 |
+
audio_features = F.normalize(audio_features, dim=-1)
|
708 |
+
|
709 |
+
text_features = self.encode_text(text, device=device)
|
710 |
+
# print("text_features", text_features)
|
711 |
+
# print("text_features.shape", text_features.shape)
|
712 |
+
# print("text_features.type", type(text_features))
|
713 |
+
text_features = F.normalize(text_features, dim=-1)
|
714 |
+
|
715 |
+
audio_features_mlp = self.audio_transform(audio_features)
|
716 |
+
text_features_mlp = self.text_transform(text_features)
|
717 |
+
# Four outputs: audio features (basic & MLP), text features (basic & MLP)
|
718 |
+
return (
|
719 |
+
audio_features,
|
720 |
+
text_features,
|
721 |
+
audio_features_mlp,
|
722 |
+
text_features_mlp,
|
723 |
+
self.logit_scale_a.exp(),
|
724 |
+
self.logit_scale_t.exp(),
|
725 |
+
)
|
726 |
+
|
727 |
+
def get_logit_scale(self):
|
728 |
+
return self.logit_scale_a.exp(), self.logit_scale_t.exp()
|
729 |
+
|
730 |
+
def get_text_embedding(self, data):
|
731 |
+
"""Get the text embedding from the model
|
732 |
+
|
733 |
+
Parameters
|
734 |
+
----------
|
735 |
+
data: torch.Tensor
|
736 |
+
a tensor of text embedding
|
737 |
+
|
738 |
+
Returns
|
739 |
+
----------
|
740 |
+
text_embed: torch.Tensor
|
741 |
+
a tensor of text_embeds (N, D)
|
742 |
+
|
743 |
+
"""
|
744 |
+
device = next(self.parameters()).device
|
745 |
+
for k in data:
|
746 |
+
data[k] = data[k].to(device)
|
747 |
+
text_embeds = self.encode_text(data, device=device)
|
748 |
+
text_embeds = F.normalize(text_embeds, dim=-1)
|
749 |
+
|
750 |
+
return text_embeds
|
751 |
+
|
752 |
+
def get_audio_embedding(self, data):
|
753 |
+
"""Get the audio embedding from the model
|
754 |
+
|
755 |
+
Parameters
|
756 |
+
----------
|
757 |
+
data: a list of dict
|
758 |
+
the audio input dict list from 'get_audio_feature' method
|
759 |
+
|
760 |
+
Returns
|
761 |
+
----------
|
762 |
+
audio_embed: torch.Tensor
|
763 |
+
a tensor of audio_embeds (N, D)
|
764 |
+
|
765 |
+
"""
|
766 |
+
device = next(self.parameters()).device
|
767 |
+
# input_dict = {}
|
768 |
+
# keys = data[0].keys()
|
769 |
+
# for k in keys:
|
770 |
+
# input_dict[k] = torch.cat([d[k].unsqueeze(0) for d in data], dim=0).to(
|
771 |
+
# device
|
772 |
+
# )
|
773 |
+
audio_embeds = self.audio_projection(
|
774 |
+
self.encode_audio(data, device=device)["embedding"]
|
775 |
+
)
|
776 |
+
audio_embeds = F.normalize(audio_embeds, dim=-1)
|
777 |
+
|
778 |
+
return audio_embeds
|
779 |
+
|
780 |
+
def audio_infer(self, audio, hopsize=None, device=None):
|
781 |
+
"""Forward one audio and produce the audio embedding
|
782 |
+
|
783 |
+
Parameters
|
784 |
+
----------
|
785 |
+
audio: (audio_length)
|
786 |
+
the time-domain audio input, notice that it must be only one input
|
787 |
+
hopsize: int
|
788 |
+
the overlap hopsize as the sliding window
|
789 |
+
|
790 |
+
Returns
|
791 |
+
----------
|
792 |
+
output_dict: {
|
793 |
+
key: [n, (embedding_shape)] if "HTS-AT"
|
794 |
+
or
|
795 |
+
key: [(embedding_shape)] if "PANN"
|
796 |
+
}
|
797 |
+
the list of key values of the audio branch
|
798 |
+
|
799 |
+
"""
|
800 |
+
|
801 |
+
assert not self.training, "the inference mode must be run at eval stage"
|
802 |
+
output_dict = {}
|
803 |
+
# PANN
|
804 |
+
if self.audio_cfg.model_type == "PANN":
|
805 |
+
audio_input = audio.unsqueeze(dim=0)
|
806 |
+
output_dict[key] = self.encode_audio(audio_input, device=device)[
|
807 |
+
key
|
808 |
+
].squeeze(dim=0)
|
809 |
+
elif self.audio_cfg.model_type == "HTSAT":
|
810 |
+
# repeat
|
811 |
+
audio_len = len(audio)
|
812 |
+
k = self.audio_cfg.clip_samples // audio_len
|
813 |
+
if k > 1:
|
814 |
+
audio = audio.repeat(k)
|
815 |
+
audio_len = len(audio)
|
816 |
+
|
817 |
+
if hopsize is None:
|
818 |
+
hopsize = min(hopsize, audio_len)
|
819 |
+
|
820 |
+
if audio_len > self.audio_cfg.clip_samples:
|
821 |
+
audio_input = [
|
822 |
+
audio[pos : pos + self.audio_cfg.clip_samples].clone()
|
823 |
+
for pos in range(
|
824 |
+
0, audio_len - self.audio_cfg.clip_samples, hopsize
|
825 |
+
)
|
826 |
+
]
|
827 |
+
audio_input.append(audio[-self.audio_cfg.clip_samples :].clone())
|
828 |
+
audio_input = torch.stack(audio_input)
|
829 |
+
output_dict[key] = self.encode_audio(audio_input, device=device)[key]
|
830 |
+
else:
|
831 |
+
audio_input = audio.unsqueeze(dim=0)
|
832 |
+
output_dict[key] = self.encode_audio(audio_input, device=device)[
|
833 |
+
key
|
834 |
+
].squeeze(dim=0)
|
835 |
+
|
836 |
+
return output_dict
|
837 |
+
|
838 |
+
|
839 |
+
def convert_weights_to_fp16(model: nn.Module):
|
840 |
+
"""Convert applicable model parameters to fp16"""
|
841 |
+
|
842 |
+
def _convert_weights_to_fp16(l):
|
843 |
+
if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
|
844 |
+
l.weight.data = l.weight.data.half()
|
845 |
+
if l.bias is not None:
|
846 |
+
l.bias.data = l.bias.data.half()
|
847 |
+
|
848 |
+
if isinstance(l, nn.MultiheadAttention):
|
849 |
+
for attr in [
|
850 |
+
*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]],
|
851 |
+
"in_proj_bias",
|
852 |
+
"bias_k",
|
853 |
+
"bias_v",
|
854 |
+
]:
|
855 |
+
tensor = getattr(l, attr)
|
856 |
+
if tensor is not None:
|
857 |
+
tensor.data = tensor.data.half()
|
858 |
+
|
859 |
+
for name in ["text_projection", "proj"]:
|
860 |
+
if hasattr(l, name):
|
861 |
+
attr = getattr(l, name)
|
862 |
+
if attr is not None:
|
863 |
+
attr.data = attr.data.half()
|
864 |
+
|
865 |
+
model.apply(_convert_weights_to_fp16)
|
866 |
+
|
867 |
+
|
868 |
+
# Ignore the state dict of the vision part
|
869 |
+
def build_model_from_openai_state_dict(
|
870 |
+
state_dict: dict, model_cfg, enable_fusion: bool = False, fusion_type: str = "None"
|
871 |
+
):
|
872 |
+
embed_dim = model_cfg["embed_dim"]
|
873 |
+
audio_cfg = model_cfg["audio_cfg"]
|
874 |
+
text_cfg = model_cfg["text_cfg"]
|
875 |
+
state_dict["positional_embedding"].shape[0]
|
876 |
+
state_dict["token_embedding.weight"].shape[0]
|
877 |
+
transformer_width = state_dict["ln_final.weight"].shape[0]
|
878 |
+
transformer_width // 64
|
879 |
+
transformer_layers = len(
|
880 |
+
set(
|
881 |
+
k.split(".")[2]
|
882 |
+
for k in state_dict
|
883 |
+
if k.startswith(f"transformer.resblocks")
|
884 |
+
)
|
885 |
+
)
|
886 |
+
|
887 |
+
audio_cfg = CLAPAudioCfp(**audio_cfg)
|
888 |
+
text_cfg = CLAPTextCfg(**text_cfg)
|
889 |
+
|
890 |
+
model = CLAP(
|
891 |
+
embed_dim,
|
892 |
+
audio_cfg=audio_cfg,
|
893 |
+
text_cfg=text_cfg,
|
894 |
+
quick_gelu=True, # OpenAI models were trained with QuickGELU
|
895 |
+
enable_fusion=enable_fusion,
|
896 |
+
fusion_type=fusion_type,
|
897 |
+
)
|
898 |
+
state_dict["logit_scale_a"] = state_dict["logit_scale"]
|
899 |
+
state_dict["logit_scale_t"] = state_dict["logit_scale"]
|
900 |
+
pop_keys = list(state_dict.keys())[::]
|
901 |
+
# pop the visual branch saved weights
|
902 |
+
for key in pop_keys:
|
903 |
+
if key.startswith("visual."):
|
904 |
+
state_dict.pop(key, None)
|
905 |
+
|
906 |
+
for key in ["logit_scale", "input_resolution", "context_length", "vocab_size"]:
|
907 |
+
state_dict.pop(key, None)
|
908 |
+
|
909 |
+
# not use fp16
|
910 |
+
# convert_weights_to_fp16(model)
|
911 |
+
model.load_state_dict(state_dict, strict=False)
|
912 |
+
return model.eval()
|
913 |
+
|
914 |
+
|
915 |
+
def trace_model(model, batch_size=256, device=torch.device("cpu")):
|
916 |
+
model.eval()
|
917 |
+
audio_length = model.audio_cfg.audio_length
|
918 |
+
example_audio = torch.ones((batch_size, audio_length), device=device)
|
919 |
+
example_text = torch.zeros(
|
920 |
+
(batch_size, model.context_length), dtype=torch.int, device=device
|
921 |
+
)
|
922 |
+
model = torch.jit.trace_module(
|
923 |
+
model,
|
924 |
+
inputs=dict(
|
925 |
+
forward=(example_audio, example_text),
|
926 |
+
encode_text=(example_text,),
|
927 |
+
encode_image=(example_audio,),
|
928 |
+
),
|
929 |
+
)
|
930 |
+
model.audio_cfg.audio_length = audio_length # Question: what does this do?
|
931 |
+
return model
|
audioldm2/clap/open_clip/model_configs/HTSAT-base.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 1024,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "HTSAT",
|
14 |
+
"model_name": "base"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/HTSAT-large.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 2048,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "HTSAT",
|
14 |
+
"model_name": "large"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/HTSAT-tiny-win-1536.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 768,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1536,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "HTSAT",
|
14 |
+
"model_name": "tiny"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/HTSAT-tiny.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 768,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "HTSAT",
|
14 |
+
"model_name": "tiny"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/PANN-10.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 1024,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "PANN",
|
14 |
+
"model_name": "Cnn10"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/PANN-14-fmax-18k.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 2048,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 18000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "PANN",
|
14 |
+
"model_name": "Cnn14"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/PANN-14-fmax-8k-20s.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 2048,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 960000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 360,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 8000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "PANN",
|
14 |
+
"model_name": "Cnn14"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/PANN-14-tiny-transformer.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 2048,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "PANN",
|
14 |
+
"model_name": "Cnn14"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 4
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/PANN-14-win-1536.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 2048,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1536,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "PANN",
|
14 |
+
"model_name": "Cnn14"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/PANN-14.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 2048,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "PANN",
|
14 |
+
"model_name": "Cnn14"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/PANN-6.json
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 512,
|
3 |
+
"audio_cfg": {
|
4 |
+
"audio_length": 1024,
|
5 |
+
"clip_samples": 480000,
|
6 |
+
"mel_bins": 64,
|
7 |
+
"sample_rate": 48000,
|
8 |
+
"window_size": 1024,
|
9 |
+
"hop_size": 480,
|
10 |
+
"fmin": 50,
|
11 |
+
"fmax": 14000,
|
12 |
+
"class_num": 527,
|
13 |
+
"model_type": "PANN",
|
14 |
+
"model_name": "Cnn6"
|
15 |
+
},
|
16 |
+
"text_cfg": {
|
17 |
+
"context_length": 77,
|
18 |
+
"vocab_size": 49408,
|
19 |
+
"width": 512,
|
20 |
+
"heads": 8,
|
21 |
+
"layers": 12
|
22 |
+
}
|
23 |
+
}
|
audioldm2/clap/open_clip/model_configs/RN101-quickgelu.json
ADDED
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 512,
|
3 |
+
"quick_gelu": true,
|
4 |
+
"vision_cfg": {
|
5 |
+
"image_size": 224,
|
6 |
+
"layers": [
|
7 |
+
3,
|
8 |
+
4,
|
9 |
+
23,
|
10 |
+
3
|
11 |
+
],
|
12 |
+
"width": 64,
|
13 |
+
"patch_size": null
|
14 |
+
},
|
15 |
+
"text_cfg": {
|
16 |
+
"context_length": 77,
|
17 |
+
"vocab_size": 49408,
|
18 |
+
"width": 512,
|
19 |
+
"heads": 8,
|
20 |
+
"layers": 12
|
21 |
+
}
|
22 |
+
}
|
audioldm2/clap/open_clip/model_configs/RN101.json
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 512,
|
3 |
+
"vision_cfg": {
|
4 |
+
"image_size": 224,
|
5 |
+
"layers": [
|
6 |
+
3,
|
7 |
+
4,
|
8 |
+
23,
|
9 |
+
3
|
10 |
+
],
|
11 |
+
"width": 64,
|
12 |
+
"patch_size": null
|
13 |
+
},
|
14 |
+
"text_cfg": {
|
15 |
+
"context_length": 77,
|
16 |
+
"vocab_size": 49408,
|
17 |
+
"width": 512,
|
18 |
+
"heads": 8,
|
19 |
+
"layers": 12
|
20 |
+
}
|
21 |
+
}
|
audioldm2/clap/open_clip/model_configs/RN50-quickgelu.json
ADDED
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 1024,
|
3 |
+
"quick_gelu": true,
|
4 |
+
"vision_cfg": {
|
5 |
+
"image_size": 224,
|
6 |
+
"layers": [
|
7 |
+
3,
|
8 |
+
4,
|
9 |
+
6,
|
10 |
+
3
|
11 |
+
],
|
12 |
+
"width": 64,
|
13 |
+
"patch_size": null
|
14 |
+
},
|
15 |
+
"text_cfg": {
|
16 |
+
"context_length": 77,
|
17 |
+
"vocab_size": 49408,
|
18 |
+
"width": 512,
|
19 |
+
"heads": 8,
|
20 |
+
"layers": 12
|
21 |
+
}
|
22 |
+
}
|
audioldm2/clap/open_clip/model_configs/RN50.json
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 1024,
|
3 |
+
"vision_cfg": {
|
4 |
+
"image_size": 224,
|
5 |
+
"layers": [
|
6 |
+
3,
|
7 |
+
4,
|
8 |
+
6,
|
9 |
+
3
|
10 |
+
],
|
11 |
+
"width": 64,
|
12 |
+
"patch_size": null
|
13 |
+
},
|
14 |
+
"text_cfg": {
|
15 |
+
"context_length": 77,
|
16 |
+
"vocab_size": 49408,
|
17 |
+
"width": 512,
|
18 |
+
"heads": 8,
|
19 |
+
"layers": 12
|
20 |
+
}
|
21 |
+
}
|
audioldm2/clap/open_clip/model_configs/RN50x16.json
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 768,
|
3 |
+
"vision_cfg": {
|
4 |
+
"image_size": 384,
|
5 |
+
"layers": [
|
6 |
+
6,
|
7 |
+
8,
|
8 |
+
18,
|
9 |
+
8
|
10 |
+
],
|
11 |
+
"width": 96,
|
12 |
+
"patch_size": null
|
13 |
+
},
|
14 |
+
"text_cfg": {
|
15 |
+
"context_length": 77,
|
16 |
+
"vocab_size": 49408,
|
17 |
+
"width": 768,
|
18 |
+
"heads": 12,
|
19 |
+
"layers": 12
|
20 |
+
}
|
21 |
+
}
|
audioldm2/clap/open_clip/model_configs/RN50x4.json
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 640,
|
3 |
+
"vision_cfg": {
|
4 |
+
"image_size": 288,
|
5 |
+
"layers": [
|
6 |
+
4,
|
7 |
+
6,
|
8 |
+
10,
|
9 |
+
6
|
10 |
+
],
|
11 |
+
"width": 80,
|
12 |
+
"patch_size": null
|
13 |
+
},
|
14 |
+
"text_cfg": {
|
15 |
+
"context_length": 77,
|
16 |
+
"vocab_size": 49408,
|
17 |
+
"width": 640,
|
18 |
+
"heads": 10,
|
19 |
+
"layers": 12
|
20 |
+
}
|
21 |
+
}
|
audioldm2/clap/open_clip/model_configs/ViT-B-16.json
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 512,
|
3 |
+
"vision_cfg": {
|
4 |
+
"image_size": 224,
|
5 |
+
"layers": 12,
|
6 |
+
"width": 768,
|
7 |
+
"patch_size": 16
|
8 |
+
},
|
9 |
+
"text_cfg": {
|
10 |
+
"context_length": 77,
|
11 |
+
"vocab_size": 49408,
|
12 |
+
"width": 512,
|
13 |
+
"heads": 8,
|
14 |
+
"layers": 12
|
15 |
+
}
|
16 |
+
}
|
audioldm2/clap/open_clip/model_configs/ViT-B-32-quickgelu.json
ADDED
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 512,
|
3 |
+
"quick_gelu": true,
|
4 |
+
"vision_cfg": {
|
5 |
+
"image_size": 224,
|
6 |
+
"layers": 12,
|
7 |
+
"width": 768,
|
8 |
+
"patch_size": 32
|
9 |
+
},
|
10 |
+
"text_cfg": {
|
11 |
+
"context_length": 77,
|
12 |
+
"vocab_size": 49408,
|
13 |
+
"width": 512,
|
14 |
+
"heads": 8,
|
15 |
+
"layers": 12
|
16 |
+
}
|
17 |
+
}
|
audioldm2/clap/open_clip/model_configs/ViT-B-32.json
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 512,
|
3 |
+
"vision_cfg": {
|
4 |
+
"image_size": 224,
|
5 |
+
"layers": 12,
|
6 |
+
"width": 768,
|
7 |
+
"patch_size": 32
|
8 |
+
},
|
9 |
+
"text_cfg": {
|
10 |
+
"context_length": 77,
|
11 |
+
"vocab_size": 49408,
|
12 |
+
"width": 512,
|
13 |
+
"heads": 8,
|
14 |
+
"layers": 12
|
15 |
+
}
|
16 |
+
}
|
audioldm2/clap/open_clip/model_configs/ViT-L-14.json
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"embed_dim": 768,
|
3 |
+
"vision_cfg": {
|
4 |
+
"image_size": 224,
|
5 |
+
"layers": 24,
|
6 |
+
"width": 1024,
|
7 |
+
"patch_size": 14
|
8 |
+
},
|
9 |
+
"text_cfg": {
|
10 |
+
"context_length": 77,
|
11 |
+
"vocab_size": 49408,
|
12 |
+
"width": 768,
|
13 |
+
"heads": 12,
|
14 |
+
"layers": 12
|
15 |
+
}
|
16 |
+
}
|
audioldm2/clap/open_clip/openai.py
ADDED
@@ -0,0 +1,156 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
""" OpenAI pretrained model functions
|
2 |
+
|
3 |
+
Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
|
4 |
+
"""
|
5 |
+
|
6 |
+
import os
|
7 |
+
import warnings
|
8 |
+
from typing import Union, List
|
9 |
+
|
10 |
+
import torch
|
11 |
+
|
12 |
+
from .model import build_model_from_openai_state_dict
|
13 |
+
from .pretrained import (
|
14 |
+
get_pretrained_url,
|
15 |
+
list_pretrained_tag_models,
|
16 |
+
download_pretrained,
|
17 |
+
)
|
18 |
+
|
19 |
+
__all__ = ["list_openai_models", "load_openai_model"]
|
20 |
+
|
21 |
+
|
22 |
+
def list_openai_models() -> List[str]:
|
23 |
+
"""Returns the names of available CLIP models"""
|
24 |
+
return list_pretrained_tag_models("openai")
|
25 |
+
|
26 |
+
|
27 |
+
def load_openai_model(
|
28 |
+
name: str,
|
29 |
+
model_cfg,
|
30 |
+
device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu",
|
31 |
+
jit=True,
|
32 |
+
cache_dir=os.path.expanduser("~/.cache/clip"),
|
33 |
+
enable_fusion: bool = False,
|
34 |
+
fusion_type: str = "None",
|
35 |
+
):
|
36 |
+
"""Load a CLIP model, preserve its text pretrained part, and set in the CLAP model
|
37 |
+
|
38 |
+
Parameters
|
39 |
+
----------
|
40 |
+
name : str
|
41 |
+
A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
|
42 |
+
device : Union[str, torch.device]
|
43 |
+
The device to put the loaded model
|
44 |
+
jit : bool
|
45 |
+
Whether to load the optimized JIT model (default) or more hackable non-JIT model.
|
46 |
+
|
47 |
+
Returns
|
48 |
+
-------
|
49 |
+
model : torch.nn.Module
|
50 |
+
The CLAP model
|
51 |
+
preprocess : Callable[[PIL.Image], torch.Tensor]
|
52 |
+
A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
|
53 |
+
"""
|
54 |
+
if get_pretrained_url(name, "openai"):
|
55 |
+
model_path = download_pretrained(
|
56 |
+
get_pretrained_url(name, "openai"), root=cache_dir
|
57 |
+
)
|
58 |
+
elif os.path.isfile(name):
|
59 |
+
model_path = name
|
60 |
+
else:
|
61 |
+
raise RuntimeError(
|
62 |
+
f"Model {name} not found; available models = {list_openai_models()}"
|
63 |
+
)
|
64 |
+
|
65 |
+
try:
|
66 |
+
# loading JIT archive
|
67 |
+
model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval()
|
68 |
+
state_dict = None
|
69 |
+
except RuntimeError:
|
70 |
+
# loading saved state dict
|
71 |
+
if jit:
|
72 |
+
warnings.warn(
|
73 |
+
f"File {model_path} is not a JIT archive. Loading as a state dict instead"
|
74 |
+
)
|
75 |
+
jit = False
|
76 |
+
state_dict = torch.load(model_path, map_location="cpu")
|
77 |
+
|
78 |
+
if not jit:
|
79 |
+
try:
|
80 |
+
model = build_model_from_openai_state_dict(
|
81 |
+
state_dict or model.state_dict(), model_cfg, enable_fusion, fusion_type
|
82 |
+
).to(device)
|
83 |
+
except KeyError:
|
84 |
+
sd = {k[7:]: v for k, v in state_dict["state_dict"].items()}
|
85 |
+
model = build_model_from_openai_state_dict(
|
86 |
+
sd, model_cfg, enable_fusion, fusion_type
|
87 |
+
).to(device)
|
88 |
+
|
89 |
+
if str(device) == "cpu":
|
90 |
+
model.float()
|
91 |
+
return model
|
92 |
+
|
93 |
+
# patch the device names
|
94 |
+
device_holder = torch.jit.trace(
|
95 |
+
lambda: torch.ones([]).to(torch.device(device)), example_inputs=[]
|
96 |
+
)
|
97 |
+
device_node = [
|
98 |
+
n
|
99 |
+
for n in device_holder.graph.findAllNodes("prim::Constant")
|
100 |
+
if "Device" in repr(n)
|
101 |
+
][-1]
|
102 |
+
|
103 |
+
def patch_device(module):
|
104 |
+
try:
|
105 |
+
graphs = [module.graph] if hasattr(module, "graph") else []
|
106 |
+
except RuntimeError:
|
107 |
+
graphs = []
|
108 |
+
|
109 |
+
if hasattr(module, "forward1"):
|
110 |
+
graphs.append(module.forward1.graph)
|
111 |
+
|
112 |
+
for graph in graphs:
|
113 |
+
for node in graph.findAllNodes("prim::Constant"):
|
114 |
+
if "value" in node.attributeNames() and str(node["value"]).startswith(
|
115 |
+
"cuda"
|
116 |
+
):
|
117 |
+
node.copyAttributes(device_node)
|
118 |
+
|
119 |
+
model.apply(patch_device)
|
120 |
+
patch_device(model.encode_audio)
|
121 |
+
patch_device(model.encode_text)
|
122 |
+
|
123 |
+
# patch dtype to float32 on CPU
|
124 |
+
if str(device) == "cpu":
|
125 |
+
float_holder = torch.jit.trace(
|
126 |
+
lambda: torch.ones([]).float(), example_inputs=[]
|
127 |
+
)
|
128 |
+
float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
|
129 |
+
float_node = float_input.node()
|
130 |
+
|
131 |
+
def patch_float(module):
|
132 |
+
try:
|
133 |
+
graphs = [module.graph] if hasattr(module, "graph") else []
|
134 |
+
except RuntimeError:
|
135 |
+
graphs = []
|
136 |
+
|
137 |
+
if hasattr(module, "forward1"):
|
138 |
+
graphs.append(module.forward1.graph)
|
139 |
+
|
140 |
+
for graph in graphs:
|
141 |
+
for node in graph.findAllNodes("aten::to"):
|
142 |
+
inputs = list(node.inputs())
|
143 |
+
for i in [
|
144 |
+
1,
|
145 |
+
2,
|
146 |
+
]: # dtype can be the second or third argument to aten::to()
|
147 |
+
if inputs[i].node()["value"] == 5:
|
148 |
+
inputs[i].node().copyAttributes(float_node)
|
149 |
+
|
150 |
+
model.apply(patch_float)
|
151 |
+
patch_float(model.encode_audio)
|
152 |
+
patch_float(model.encode_text)
|
153 |
+
model.float()
|
154 |
+
|
155 |
+
model.audio_branch.audio_length = model.audio_cfg.audio_length
|
156 |
+
return model
|
audioldm2/clap/open_clip/pann_model.py
ADDED
@@ -0,0 +1,697 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# PANNs: Large-Scale Pretrained Audio Neural Networks for Audio Pattern Recognition
|
2 |
+
# Reference from https://github.com/qiuqiangkong/audioset_tagging_cnn
|
3 |
+
# Some layers are re-designed for CLAP
|
4 |
+
import os
|
5 |
+
|
6 |
+
os.environ["NUMBA_CACHE_DIR"] = "/tmp/"
|
7 |
+
|
8 |
+
import torch
|
9 |
+
import torch.nn as nn
|
10 |
+
import torch.nn.functional as F
|
11 |
+
from torchlibrosa.stft import Spectrogram, LogmelFilterBank
|
12 |
+
from torchlibrosa.augmentation import SpecAugmentation
|
13 |
+
|
14 |
+
from .utils import do_mixup, interpolate
|
15 |
+
from .feature_fusion import iAFF, AFF, DAF
|
16 |
+
|
17 |
+
|
18 |
+
def init_layer(layer):
|
19 |
+
"""Initialize a Linear or Convolutional layer."""
|
20 |
+
nn.init.xavier_uniform_(layer.weight)
|
21 |
+
|
22 |
+
if hasattr(layer, "bias"):
|
23 |
+
if layer.bias is not None:
|
24 |
+
layer.bias.data.fill_(0.0)
|
25 |
+
|
26 |
+
|
27 |
+
def init_bn(bn):
|
28 |
+
"""Initialize a Batchnorm layer."""
|
29 |
+
bn.bias.data.fill_(0.0)
|
30 |
+
bn.weight.data.fill_(1.0)
|
31 |
+
|
32 |
+
|
33 |
+
class ConvBlock(nn.Module):
|
34 |
+
def __init__(self, in_channels, out_channels):
|
35 |
+
super(ConvBlock, self).__init__()
|
36 |
+
|
37 |
+
self.conv1 = nn.Conv2d(
|
38 |
+
in_channels=in_channels,
|
39 |
+
out_channels=out_channels,
|
40 |
+
kernel_size=(3, 3),
|
41 |
+
stride=(1, 1),
|
42 |
+
padding=(1, 1),
|
43 |
+
bias=False,
|
44 |
+
)
|
45 |
+
|
46 |
+
self.conv2 = nn.Conv2d(
|
47 |
+
in_channels=out_channels,
|
48 |
+
out_channels=out_channels,
|
49 |
+
kernel_size=(3, 3),
|
50 |
+
stride=(1, 1),
|
51 |
+
padding=(1, 1),
|
52 |
+
bias=False,
|
53 |
+
)
|
54 |
+
|
55 |
+
self.bn1 = nn.BatchNorm2d(out_channels)
|
56 |
+
self.bn2 = nn.BatchNorm2d(out_channels)
|
57 |
+
|
58 |
+
self.init_weight()
|
59 |
+
|
60 |
+
def init_weight(self):
|
61 |
+
init_layer(self.conv1)
|
62 |
+
init_layer(self.conv2)
|
63 |
+
init_bn(self.bn1)
|
64 |
+
init_bn(self.bn2)
|
65 |
+
|
66 |
+
def forward(self, input, pool_size=(2, 2), pool_type="avg"):
|
67 |
+
x = input
|
68 |
+
x = F.relu_(self.bn1(self.conv1(x)))
|
69 |
+
x = F.relu_(self.bn2(self.conv2(x)))
|
70 |
+
if pool_type == "max":
|
71 |
+
x = F.max_pool2d(x, kernel_size=pool_size)
|
72 |
+
elif pool_type == "avg":
|
73 |
+
x = F.avg_pool2d(x, kernel_size=pool_size)
|
74 |
+
elif pool_type == "avg+max":
|
75 |
+
x1 = F.avg_pool2d(x, kernel_size=pool_size)
|
76 |
+
x2 = F.max_pool2d(x, kernel_size=pool_size)
|
77 |
+
x = x1 + x2
|
78 |
+
else:
|
79 |
+
raise Exception("Incorrect argument!")
|
80 |
+
|
81 |
+
return x
|
82 |
+
|
83 |
+
|
84 |
+
class ConvBlock5x5(nn.Module):
|
85 |
+
def __init__(self, in_channels, out_channels):
|
86 |
+
super(ConvBlock5x5, self).__init__()
|
87 |
+
|
88 |
+
self.conv1 = nn.Conv2d(
|
89 |
+
in_channels=in_channels,
|
90 |
+
out_channels=out_channels,
|
91 |
+
kernel_size=(5, 5),
|
92 |
+
stride=(1, 1),
|
93 |
+
padding=(2, 2),
|
94 |
+
bias=False,
|
95 |
+
)
|
96 |
+
|
97 |
+
self.bn1 = nn.BatchNorm2d(out_channels)
|
98 |
+
|
99 |
+
self.init_weight()
|
100 |
+
|
101 |
+
def init_weight(self):
|
102 |
+
init_layer(self.conv1)
|
103 |
+
init_bn(self.bn1)
|
104 |
+
|
105 |
+
def forward(self, input, pool_size=(2, 2), pool_type="avg"):
|
106 |
+
x = input
|
107 |
+
x = F.relu_(self.bn1(self.conv1(x)))
|
108 |
+
if pool_type == "max":
|
109 |
+
x = F.max_pool2d(x, kernel_size=pool_size)
|
110 |
+
elif pool_type == "avg":
|
111 |
+
x = F.avg_pool2d(x, kernel_size=pool_size)
|
112 |
+
elif pool_type == "avg+max":
|
113 |
+
x1 = F.avg_pool2d(x, kernel_size=pool_size)
|
114 |
+
x2 = F.max_pool2d(x, kernel_size=pool_size)
|
115 |
+
x = x1 + x2
|
116 |
+
else:
|
117 |
+
raise Exception("Incorrect argument!")
|
118 |
+
|
119 |
+
return x
|
120 |
+
|
121 |
+
|
122 |
+
class AttBlock(nn.Module):
|
123 |
+
def __init__(self, n_in, n_out, activation="linear", temperature=1.0):
|
124 |
+
super(AttBlock, self).__init__()
|
125 |
+
|
126 |
+
self.activation = activation
|
127 |
+
self.temperature = temperature
|
128 |
+
self.att = nn.Conv1d(
|
129 |
+
in_channels=n_in,
|
130 |
+
out_channels=n_out,
|
131 |
+
kernel_size=1,
|
132 |
+
stride=1,
|
133 |
+
padding=0,
|
134 |
+
bias=True,
|
135 |
+
)
|
136 |
+
self.cla = nn.Conv1d(
|
137 |
+
in_channels=n_in,
|
138 |
+
out_channels=n_out,
|
139 |
+
kernel_size=1,
|
140 |
+
stride=1,
|
141 |
+
padding=0,
|
142 |
+
bias=True,
|
143 |
+
)
|
144 |
+
|
145 |
+
self.bn_att = nn.BatchNorm1d(n_out)
|
146 |
+
self.init_weights()
|
147 |
+
|
148 |
+
def init_weights(self):
|
149 |
+
init_layer(self.att)
|
150 |
+
init_layer(self.cla)
|
151 |
+
init_bn(self.bn_att)
|
152 |
+
|
153 |
+
def forward(self, x):
|
154 |
+
# x: (n_samples, n_in, n_time)
|
155 |
+
norm_att = torch.softmax(torch.clamp(self.att(x), -10, 10), dim=-1)
|
156 |
+
cla = self.nonlinear_transform(self.cla(x))
|
157 |
+
x = torch.sum(norm_att * cla, dim=2)
|
158 |
+
return x, norm_att, cla
|
159 |
+
|
160 |
+
def nonlinear_transform(self, x):
|
161 |
+
if self.activation == "linear":
|
162 |
+
return x
|
163 |
+
elif self.activation == "sigmoid":
|
164 |
+
return torch.sigmoid(x)
|
165 |
+
|
166 |
+
|
167 |
+
class Cnn14(nn.Module):
|
168 |
+
def __init__(
|
169 |
+
self,
|
170 |
+
sample_rate,
|
171 |
+
window_size,
|
172 |
+
hop_size,
|
173 |
+
mel_bins,
|
174 |
+
fmin,
|
175 |
+
fmax,
|
176 |
+
classes_num,
|
177 |
+
enable_fusion=False,
|
178 |
+
fusion_type="None",
|
179 |
+
):
|
180 |
+
super(Cnn14, self).__init__()
|
181 |
+
|
182 |
+
window = "hann"
|
183 |
+
center = True
|
184 |
+
pad_mode = "reflect"
|
185 |
+
ref = 1.0
|
186 |
+
amin = 1e-10
|
187 |
+
top_db = None
|
188 |
+
|
189 |
+
self.enable_fusion = enable_fusion
|
190 |
+
self.fusion_type = fusion_type
|
191 |
+
|
192 |
+
# Spectrogram extractor
|
193 |
+
self.spectrogram_extractor = Spectrogram(
|
194 |
+
n_fft=window_size,
|
195 |
+
hop_length=hop_size,
|
196 |
+
win_length=window_size,
|
197 |
+
window=window,
|
198 |
+
center=center,
|
199 |
+
pad_mode=pad_mode,
|
200 |
+
freeze_parameters=True,
|
201 |
+
)
|
202 |
+
|
203 |
+
# Logmel feature extractor
|
204 |
+
self.logmel_extractor = LogmelFilterBank(
|
205 |
+
sr=sample_rate,
|
206 |
+
n_fft=window_size,
|
207 |
+
n_mels=mel_bins,
|
208 |
+
fmin=fmin,
|
209 |
+
fmax=fmax,
|
210 |
+
ref=ref,
|
211 |
+
amin=amin,
|
212 |
+
top_db=top_db,
|
213 |
+
freeze_parameters=True,
|
214 |
+
)
|
215 |
+
|
216 |
+
# Spec augmenter
|
217 |
+
self.spec_augmenter = SpecAugmentation(
|
218 |
+
time_drop_width=64,
|
219 |
+
time_stripes_num=2,
|
220 |
+
freq_drop_width=8,
|
221 |
+
freq_stripes_num=2,
|
222 |
+
)
|
223 |
+
|
224 |
+
self.bn0 = nn.BatchNorm2d(64)
|
225 |
+
|
226 |
+
if (self.enable_fusion) and (self.fusion_type == "channel_map"):
|
227 |
+
self.conv_block1 = ConvBlock(in_channels=4, out_channels=64)
|
228 |
+
else:
|
229 |
+
self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
|
230 |
+
self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
|
231 |
+
self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
|
232 |
+
self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
|
233 |
+
self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)
|
234 |
+
self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)
|
235 |
+
|
236 |
+
self.fc1 = nn.Linear(2048, 2048, bias=True)
|
237 |
+
self.fc_audioset = nn.Linear(2048, classes_num, bias=True)
|
238 |
+
|
239 |
+
if (self.enable_fusion) and (
|
240 |
+
self.fusion_type in ["daf_1d", "aff_1d", "iaff_1d"]
|
241 |
+
):
|
242 |
+
self.mel_conv1d = nn.Sequential(
|
243 |
+
nn.Conv1d(64, 64, kernel_size=5, stride=3, padding=2),
|
244 |
+
nn.BatchNorm1d(64), # No Relu
|
245 |
+
)
|
246 |
+
if self.fusion_type == "daf_1d":
|
247 |
+
self.fusion_model = DAF()
|
248 |
+
elif self.fusion_type == "aff_1d":
|
249 |
+
self.fusion_model = AFF(channels=64, type="1D")
|
250 |
+
elif self.fusion_type == "iaff_1d":
|
251 |
+
self.fusion_model = iAFF(channels=64, type="1D")
|
252 |
+
|
253 |
+
if (self.enable_fusion) and (
|
254 |
+
self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d"]
|
255 |
+
):
|
256 |
+
self.mel_conv2d = nn.Sequential(
|
257 |
+
nn.Conv2d(1, 64, kernel_size=(5, 5), stride=(6, 2), padding=(2, 2)),
|
258 |
+
nn.BatchNorm2d(64),
|
259 |
+
nn.ReLU(inplace=True),
|
260 |
+
)
|
261 |
+
|
262 |
+
if self.fusion_type == "daf_2d":
|
263 |
+
self.fusion_model = DAF()
|
264 |
+
elif self.fusion_type == "aff_2d":
|
265 |
+
self.fusion_model = AFF(channels=64, type="2D")
|
266 |
+
elif self.fusion_type == "iaff_2d":
|
267 |
+
self.fusion_model = iAFF(channels=64, type="2D")
|
268 |
+
self.init_weight()
|
269 |
+
|
270 |
+
def init_weight(self):
|
271 |
+
init_bn(self.bn0)
|
272 |
+
init_layer(self.fc1)
|
273 |
+
init_layer(self.fc_audioset)
|
274 |
+
|
275 |
+
def forward(self, input, mixup_lambda=None, device=None):
|
276 |
+
"""
|
277 |
+
Input: (batch_size, data_length)"""
|
278 |
+
|
279 |
+
if self.enable_fusion and input["longer"].sum() == 0:
|
280 |
+
# if no audio is longer than 10s, then randomly select one audio to be longer
|
281 |
+
input["longer"][torch.randint(0, input["longer"].shape[0], (1,))] = True
|
282 |
+
|
283 |
+
if not self.enable_fusion:
|
284 |
+
x = self.spectrogram_extractor(
|
285 |
+
input["waveform"].to(device=device, non_blocking=True)
|
286 |
+
) # (batch_size, 1, time_steps, freq_bins)
|
287 |
+
x = self.logmel_extractor(x) # (batch_size, 1, time_steps, mel_bins)
|
288 |
+
|
289 |
+
x = x.transpose(1, 3)
|
290 |
+
x = self.bn0(x)
|
291 |
+
x = x.transpose(1, 3)
|
292 |
+
else:
|
293 |
+
longer_list = input["longer"].to(device=device, non_blocking=True)
|
294 |
+
x = input["mel_fusion"].to(device=device, non_blocking=True)
|
295 |
+
longer_list_idx = torch.where(longer_list)[0]
|
296 |
+
x = x.transpose(1, 3)
|
297 |
+
x = self.bn0(x)
|
298 |
+
x = x.transpose(1, 3)
|
299 |
+
if self.fusion_type in ["daf_1d", "aff_1d", "iaff_1d"]:
|
300 |
+
new_x = x[:, 0:1, :, :].clone().contiguous()
|
301 |
+
# local processing
|
302 |
+
if len(longer_list_idx) > 0:
|
303 |
+
fusion_x_local = x[longer_list_idx, 1:, :, :].clone().contiguous()
|
304 |
+
FB, FC, FT, FF = fusion_x_local.size()
|
305 |
+
fusion_x_local = fusion_x_local.view(FB * FC, FT, FF)
|
306 |
+
fusion_x_local = torch.permute(
|
307 |
+
fusion_x_local, (0, 2, 1)
|
308 |
+
).contiguous()
|
309 |
+
fusion_x_local = self.mel_conv1d(fusion_x_local)
|
310 |
+
fusion_x_local = fusion_x_local.view(
|
311 |
+
FB, FC, FF, fusion_x_local.size(-1)
|
312 |
+
)
|
313 |
+
fusion_x_local = (
|
314 |
+
torch.permute(fusion_x_local, (0, 2, 1, 3))
|
315 |
+
.contiguous()
|
316 |
+
.flatten(2)
|
317 |
+
)
|
318 |
+
if fusion_x_local.size(-1) < FT:
|
319 |
+
fusion_x_local = torch.cat(
|
320 |
+
[
|
321 |
+
fusion_x_local,
|
322 |
+
torch.zeros(
|
323 |
+
(FB, FF, FT - fusion_x_local.size(-1)),
|
324 |
+
device=device,
|
325 |
+
),
|
326 |
+
],
|
327 |
+
dim=-1,
|
328 |
+
)
|
329 |
+
else:
|
330 |
+
fusion_x_local = fusion_x_local[:, :, :FT]
|
331 |
+
# 1D fusion
|
332 |
+
new_x = new_x.squeeze(1).permute((0, 2, 1)).contiguous()
|
333 |
+
new_x[longer_list_idx] = self.fusion_model(
|
334 |
+
new_x[longer_list_idx], fusion_x_local
|
335 |
+
)
|
336 |
+
x = new_x.permute((0, 2, 1)).contiguous()[:, None, :, :]
|
337 |
+
else:
|
338 |
+
x = new_x
|
339 |
+
elif self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d", "channel_map"]:
|
340 |
+
x = x # no change
|
341 |
+
|
342 |
+
if self.training:
|
343 |
+
x = self.spec_augmenter(x)
|
344 |
+
# Mixup on spectrogram
|
345 |
+
if self.training and mixup_lambda is not None:
|
346 |
+
x = do_mixup(x, mixup_lambda)
|
347 |
+
if (self.enable_fusion) and (
|
348 |
+
self.fusion_type in ["daf_2d", "aff_2d", "iaff_2d"]
|
349 |
+
):
|
350 |
+
global_x = x[:, 0:1, :, :]
|
351 |
+
|
352 |
+
# global processing
|
353 |
+
B, C, H, W = global_x.shape
|
354 |
+
global_x = self.conv_block1(global_x, pool_size=(2, 2), pool_type="avg")
|
355 |
+
if len(longer_list_idx) > 0:
|
356 |
+
local_x = x[longer_list_idx, 1:, :, :].contiguous()
|
357 |
+
TH = global_x.size(-2)
|
358 |
+
# local processing
|
359 |
+
B, C, H, W = local_x.shape
|
360 |
+
local_x = local_x.view(B * C, 1, H, W)
|
361 |
+
local_x = self.mel_conv2d(local_x)
|
362 |
+
local_x = local_x.view(
|
363 |
+
B, C, local_x.size(1), local_x.size(2), local_x.size(3)
|
364 |
+
)
|
365 |
+
local_x = local_x.permute((0, 2, 1, 3, 4)).contiguous().flatten(2, 3)
|
366 |
+
TB, TC, _, TW = local_x.size()
|
367 |
+
if local_x.size(-2) < TH:
|
368 |
+
local_x = torch.cat(
|
369 |
+
[
|
370 |
+
local_x,
|
371 |
+
torch.zeros(
|
372 |
+
(TB, TC, TH - local_x.size(-2), TW),
|
373 |
+
device=global_x.device,
|
374 |
+
),
|
375 |
+
],
|
376 |
+
dim=-2,
|
377 |
+
)
|
378 |
+
else:
|
379 |
+
local_x = local_x[:, :, :TH, :]
|
380 |
+
|
381 |
+
global_x[longer_list_idx] = self.fusion_model(
|
382 |
+
global_x[longer_list_idx], local_x
|
383 |
+
)
|
384 |
+
x = global_x
|
385 |
+
else:
|
386 |
+
x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
|
387 |
+
|
388 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
389 |
+
x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
|
390 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
391 |
+
x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
|
392 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
393 |
+
x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
|
394 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
395 |
+
x = self.conv_block5(x, pool_size=(2, 2), pool_type="avg")
|
396 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
397 |
+
x = self.conv_block6(x, pool_size=(1, 1), pool_type="avg")
|
398 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
399 |
+
x = torch.mean(x, dim=3)
|
400 |
+
|
401 |
+
latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
|
402 |
+
latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
|
403 |
+
latent_x = latent_x1 + latent_x2
|
404 |
+
latent_x = latent_x.transpose(1, 2)
|
405 |
+
latent_x = F.relu_(self.fc1(latent_x))
|
406 |
+
latent_output = interpolate(latent_x, 32)
|
407 |
+
|
408 |
+
(x1, _) = torch.max(x, dim=2)
|
409 |
+
x2 = torch.mean(x, dim=2)
|
410 |
+
x = x1 + x2
|
411 |
+
x = F.dropout(x, p=0.5, training=self.training)
|
412 |
+
x = F.relu_(self.fc1(x))
|
413 |
+
embedding = F.dropout(x, p=0.5, training=self.training)
|
414 |
+
clipwise_output = torch.sigmoid(self.fc_audioset(x))
|
415 |
+
|
416 |
+
output_dict = {
|
417 |
+
"clipwise_output": clipwise_output,
|
418 |
+
"embedding": embedding,
|
419 |
+
"fine_grained_embedding": latent_output,
|
420 |
+
}
|
421 |
+
return output_dict
|
422 |
+
|
423 |
+
|
424 |
+
class Cnn6(nn.Module):
|
425 |
+
def __init__(
|
426 |
+
self,
|
427 |
+
sample_rate,
|
428 |
+
window_size,
|
429 |
+
hop_size,
|
430 |
+
mel_bins,
|
431 |
+
fmin,
|
432 |
+
fmax,
|
433 |
+
classes_num,
|
434 |
+
enable_fusion=False,
|
435 |
+
fusion_type="None",
|
436 |
+
):
|
437 |
+
super(Cnn6, self).__init__()
|
438 |
+
|
439 |
+
window = "hann"
|
440 |
+
center = True
|
441 |
+
pad_mode = "reflect"
|
442 |
+
ref = 1.0
|
443 |
+
amin = 1e-10
|
444 |
+
top_db = None
|
445 |
+
|
446 |
+
self.enable_fusion = enable_fusion
|
447 |
+
self.fusion_type = fusion_type
|
448 |
+
|
449 |
+
# Spectrogram extractor
|
450 |
+
self.spectrogram_extractor = Spectrogram(
|
451 |
+
n_fft=window_size,
|
452 |
+
hop_length=hop_size,
|
453 |
+
win_length=window_size,
|
454 |
+
window=window,
|
455 |
+
center=center,
|
456 |
+
pad_mode=pad_mode,
|
457 |
+
freeze_parameters=True,
|
458 |
+
)
|
459 |
+
|
460 |
+
# Logmel feature extractor
|
461 |
+
self.logmel_extractor = LogmelFilterBank(
|
462 |
+
sr=sample_rate,
|
463 |
+
n_fft=window_size,
|
464 |
+
n_mels=mel_bins,
|
465 |
+
fmin=fmin,
|
466 |
+
fmax=fmax,
|
467 |
+
ref=ref,
|
468 |
+
amin=amin,
|
469 |
+
top_db=top_db,
|
470 |
+
freeze_parameters=True,
|
471 |
+
)
|
472 |
+
|
473 |
+
# Spec augmenter
|
474 |
+
self.spec_augmenter = SpecAugmentation(
|
475 |
+
time_drop_width=64,
|
476 |
+
time_stripes_num=2,
|
477 |
+
freq_drop_width=8,
|
478 |
+
freq_stripes_num=2,
|
479 |
+
)
|
480 |
+
|
481 |
+
self.bn0 = nn.BatchNorm2d(64)
|
482 |
+
|
483 |
+
self.conv_block1 = ConvBlock5x5(in_channels=1, out_channels=64)
|
484 |
+
self.conv_block2 = ConvBlock5x5(in_channels=64, out_channels=128)
|
485 |
+
self.conv_block3 = ConvBlock5x5(in_channels=128, out_channels=256)
|
486 |
+
self.conv_block4 = ConvBlock5x5(in_channels=256, out_channels=512)
|
487 |
+
|
488 |
+
self.fc1 = nn.Linear(512, 512, bias=True)
|
489 |
+
self.fc_audioset = nn.Linear(512, classes_num, bias=True)
|
490 |
+
|
491 |
+
self.init_weight()
|
492 |
+
|
493 |
+
def init_weight(self):
|
494 |
+
init_bn(self.bn0)
|
495 |
+
init_layer(self.fc1)
|
496 |
+
init_layer(self.fc_audioset)
|
497 |
+
|
498 |
+
def forward(self, input, mixup_lambda=None, device=None):
|
499 |
+
"""
|
500 |
+
Input: (batch_size, data_length)"""
|
501 |
+
|
502 |
+
x = self.spectrogram_extractor(input) # (batch_size, 1, time_steps, freq_bins)
|
503 |
+
x = self.logmel_extractor(x) # (batch_size, 1, time_steps, mel_bins)
|
504 |
+
|
505 |
+
x = x.transpose(1, 3)
|
506 |
+
x = self.bn0(x)
|
507 |
+
x = x.transpose(1, 3)
|
508 |
+
|
509 |
+
if self.training:
|
510 |
+
x = self.spec_augmenter(x)
|
511 |
+
|
512 |
+
# Mixup on spectrogram
|
513 |
+
if self.training and mixup_lambda is not None:
|
514 |
+
x = do_mixup(x, mixup_lambda)
|
515 |
+
|
516 |
+
x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
|
517 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
518 |
+
x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
|
519 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
520 |
+
x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
|
521 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
522 |
+
x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
|
523 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
524 |
+
x = torch.mean(x, dim=3)
|
525 |
+
|
526 |
+
latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
|
527 |
+
latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
|
528 |
+
latent_x = latent_x1 + latent_x2
|
529 |
+
latent_x = latent_x.transpose(1, 2)
|
530 |
+
latent_x = F.relu_(self.fc1(latent_x))
|
531 |
+
latent_output = interpolate(latent_x, 16)
|
532 |
+
|
533 |
+
(x1, _) = torch.max(x, dim=2)
|
534 |
+
x2 = torch.mean(x, dim=2)
|
535 |
+
x = x1 + x2
|
536 |
+
x = F.dropout(x, p=0.5, training=self.training)
|
537 |
+
x = F.relu_(self.fc1(x))
|
538 |
+
embedding = F.dropout(x, p=0.5, training=self.training)
|
539 |
+
clipwise_output = torch.sigmoid(self.fc_audioset(x))
|
540 |
+
|
541 |
+
output_dict = {
|
542 |
+
"clipwise_output": clipwise_output,
|
543 |
+
"embedding": embedding,
|
544 |
+
"fine_grained_embedding": latent_output,
|
545 |
+
}
|
546 |
+
|
547 |
+
return output_dict
|
548 |
+
|
549 |
+
|
550 |
+
class Cnn10(nn.Module):
|
551 |
+
def __init__(
|
552 |
+
self,
|
553 |
+
sample_rate,
|
554 |
+
window_size,
|
555 |
+
hop_size,
|
556 |
+
mel_bins,
|
557 |
+
fmin,
|
558 |
+
fmax,
|
559 |
+
classes_num,
|
560 |
+
enable_fusion=False,
|
561 |
+
fusion_type="None",
|
562 |
+
):
|
563 |
+
super(Cnn10, self).__init__()
|
564 |
+
|
565 |
+
window = "hann"
|
566 |
+
center = True
|
567 |
+
pad_mode = "reflect"
|
568 |
+
ref = 1.0
|
569 |
+
amin = 1e-10
|
570 |
+
top_db = None
|
571 |
+
|
572 |
+
self.enable_fusion = enable_fusion
|
573 |
+
self.fusion_type = fusion_type
|
574 |
+
|
575 |
+
# Spectrogram extractor
|
576 |
+
self.spectrogram_extractor = Spectrogram(
|
577 |
+
n_fft=window_size,
|
578 |
+
hop_length=hop_size,
|
579 |
+
win_length=window_size,
|
580 |
+
window=window,
|
581 |
+
center=center,
|
582 |
+
pad_mode=pad_mode,
|
583 |
+
freeze_parameters=True,
|
584 |
+
)
|
585 |
+
|
586 |
+
# Logmel feature extractor
|
587 |
+
self.logmel_extractor = LogmelFilterBank(
|
588 |
+
sr=sample_rate,
|
589 |
+
n_fft=window_size,
|
590 |
+
n_mels=mel_bins,
|
591 |
+
fmin=fmin,
|
592 |
+
fmax=fmax,
|
593 |
+
ref=ref,
|
594 |
+
amin=amin,
|
595 |
+
top_db=top_db,
|
596 |
+
freeze_parameters=True,
|
597 |
+
)
|
598 |
+
|
599 |
+
# Spec augmenter
|
600 |
+
self.spec_augmenter = SpecAugmentation(
|
601 |
+
time_drop_width=64,
|
602 |
+
time_stripes_num=2,
|
603 |
+
freq_drop_width=8,
|
604 |
+
freq_stripes_num=2,
|
605 |
+
)
|
606 |
+
|
607 |
+
self.bn0 = nn.BatchNorm2d(64)
|
608 |
+
|
609 |
+
self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
|
610 |
+
self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
|
611 |
+
self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
|
612 |
+
self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
|
613 |
+
self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)
|
614 |
+
|
615 |
+
self.fc1 = nn.Linear(1024, 1024, bias=True)
|
616 |
+
self.fc_audioset = nn.Linear(1024, classes_num, bias=True)
|
617 |
+
|
618 |
+
self.init_weight()
|
619 |
+
|
620 |
+
def init_weight(self):
|
621 |
+
init_bn(self.bn0)
|
622 |
+
init_layer(self.fc1)
|
623 |
+
init_layer(self.fc_audioset)
|
624 |
+
|
625 |
+
def forward(self, input, mixup_lambda=None, device=None):
|
626 |
+
"""
|
627 |
+
Input: (batch_size, data_length)"""
|
628 |
+
|
629 |
+
x = self.spectrogram_extractor(input) # (batch_size, 1, time_steps, freq_bins)
|
630 |
+
x = self.logmel_extractor(x) # (batch_size, 1, time_steps, mel_bins)
|
631 |
+
|
632 |
+
x = x.transpose(1, 3)
|
633 |
+
x = self.bn0(x)
|
634 |
+
x = x.transpose(1, 3)
|
635 |
+
|
636 |
+
if self.training:
|
637 |
+
x = self.spec_augmenter(x)
|
638 |
+
|
639 |
+
# Mixup on spectrogram
|
640 |
+
if self.training and mixup_lambda is not None:
|
641 |
+
x = do_mixup(x, mixup_lambda)
|
642 |
+
|
643 |
+
x = self.conv_block1(x, pool_size=(2, 2), pool_type="avg")
|
644 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
645 |
+
x = self.conv_block2(x, pool_size=(2, 2), pool_type="avg")
|
646 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
647 |
+
x = self.conv_block3(x, pool_size=(2, 2), pool_type="avg")
|
648 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
649 |
+
x = self.conv_block4(x, pool_size=(2, 2), pool_type="avg")
|
650 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
651 |
+
x = self.conv_block5(x, pool_size=(2, 2), pool_type="avg")
|
652 |
+
x = F.dropout(x, p=0.2, training=self.training)
|
653 |
+
x = torch.mean(x, dim=3)
|
654 |
+
|
655 |
+
latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
|
656 |
+
latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
|
657 |
+
latent_x = latent_x1 + latent_x2
|
658 |
+
latent_x = latent_x.transpose(1, 2)
|
659 |
+
latent_x = F.relu_(self.fc1(latent_x))
|
660 |
+
latent_output = interpolate(latent_x, 32)
|
661 |
+
|
662 |
+
(x1, _) = torch.max(x, dim=2)
|
663 |
+
x2 = torch.mean(x, dim=2)
|
664 |
+
x = x1 + x2
|
665 |
+
x = F.dropout(x, p=0.5, training=self.training)
|
666 |
+
x = F.relu_(self.fc1(x))
|
667 |
+
embedding = F.dropout(x, p=0.5, training=self.training)
|
668 |
+
clipwise_output = torch.sigmoid(self.fc_audioset(x))
|
669 |
+
|
670 |
+
output_dict = {
|
671 |
+
"clipwise_output": clipwise_output,
|
672 |
+
"embedding": embedding,
|
673 |
+
"fine_grained_embedding": latent_output,
|
674 |
+
}
|
675 |
+
|
676 |
+
return output_dict
|
677 |
+
|
678 |
+
|
679 |
+
def create_pann_model(audio_cfg, enable_fusion=False, fusion_type="None"):
|
680 |
+
try:
|
681 |
+
ModelProto = eval(audio_cfg.model_name)
|
682 |
+
model = ModelProto(
|
683 |
+
sample_rate=audio_cfg.sample_rate,
|
684 |
+
window_size=audio_cfg.window_size,
|
685 |
+
hop_size=audio_cfg.hop_size,
|
686 |
+
mel_bins=audio_cfg.mel_bins,
|
687 |
+
fmin=audio_cfg.fmin,
|
688 |
+
fmax=audio_cfg.fmax,
|
689 |
+
classes_num=audio_cfg.class_num,
|
690 |
+
enable_fusion=enable_fusion,
|
691 |
+
fusion_type=fusion_type,
|
692 |
+
)
|
693 |
+
return model
|
694 |
+
except:
|
695 |
+
raise RuntimeError(
|
696 |
+
f"Import Model for {audio_cfg.model_name} not found, or the audio cfg parameters are not enough."
|
697 |
+
)
|
audioldm2/clap/open_clip/pretrained.py
ADDED
@@ -0,0 +1,167 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import hashlib
|
2 |
+
import os
|
3 |
+
import urllib
|
4 |
+
import warnings
|
5 |
+
|
6 |
+
from tqdm import tqdm
|
7 |
+
|
8 |
+
_RN50 = dict(
|
9 |
+
openai="https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
|
10 |
+
yfcc15m="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn50-quickgelu-yfcc15m-455df137.pt",
|
11 |
+
cc12m="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn50-quickgelu-cc12m-f000538c.pt",
|
12 |
+
)
|
13 |
+
|
14 |
+
_RN50_quickgelu = dict(
|
15 |
+
openai="https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
|
16 |
+
yfcc15m="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn50-quickgelu-yfcc15m-455df137.pt",
|
17 |
+
cc12m="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn50-quickgelu-cc12m-f000538c.pt",
|
18 |
+
)
|
19 |
+
|
20 |
+
_RN101 = dict(
|
21 |
+
openai="https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
|
22 |
+
yfcc15m="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn101-quickgelu-yfcc15m-3e04b30e.pt",
|
23 |
+
)
|
24 |
+
|
25 |
+
_RN101_quickgelu = dict(
|
26 |
+
openai="https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
|
27 |
+
yfcc15m="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn101-quickgelu-yfcc15m-3e04b30e.pt",
|
28 |
+
)
|
29 |
+
|
30 |
+
_RN50x4 = dict(
|
31 |
+
openai="https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
|
32 |
+
)
|
33 |
+
|
34 |
+
_RN50x16 = dict(
|
35 |
+
openai="https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
|
36 |
+
)
|
37 |
+
|
38 |
+
_RN50x64 = dict(
|
39 |
+
openai="https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
|
40 |
+
)
|
41 |
+
|
42 |
+
_VITB32 = dict(
|
43 |
+
openai="https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
|
44 |
+
laion400m_e31="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt",
|
45 |
+
laion400m_e32="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt",
|
46 |
+
laion400m_avg="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_avg-8a00ab3c.pt",
|
47 |
+
)
|
48 |
+
|
49 |
+
_VITB32_quickgelu = dict(
|
50 |
+
openai="https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
|
51 |
+
laion400m_e31="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt",
|
52 |
+
laion400m_e32="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt",
|
53 |
+
laion400m_avg="https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_avg-8a00ab3c.pt",
|
54 |
+
)
|
55 |
+
|
56 |
+
_VITB16 = dict(
|
57 |
+
openai="https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
|
58 |
+
)
|
59 |
+
|
60 |
+
_VITL14 = dict(
|
61 |
+
openai="https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
|
62 |
+
)
|
63 |
+
|
64 |
+
_PRETRAINED = {
|
65 |
+
"RN50": _RN50,
|
66 |
+
"RN50-quickgelu": _RN50_quickgelu,
|
67 |
+
"RN101": _RN101,
|
68 |
+
"RN101-quickgelu": _RN101_quickgelu,
|
69 |
+
"RN50x4": _RN50x4,
|
70 |
+
"RN50x16": _RN50x16,
|
71 |
+
"ViT-B-32": _VITB32,
|
72 |
+
"ViT-B-32-quickgelu": _VITB32_quickgelu,
|
73 |
+
"ViT-B-16": _VITB16,
|
74 |
+
"ViT-L-14": _VITL14,
|
75 |
+
}
|
76 |
+
|
77 |
+
|
78 |
+
def list_pretrained(as_str: bool = False):
|
79 |
+
"""returns list of pretrained models
|
80 |
+
Returns a tuple (model_name, pretrain_tag) by default or 'name:tag' if as_str == True
|
81 |
+
"""
|
82 |
+
return [
|
83 |
+
":".join([k, t]) if as_str else (k, t)
|
84 |
+
for k in _PRETRAINED.keys()
|
85 |
+
for t in _PRETRAINED[k].keys()
|
86 |
+
]
|
87 |
+
|
88 |
+
|
89 |
+
def list_pretrained_tag_models(tag: str):
|
90 |
+
"""return all models having the specified pretrain tag"""
|
91 |
+
models = []
|
92 |
+
for k in _PRETRAINED.keys():
|
93 |
+
if tag in _PRETRAINED[k]:
|
94 |
+
models.append(k)
|
95 |
+
return models
|
96 |
+
|
97 |
+
|
98 |
+
def list_pretrained_model_tags(model: str):
|
99 |
+
"""return all pretrain tags for the specified model architecture"""
|
100 |
+
tags = []
|
101 |
+
if model in _PRETRAINED:
|
102 |
+
tags.extend(_PRETRAINED[model].keys())
|
103 |
+
return tags
|
104 |
+
|
105 |
+
|
106 |
+
def get_pretrained_url(model: str, tag: str):
|
107 |
+
if model not in _PRETRAINED:
|
108 |
+
return ""
|
109 |
+
model_pretrained = _PRETRAINED[model]
|
110 |
+
if tag not in model_pretrained:
|
111 |
+
return ""
|
112 |
+
return model_pretrained[tag]
|
113 |
+
|
114 |
+
|
115 |
+
def download_pretrained(url: str, root: str = os.path.expanduser("~/.cache/clip")):
|
116 |
+
os.makedirs(root, exist_ok=True)
|
117 |
+
filename = os.path.basename(url)
|
118 |
+
|
119 |
+
if "openaipublic" in url:
|
120 |
+
expected_sha256 = url.split("/")[-2]
|
121 |
+
else:
|
122 |
+
expected_sha256 = ""
|
123 |
+
|
124 |
+
download_target = os.path.join(root, filename)
|
125 |
+
|
126 |
+
if os.path.exists(download_target) and not os.path.isfile(download_target):
|
127 |
+
raise RuntimeError(f"{download_target} exists and is not a regular file")
|
128 |
+
|
129 |
+
if os.path.isfile(download_target):
|
130 |
+
if expected_sha256:
|
131 |
+
if (
|
132 |
+
hashlib.sha256(open(download_target, "rb").read()).hexdigest()
|
133 |
+
== expected_sha256
|
134 |
+
):
|
135 |
+
return download_target
|
136 |
+
else:
|
137 |
+
warnings.warn(
|
138 |
+
f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file"
|
139 |
+
)
|
140 |
+
else:
|
141 |
+
return download_target
|
142 |
+
|
143 |
+
with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
|
144 |
+
with tqdm(
|
145 |
+
total=int(source.info().get("Content-Length")),
|
146 |
+
ncols=80,
|
147 |
+
unit="iB",
|
148 |
+
unit_scale=True,
|
149 |
+
) as loop:
|
150 |
+
while True:
|
151 |
+
buffer = source.read(8192)
|
152 |
+
if not buffer:
|
153 |
+
break
|
154 |
+
|
155 |
+
output.write(buffer)
|
156 |
+
loop.update(len(buffer))
|
157 |
+
|
158 |
+
if (
|
159 |
+
expected_sha256
|
160 |
+
and hashlib.sha256(open(download_target, "rb").read()).hexdigest()
|
161 |
+
!= expected_sha256
|
162 |
+
):
|
163 |
+
raise RuntimeError(
|
164 |
+
f"Model has been downloaded but the SHA256 checksum does not not match"
|
165 |
+
)
|
166 |
+
|
167 |
+
return download_target
|
audioldm2/clap/open_clip/timm_model.py
ADDED
@@ -0,0 +1,112 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
""" timm model adapter
|
2 |
+
|
3 |
+
Wraps timm (https://github.com/rwightman/pytorch-image-models) models for use as a vision tower in CLIP model.
|
4 |
+
"""
|
5 |
+
from collections import OrderedDict
|
6 |
+
|
7 |
+
import torch.nn as nn
|
8 |
+
|
9 |
+
try:
|
10 |
+
import timm
|
11 |
+
from timm.models.layers import Mlp, to_2tuple
|
12 |
+
from timm.models.layers.attention_pool2d import RotAttentionPool2d
|
13 |
+
from timm.models.layers.attention_pool2d import (
|
14 |
+
AttentionPool2d as AbsAttentionPool2d,
|
15 |
+
)
|
16 |
+
except ImportError:
|
17 |
+
timm = None
|
18 |
+
|
19 |
+
from .utils import freeze_batch_norm_2d
|
20 |
+
|
21 |
+
|
22 |
+
class TimmModel(nn.Module):
|
23 |
+
"""timm model adapter
|
24 |
+
# FIXME this adapter is a work in progress, may change in ways that break weight compat
|
25 |
+
"""
|
26 |
+
|
27 |
+
def __init__(
|
28 |
+
self,
|
29 |
+
model_name,
|
30 |
+
embed_dim,
|
31 |
+
image_size=224,
|
32 |
+
pool="avg",
|
33 |
+
proj="linear",
|
34 |
+
drop=0.0,
|
35 |
+
pretrained=False,
|
36 |
+
):
|
37 |
+
super().__init__()
|
38 |
+
if timm is None:
|
39 |
+
raise RuntimeError("Please `pip install timm` to use timm models.")
|
40 |
+
|
41 |
+
self.image_size = to_2tuple(image_size)
|
42 |
+
self.trunk = timm.create_model(model_name, pretrained=pretrained)
|
43 |
+
feat_size = self.trunk.default_cfg.get("pool_size", None)
|
44 |
+
feature_ndim = 1 if not feat_size else 2
|
45 |
+
if pool in ("abs_attn", "rot_attn"):
|
46 |
+
assert feature_ndim == 2
|
47 |
+
# if attn pooling used, remove both classifier and default pool
|
48 |
+
self.trunk.reset_classifier(0, global_pool="")
|
49 |
+
else:
|
50 |
+
# reset global pool if pool config set, otherwise leave as network default
|
51 |
+
reset_kwargs = dict(global_pool=pool) if pool else {}
|
52 |
+
self.trunk.reset_classifier(0, **reset_kwargs)
|
53 |
+
prev_chs = self.trunk.num_features
|
54 |
+
|
55 |
+
head_layers = OrderedDict()
|
56 |
+
if pool == "abs_attn":
|
57 |
+
head_layers["pool"] = AbsAttentionPool2d(
|
58 |
+
prev_chs, feat_size=feat_size, out_features=embed_dim
|
59 |
+
)
|
60 |
+
prev_chs = embed_dim
|
61 |
+
elif pool == "rot_attn":
|
62 |
+
head_layers["pool"] = RotAttentionPool2d(prev_chs, out_features=embed_dim)
|
63 |
+
prev_chs = embed_dim
|
64 |
+
else:
|
65 |
+
assert proj, "projection layer needed if non-attention pooling is used."
|
66 |
+
|
67 |
+
# NOTE attention pool ends with a projection layer, so proj should usually be set to '' if such pooling is used
|
68 |
+
if proj == "linear":
|
69 |
+
head_layers["drop"] = nn.Dropout(drop)
|
70 |
+
head_layers["proj"] = nn.Linear(prev_chs, embed_dim)
|
71 |
+
elif proj == "mlp":
|
72 |
+
head_layers["mlp"] = Mlp(prev_chs, 2 * embed_dim, embed_dim, drop=drop)
|
73 |
+
|
74 |
+
self.head = nn.Sequential(head_layers)
|
75 |
+
|
76 |
+
def lock(self, unlocked_groups=0, freeze_bn_stats=False):
|
77 |
+
"""lock modules
|
78 |
+
Args:
|
79 |
+
unlocked_groups (int): leave last n layer groups unlocked (default: 0)
|
80 |
+
"""
|
81 |
+
if not unlocked_groups:
|
82 |
+
# lock full model
|
83 |
+
for param in self.trunk.parameters():
|
84 |
+
param.requires_grad = False
|
85 |
+
if freeze_bn_stats:
|
86 |
+
freeze_batch_norm_2d(self.trunk)
|
87 |
+
else:
|
88 |
+
# NOTE: partial freeze requires latest timm (master) branch and is subject to change
|
89 |
+
try:
|
90 |
+
# FIXME import here until API stable and in an official release
|
91 |
+
from timm.models.helpers import group_parameters, group_modules
|
92 |
+
except ImportError:
|
93 |
+
raise RuntimeError(
|
94 |
+
"Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`"
|
95 |
+
)
|
96 |
+
matcher = self.trunk.group_matcher()
|
97 |
+
gparams = group_parameters(self.trunk, matcher)
|
98 |
+
max_layer_id = max(gparams.keys())
|
99 |
+
max_layer_id = max_layer_id - unlocked_groups
|
100 |
+
for group_idx in range(max_layer_id + 1):
|
101 |
+
group = gparams[group_idx]
|
102 |
+
for param in group:
|
103 |
+
self.trunk.get_parameter(param).requires_grad = False
|
104 |
+
if freeze_bn_stats:
|
105 |
+
gmodules = group_modules(self.trunk, matcher, reverse=True)
|
106 |
+
gmodules = {k for k, v in gmodules.items() if v <= max_layer_id}
|
107 |
+
freeze_batch_norm_2d(self.trunk, gmodules)
|
108 |
+
|
109 |
+
def forward(self, x):
|
110 |
+
x = self.trunk(x)
|
111 |
+
x = self.head(x)
|
112 |
+
return x
|
audioldm2/clap/open_clip/tokenizer.py
ADDED
@@ -0,0 +1,197 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
""" CLIP tokenizer
|
2 |
+
|
3 |
+
Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
|
4 |
+
"""
|
5 |
+
import gzip
|
6 |
+
import html
|
7 |
+
import os
|
8 |
+
from functools import lru_cache
|
9 |
+
from typing import Union, List
|
10 |
+
|
11 |
+
import ftfy
|
12 |
+
import regex as re
|
13 |
+
import torch
|
14 |
+
|
15 |
+
|
16 |
+
@lru_cache()
|
17 |
+
def default_bpe():
|
18 |
+
return os.path.join(
|
19 |
+
os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz"
|
20 |
+
)
|
21 |
+
|
22 |
+
|
23 |
+
@lru_cache()
|
24 |
+
def bytes_to_unicode():
|
25 |
+
"""
|
26 |
+
Returns list of utf-8 byte and a corresponding list of unicode strings.
|
27 |
+
The reversible bpe codes work on unicode strings.
|
28 |
+
This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
|
29 |
+
When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
|
30 |
+
This is a signficant percentage of your normal, say, 32K bpe vocab.
|
31 |
+
To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
|
32 |
+
And avoids mapping to whitespace/control characters the bpe code barfs on.
|
33 |
+
"""
|
34 |
+
bs = (
|
35 |
+
list(range(ord("!"), ord("~") + 1))
|
36 |
+
+ list(range(ord("¡"), ord("¬") + 1))
|
37 |
+
+ list(range(ord("®"), ord("ÿ") + 1))
|
38 |
+
)
|
39 |
+
cs = bs[:]
|
40 |
+
n = 0
|
41 |
+
for b in range(2**8):
|
42 |
+
if b not in bs:
|
43 |
+
bs.append(b)
|
44 |
+
cs.append(2**8 + n)
|
45 |
+
n += 1
|
46 |
+
cs = [chr(n) for n in cs]
|
47 |
+
return dict(zip(bs, cs))
|
48 |
+
|
49 |
+
|
50 |
+
def get_pairs(word):
|
51 |
+
"""Return set of symbol pairs in a word.
|
52 |
+
Word is represented as tuple of symbols (symbols being variable-length strings).
|
53 |
+
"""
|
54 |
+
pairs = set()
|
55 |
+
prev_char = word[0]
|
56 |
+
for char in word[1:]:
|
57 |
+
pairs.add((prev_char, char))
|
58 |
+
prev_char = char
|
59 |
+
return pairs
|
60 |
+
|
61 |
+
|
62 |
+
def basic_clean(text):
|
63 |
+
text = ftfy.fix_text(text)
|
64 |
+
text = html.unescape(html.unescape(text))
|
65 |
+
return text.strip()
|
66 |
+
|
67 |
+
|
68 |
+
def whitespace_clean(text):
|
69 |
+
text = re.sub(r"\s+", " ", text)
|
70 |
+
text = text.strip()
|
71 |
+
return text
|
72 |
+
|
73 |
+
|
74 |
+
class SimpleTokenizer(object):
|
75 |
+
def __init__(self, bpe_path: str = default_bpe(), special_tokens=None):
|
76 |
+
self.byte_encoder = bytes_to_unicode()
|
77 |
+
self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
|
78 |
+
merges = gzip.open(bpe_path).read().decode("utf-8").split("\n")
|
79 |
+
merges = merges[1 : 49152 - 256 - 2 + 1]
|
80 |
+
merges = [tuple(merge.split()) for merge in merges]
|
81 |
+
vocab = list(bytes_to_unicode().values())
|
82 |
+
vocab = vocab + [v + "</w>" for v in vocab]
|
83 |
+
for merge in merges:
|
84 |
+
vocab.append("".join(merge))
|
85 |
+
if not special_tokens:
|
86 |
+
special_tokens = ["<start_of_text>", "<end_of_text>"]
|
87 |
+
else:
|
88 |
+
special_tokens = ["<start_of_text>", "<end_of_text>"] + special_tokens
|
89 |
+
vocab.extend(special_tokens)
|
90 |
+
self.encoder = dict(zip(vocab, range(len(vocab))))
|
91 |
+
self.decoder = {v: k for k, v in self.encoder.items()}
|
92 |
+
self.bpe_ranks = dict(zip(merges, range(len(merges))))
|
93 |
+
self.cache = {t: t for t in special_tokens}
|
94 |
+
special = "|".join(special_tokens)
|
95 |
+
self.pat = re.compile(
|
96 |
+
special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
|
97 |
+
re.IGNORECASE,
|
98 |
+
)
|
99 |
+
|
100 |
+
self.vocab_size = len(self.encoder)
|
101 |
+
self.all_special_ids = [self.encoder[t] for t in special_tokens]
|
102 |
+
|
103 |
+
def bpe(self, token):
|
104 |
+
if token in self.cache:
|
105 |
+
return self.cache[token]
|
106 |
+
word = tuple(token[:-1]) + (token[-1] + "</w>",)
|
107 |
+
pairs = get_pairs(word)
|
108 |
+
|
109 |
+
if not pairs:
|
110 |
+
return token + "</w>"
|
111 |
+
|
112 |
+
while True:
|
113 |
+
bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
|
114 |
+
if bigram not in self.bpe_ranks:
|
115 |
+
break
|
116 |
+
first, second = bigram
|
117 |
+
new_word = []
|
118 |
+
i = 0
|
119 |
+
while i < len(word):
|
120 |
+
try:
|
121 |
+
j = word.index(first, i)
|
122 |
+
new_word.extend(word[i:j])
|
123 |
+
i = j
|
124 |
+
except:
|
125 |
+
new_word.extend(word[i:])
|
126 |
+
break
|
127 |
+
|
128 |
+
if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
|
129 |
+
new_word.append(first + second)
|
130 |
+
i += 2
|
131 |
+
else:
|
132 |
+
new_word.append(word[i])
|
133 |
+
i += 1
|
134 |
+
new_word = tuple(new_word)
|
135 |
+
word = new_word
|
136 |
+
if len(word) == 1:
|
137 |
+
break
|
138 |
+
else:
|
139 |
+
pairs = get_pairs(word)
|
140 |
+
word = " ".join(word)
|
141 |
+
self.cache[token] = word
|
142 |
+
return word
|
143 |
+
|
144 |
+
def encode(self, text):
|
145 |
+
bpe_tokens = []
|
146 |
+
text = whitespace_clean(basic_clean(text)).lower()
|
147 |
+
for token in re.findall(self.pat, text):
|
148 |
+
token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
|
149 |
+
bpe_tokens.extend(
|
150 |
+
self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" ")
|
151 |
+
)
|
152 |
+
return bpe_tokens
|
153 |
+
|
154 |
+
def decode(self, tokens):
|
155 |
+
text = "".join([self.decoder[token] for token in tokens])
|
156 |
+
text = (
|
157 |
+
bytearray([self.byte_decoder[c] for c in text])
|
158 |
+
.decode("utf-8", errors="replace")
|
159 |
+
.replace("</w>", " ")
|
160 |
+
)
|
161 |
+
return text
|
162 |
+
|
163 |
+
|
164 |
+
_tokenizer = SimpleTokenizer()
|
165 |
+
|
166 |
+
|
167 |
+
def tokenize(
|
168 |
+
texts: Union[str, List[str]], context_length: int = 77
|
169 |
+
) -> torch.LongTensor:
|
170 |
+
"""
|
171 |
+
Returns the tokenized representation of given input string(s)
|
172 |
+
|
173 |
+
Parameters
|
174 |
+
----------
|
175 |
+
texts : Union[str, List[str]]
|
176 |
+
An input string or a list of input strings to tokenize
|
177 |
+
context_length : int
|
178 |
+
The context length to use; all CLIP models use 77 as the context length
|
179 |
+
|
180 |
+
Returns
|
181 |
+
-------
|
182 |
+
A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
|
183 |
+
"""
|
184 |
+
if isinstance(texts, str):
|
185 |
+
texts = [texts]
|
186 |
+
|
187 |
+
sot_token = _tokenizer.encoder["<start_of_text>"]
|
188 |
+
eot_token = _tokenizer.encoder["<end_of_text>"]
|
189 |
+
all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
|
190 |
+
result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
|
191 |
+
|
192 |
+
for i, tokens in enumerate(all_tokens):
|
193 |
+
if len(tokens) > context_length:
|
194 |
+
tokens = tokens[:context_length] # Truncate
|
195 |
+
result[i, : len(tokens)] = torch.tensor(tokens)
|
196 |
+
|
197 |
+
return result
|
audioldm2/clap/open_clip/transform.py
ADDED
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from torchvision.transforms import (
|
2 |
+
Normalize,
|
3 |
+
Compose,
|
4 |
+
RandomResizedCrop,
|
5 |
+
InterpolationMode,
|
6 |
+
ToTensor,
|
7 |
+
Resize,
|
8 |
+
CenterCrop,
|
9 |
+
)
|
10 |
+
|
11 |
+
|
12 |
+
def _convert_to_rgb(image):
|
13 |
+
return image.convert("RGB")
|
14 |
+
|
15 |
+
|
16 |
+
def image_transform(
|
17 |
+
image_size: int,
|
18 |
+
is_train: bool,
|
19 |
+
mean=(0.48145466, 0.4578275, 0.40821073),
|
20 |
+
std=(0.26862954, 0.26130258, 0.27577711),
|
21 |
+
):
|
22 |
+
normalize = Normalize(mean=mean, std=std)
|
23 |
+
if is_train:
|
24 |
+
return Compose(
|
25 |
+
[
|
26 |
+
RandomResizedCrop(
|
27 |
+
image_size,
|
28 |
+
scale=(0.9, 1.0),
|
29 |
+
interpolation=InterpolationMode.BICUBIC,
|
30 |
+
),
|
31 |
+
_convert_to_rgb,
|
32 |
+
ToTensor(),
|
33 |
+
normalize,
|
34 |
+
]
|
35 |
+
)
|
36 |
+
else:
|
37 |
+
return Compose(
|
38 |
+
[
|
39 |
+
Resize(image_size, interpolation=InterpolationMode.BICUBIC),
|
40 |
+
CenterCrop(image_size),
|
41 |
+
_convert_to_rgb,
|
42 |
+
ToTensor(),
|
43 |
+
normalize,
|
44 |
+
]
|
45 |
+
)
|
audioldm2/clap/open_clip/utils.py
ADDED
@@ -0,0 +1,356 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import torch
|
3 |
+
from torch import nn as nn
|
4 |
+
from torchvision.ops.misc import FrozenBatchNorm2d
|
5 |
+
import logging
|
6 |
+
import h5py
|
7 |
+
from tqdm import tqdm
|
8 |
+
import random
|
9 |
+
import json
|
10 |
+
import os
|
11 |
+
import pathlib
|
12 |
+
|
13 |
+
# TODO: (yusong) this not a good place to store those information and does not scale. Need to be fixed later.
|
14 |
+
dataset_split = {
|
15 |
+
"audiocaps": ["train", "valid", "test"],
|
16 |
+
"audioset": ["balanced_train", "unbalanced_train", "eval"],
|
17 |
+
"BBCSoundEffects": ["train", "test"],
|
18 |
+
"Clotho": ["train", "test", "valid"],
|
19 |
+
"free_to_use_sounds": ["train", "test"],
|
20 |
+
"paramount_motion": ["train", "test"],
|
21 |
+
"sonniss_game_effects": ["train", "test"],
|
22 |
+
"wesoundeffects": ["train", "test"],
|
23 |
+
"MACS": ["train", "test"],
|
24 |
+
"freesound": ["train", "test"],
|
25 |
+
"FSD50K": ["train", "test", "valid"],
|
26 |
+
"fsd50k_class_label": ["train", "test", "valid"],
|
27 |
+
"esc50": ["train", "test"],
|
28 |
+
"audiostock": ["train", "test"],
|
29 |
+
"freesound_no_overlap_noesc50": ["train", "test"],
|
30 |
+
"epidemic_sound_effects": ["train", "test"],
|
31 |
+
"VGGSound": ["train", "test"],
|
32 |
+
"urbansound8k_class_label": ["train", "test"],
|
33 |
+
"audioset_t5": ["balanced_train", "unbalanced_train", "eval"],
|
34 |
+
"epidemic_sound_effects_t5": ["train", "test"],
|
35 |
+
"WavText5K": ["train", "test"],
|
36 |
+
"esc50_no_overlap": ["train", "test"],
|
37 |
+
"usd8k_no_overlap": ["train", "test"],
|
38 |
+
"fsd50k_200_class_label": ["train", "test", "valid"],
|
39 |
+
}
|
40 |
+
|
41 |
+
|
42 |
+
def freeze_batch_norm_2d(module, module_match={}, name=""):
|
43 |
+
"""
|
44 |
+
Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. If `module` is
|
45 |
+
itself an instance of either `BatchNorm2d` or `SyncBatchNorm`, it is converted into `FrozenBatchNorm2d` and
|
46 |
+
returned. Otherwise, the module is walked recursively and submodules are converted in place.
|
47 |
+
|
48 |
+
Args:
|
49 |
+
module (torch.nn.Module): Any PyTorch module.
|
50 |
+
module_match (dict): Dictionary of full module names to freeze (all if empty)
|
51 |
+
name (str): Full module name (prefix)
|
52 |
+
|
53 |
+
Returns:
|
54 |
+
torch.nn.Module: Resulting module
|
55 |
+
|
56 |
+
Inspired by https://github.com/pytorch/pytorch/blob/a5895f85be0f10212791145bfedc0261d364f103/torch/nn/modules/batchnorm.py#L762
|
57 |
+
"""
|
58 |
+
res = module
|
59 |
+
is_match = True
|
60 |
+
if module_match:
|
61 |
+
is_match = name in module_match
|
62 |
+
if is_match and isinstance(
|
63 |
+
module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)
|
64 |
+
):
|
65 |
+
res = FrozenBatchNorm2d(module.num_features)
|
66 |
+
res.num_features = module.num_features
|
67 |
+
res.affine = module.affine
|
68 |
+
if module.affine:
|
69 |
+
res.weight.data = module.weight.data.clone().detach()
|
70 |
+
res.bias.data = module.bias.data.clone().detach()
|
71 |
+
res.running_mean.data = module.running_mean.data
|
72 |
+
res.running_var.data = module.running_var.data
|
73 |
+
res.eps = module.eps
|
74 |
+
else:
|
75 |
+
for child_name, child in module.named_children():
|
76 |
+
full_child_name = ".".join([name, child_name]) if name else child_name
|
77 |
+
new_child = freeze_batch_norm_2d(child, module_match, full_child_name)
|
78 |
+
if new_child is not child:
|
79 |
+
res.add_module(child_name, new_child)
|
80 |
+
return res
|
81 |
+
|
82 |
+
|
83 |
+
def exist(dataset_name, dataset_type):
|
84 |
+
"""
|
85 |
+
Check if dataset exists
|
86 |
+
"""
|
87 |
+
if dataset_type in dataset_split[dataset_name]:
|
88 |
+
return True
|
89 |
+
else:
|
90 |
+
return False
|
91 |
+
|
92 |
+
|
93 |
+
def get_tar_path_from_dataset_name(
|
94 |
+
dataset_names, dataset_types, islocal, dataset_path, proportion=1, full_dataset=None
|
95 |
+
):
|
96 |
+
"""
|
97 |
+
Get tar path from dataset name and type
|
98 |
+
"""
|
99 |
+
output = []
|
100 |
+
for n in dataset_names:
|
101 |
+
if full_dataset is not None and n in full_dataset:
|
102 |
+
current_dataset_types = dataset_split[n]
|
103 |
+
else:
|
104 |
+
current_dataset_types = dataset_types
|
105 |
+
for s in current_dataset_types:
|
106 |
+
tmp = []
|
107 |
+
if islocal:
|
108 |
+
sizefilepath_ = f"{dataset_path}/{n}/{s}/sizes.json"
|
109 |
+
if not os.path.exists(sizefilepath_):
|
110 |
+
sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
|
111 |
+
else:
|
112 |
+
sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
|
113 |
+
if not os.path.exists(sizefilepath_):
|
114 |
+
continue
|
115 |
+
sizes = json.load(open(sizefilepath_, "r"))
|
116 |
+
for k in sizes.keys():
|
117 |
+
if islocal:
|
118 |
+
tmp.append(f"{dataset_path}/{n}/{s}/{k}")
|
119 |
+
else:
|
120 |
+
tmp.append(
|
121 |
+
f"pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/{n}/{s}/{k} -"
|
122 |
+
)
|
123 |
+
if proportion != 1:
|
124 |
+
tmp = random.sample(tmp, int(proportion * len(tmp)))
|
125 |
+
output.append(tmp)
|
126 |
+
return sum(output, [])
|
127 |
+
|
128 |
+
|
129 |
+
def get_tar_path_from_txts(txt_path, islocal, proportion=1):
|
130 |
+
"""
|
131 |
+
Get tar path from txt path
|
132 |
+
"""
|
133 |
+
if isinstance(txt_path, (list, tuple)):
|
134 |
+
return sum(
|
135 |
+
[
|
136 |
+
get_tar_path_from_txts(
|
137 |
+
txt_path[i], islocal=islocal, proportion=proportion
|
138 |
+
)
|
139 |
+
for i in range(len(txt_path))
|
140 |
+
],
|
141 |
+
[],
|
142 |
+
)
|
143 |
+
if isinstance(txt_path, str):
|
144 |
+
with open(txt_path) as f:
|
145 |
+
lines = f.readlines()
|
146 |
+
if islocal:
|
147 |
+
lines = [
|
148 |
+
lines[i]
|
149 |
+
.split("\n")[0]
|
150 |
+
.replace("pipe:aws s3 cp s3://s-laion-audio/", "/mnt/audio_clip/")
|
151 |
+
for i in range(len(lines))
|
152 |
+
]
|
153 |
+
else:
|
154 |
+
lines = [
|
155 |
+
lines[i].split("\n")[0].replace(".tar", ".tar -")
|
156 |
+
for i in range(len(lines))
|
157 |
+
]
|
158 |
+
if proportion != 1:
|
159 |
+
print("Sampling tars with proportion of {}".format(proportion))
|
160 |
+
lines = random.sample(lines, int(proportion * len(lines)))
|
161 |
+
return lines
|
162 |
+
|
163 |
+
|
164 |
+
def get_mix_lambda(mixup_alpha, batch_size):
|
165 |
+
mixup_lambdas = [
|
166 |
+
np.random.beta(mixup_alpha, mixup_alpha, 1)[0] for _ in range(batch_size)
|
167 |
+
]
|
168 |
+
return np.array(mixup_lambdas).astype(np.float32)
|
169 |
+
|
170 |
+
|
171 |
+
def do_mixup(x, mixup_lambda):
|
172 |
+
"""
|
173 |
+
Args:
|
174 |
+
x: (batch_size , ...)
|
175 |
+
mixup_lambda: (batch_size,)
|
176 |
+
Returns:
|
177 |
+
out: (batch_size, ...)
|
178 |
+
"""
|
179 |
+
out = (
|
180 |
+
x.transpose(0, -1) * mixup_lambda
|
181 |
+
+ torch.flip(x, dims=[0]).transpose(0, -1) * (1 - mixup_lambda)
|
182 |
+
).transpose(0, -1)
|
183 |
+
return out
|
184 |
+
|
185 |
+
|
186 |
+
def interpolate(x, ratio):
|
187 |
+
"""Interpolate data in time domain. This is used to compensate the
|
188 |
+
resolution reduction in downsampling of a CNN.
|
189 |
+
|
190 |
+
Args:
|
191 |
+
x: (batch_size, time_steps, classes_num)
|
192 |
+
ratio: int, ratio to interpolate
|
193 |
+
Returns:
|
194 |
+
upsampled: (batch_size, time_steps * ratio, classes_num)
|
195 |
+
"""
|
196 |
+
(batch_size, time_steps, classes_num) = x.shape
|
197 |
+
upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)
|
198 |
+
upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)
|
199 |
+
return upsampled
|
200 |
+
|
201 |
+
|
202 |
+
def pad_framewise_output(framewise_output, frames_num):
|
203 |
+
"""Pad framewise_output to the same length as input frames. The pad value
|
204 |
+
is the same as the value of the last frame.
|
205 |
+
Args:
|
206 |
+
framewise_output: (batch_size, frames_num, classes_num)
|
207 |
+
frames_num: int, number of frames to pad
|
208 |
+
Outputs:
|
209 |
+
output: (batch_size, frames_num, classes_num)
|
210 |
+
"""
|
211 |
+
pad = framewise_output[:, -1:, :].repeat(
|
212 |
+
1, frames_num - framewise_output.shape[1], 1
|
213 |
+
)
|
214 |
+
"""tensor for padding"""
|
215 |
+
|
216 |
+
output = torch.cat((framewise_output, pad), dim=1)
|
217 |
+
"""(batch_size, frames_num, classes_num)"""
|
218 |
+
|
219 |
+
|
220 |
+
def process_ipc(index_path, classes_num, filename):
|
221 |
+
# load data
|
222 |
+
logging.info("Load Data...............")
|
223 |
+
ipc = [[] for _ in range(classes_num)]
|
224 |
+
with h5py.File(index_path, "r") as f:
|
225 |
+
for i in tqdm(range(len(f["target"]))):
|
226 |
+
t_class = np.where(f["target"][i])[0]
|
227 |
+
for t in t_class:
|
228 |
+
ipc[t].append(i)
|
229 |
+
print(ipc)
|
230 |
+
np.save(filename, ipc)
|
231 |
+
logging.info("Load Data Succeed...............")
|
232 |
+
|
233 |
+
|
234 |
+
def save_to_dict(s, o_={}):
|
235 |
+
sp = s.split(": ")
|
236 |
+
o_.update({sp[0]: float(sp[1])})
|
237 |
+
return o_
|
238 |
+
|
239 |
+
|
240 |
+
def get_data_from_log(txt_path):
|
241 |
+
"""
|
242 |
+
Output dictionary from out.txt log file
|
243 |
+
"""
|
244 |
+
with open(txt_path) as f:
|
245 |
+
lines = f.readlines()
|
246 |
+
val_data = {}
|
247 |
+
train_data = {}
|
248 |
+
train_losses = []
|
249 |
+
train_losses_epoch = []
|
250 |
+
for i in range(len(lines)):
|
251 |
+
if "| INFO |" in lines[i]:
|
252 |
+
if "Eval Epoch" in lines[i]:
|
253 |
+
if "val_loss" in lines[i]:
|
254 |
+
# float(regex.sub("", lines[310].split(" ")[-1]).replace(" ", ""))
|
255 |
+
line = lines[i].split("Eval Epoch: ")[-1]
|
256 |
+
num_epoch = int(line.split(" ")[0].split(" ")[0])
|
257 |
+
d = {
|
258 |
+
line.split(" ")[0]
|
259 |
+
.split(" ")[1]
|
260 |
+
.replace(":", ""): float(line.split(" ")[0].split(" ")[-1])
|
261 |
+
}
|
262 |
+
for i in range(1, len(line.split(" "))):
|
263 |
+
d = save_to_dict(line.split(" ")[i], d)
|
264 |
+
val_data[num_epoch] = d
|
265 |
+
elif "Train Epoch" in lines[i]:
|
266 |
+
num_epoch = int(lines[i].split("Train Epoch: ")[1][0])
|
267 |
+
loss = float(lines[i].split("Loss: ")[-1].split(" (")[0])
|
268 |
+
train_losses.append(loss)
|
269 |
+
train_losses_epoch.append(num_epoch)
|
270 |
+
for i in range(len(train_losses)):
|
271 |
+
train_data[i] = {
|
272 |
+
"num_epoch": train_losses_epoch[i],
|
273 |
+
"train_loss": train_losses[i],
|
274 |
+
}
|
275 |
+
return train_data, val_data
|
276 |
+
|
277 |
+
|
278 |
+
def save_p(obj, filename):
|
279 |
+
import pickle
|
280 |
+
|
281 |
+
try:
|
282 |
+
from deepdiff import DeepDiff
|
283 |
+
except:
|
284 |
+
os.system("pip install deepdiff")
|
285 |
+
from deepdiff import DeepDiff
|
286 |
+
with open(filename, "wb") as file:
|
287 |
+
pickle.dump(obj, file, protocol=pickle.HIGHEST_PROTOCOL) # highest protocol
|
288 |
+
with open(filename, "rb") as file:
|
289 |
+
z = pickle.load(file)
|
290 |
+
assert (
|
291 |
+
DeepDiff(obj, z, ignore_string_case=True) == {}
|
292 |
+
), "there is something wrong with the saving process"
|
293 |
+
return
|
294 |
+
|
295 |
+
|
296 |
+
def load_p(filename):
|
297 |
+
import pickle
|
298 |
+
|
299 |
+
with open(filename, "rb") as file:
|
300 |
+
z = pickle.load(file)
|
301 |
+
return z
|
302 |
+
|
303 |
+
|
304 |
+
def save_json(data, name="data.json"):
|
305 |
+
import json
|
306 |
+
|
307 |
+
with open(name, "w") as fp:
|
308 |
+
json.dump(data, fp)
|
309 |
+
return
|
310 |
+
|
311 |
+
|
312 |
+
def load_json(name):
|
313 |
+
import json
|
314 |
+
|
315 |
+
with open(name, "r") as fp:
|
316 |
+
data = json.load(fp)
|
317 |
+
return data
|
318 |
+
|
319 |
+
|
320 |
+
def load_class_label(path):
|
321 |
+
# https://stackoverflow.com/questions/48004243/how-to-share-large-read-only-dictionary-list-across-processes-in-multiprocessing
|
322 |
+
# https://stackoverflow.com/questions/45693949/storing-strings-in-a-multiprocessing-sharedctypes-array
|
323 |
+
out = None
|
324 |
+
if path is not None:
|
325 |
+
if pathlib.Path(path).suffix in [".pkl", ".pickle"]:
|
326 |
+
out = load_p(path)
|
327 |
+
elif pathlib.Path(path).suffix in [".json", ".txt"]:
|
328 |
+
out = load_json(path)
|
329 |
+
elif pathlib.Path(path).suffix in [".npy", ".npz"]:
|
330 |
+
out = np.load(path)
|
331 |
+
elif pathlib.Path(path).suffix in [".csv"]:
|
332 |
+
import pandas as pd
|
333 |
+
|
334 |
+
out = pd.read_csv(path)
|
335 |
+
return out
|
336 |
+
# if out is None:
|
337 |
+
# return None
|
338 |
+
# else:
|
339 |
+
# key = Array(c_wchar, '\n'.join(list(out.keys())), lock=False)
|
340 |
+
# val = Array('i', out.values(), lock=False)
|
341 |
+
# return (key, val)
|
342 |
+
|
343 |
+
|
344 |
+
from torch import optim
|
345 |
+
|
346 |
+
|
347 |
+
def get_optimizer(params, lr, betas, eps, momentum, optimizer_name):
|
348 |
+
if optimizer_name.lower() == "adamw":
|
349 |
+
optimizer = optim.AdamW(params, lr=lr, betas=betas, eps=eps)
|
350 |
+
elif optimizer_name.lower() == "sgd":
|
351 |
+
optimizer = optim.SGD(params, lr=lr, momentum=momentum)
|
352 |
+
elif optimizer_name.lower() == "adam":
|
353 |
+
optimizer = optim.Adam(params, lr=lr, betas=betas, eps=eps)
|
354 |
+
else:
|
355 |
+
raise ValueError("optimizer name is not correct")
|
356 |
+
return optimizer
|
audioldm2/clap/training/__init__.py
ADDED
File without changes
|
audioldm2/clap/training/audioset_textmap.npy
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:bada103070d92f9eadd33e1b4f45ec8583f59080ef218c966b43294bd4c86d5b
|
3 |
+
size 84448
|
audioldm2/clap/training/bpe_simple_vocab_16e6.txt.gz
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
|
3 |
+
size 1356917
|
audioldm2/clap/training/data.py
ADDED
@@ -0,0 +1,865 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import json
|
2 |
+
import logging
|
3 |
+
import os
|
4 |
+
import random
|
5 |
+
import h5py
|
6 |
+
from dataclasses import dataclass
|
7 |
+
import numpy as np
|
8 |
+
import pandas as pd
|
9 |
+
import torch
|
10 |
+
import torchvision.datasets as datasets
|
11 |
+
from PIL import Image
|
12 |
+
from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
|
13 |
+
from torch.utils.data.distributed import DistributedSampler
|
14 |
+
import soundfile as sf
|
15 |
+
import io
|
16 |
+
from pathlib import Path
|
17 |
+
# import wget
|
18 |
+
|
19 |
+
from audioldm2.clap.open_clip.utils import get_tar_path_from_dataset_name
|
20 |
+
from audioldm2.clap.open_clip.utils import load_class_label
|
21 |
+
|
22 |
+
try:
|
23 |
+
import horovod.torch as hvd
|
24 |
+
except ImportError:
|
25 |
+
hvd = None
|
26 |
+
|
27 |
+
try:
|
28 |
+
import torchaudio
|
29 |
+
except ImportError:
|
30 |
+
torchaudio = None
|
31 |
+
|
32 |
+
from audioldm2.clap.open_clip import tokenize
|
33 |
+
|
34 |
+
|
35 |
+
def tokenizer(text):
|
36 |
+
return tokenize(text).squeeze(0)
|
37 |
+
|
38 |
+
|
39 |
+
from transformers import RobertaTokenizer
|
40 |
+
|
41 |
+
tokenize = RobertaTokenizer.from_pretrained("roberta-base")
|
42 |
+
|
43 |
+
|
44 |
+
def tokenizer(text):
|
45 |
+
result = tokenize(
|
46 |
+
text,
|
47 |
+
padding="max_length",
|
48 |
+
truncation=True,
|
49 |
+
max_length=77,
|
50 |
+
return_tensors="pt",
|
51 |
+
)
|
52 |
+
return {k: v.squeeze(0) for k, v in result.items()}
|
53 |
+
|
54 |
+
|
55 |
+
# initizlied the audioset map
|
56 |
+
_AUDIOSET_MAP_PATH = os.path.join(Path(__file__).parent, "audioset_textmap.npy")
|
57 |
+
_AUDIOSET_MAP = np.load(_AUDIOSET_MAP_PATH, allow_pickle=True)
|
58 |
+
|
59 |
+
|
60 |
+
def int16_to_float32(x):
|
61 |
+
return (x / 32767.0).astype(np.float32)
|
62 |
+
|
63 |
+
|
64 |
+
def float32_to_int16(x):
|
65 |
+
x = np.clip(x, a_min=-1.0, a_max=1.0)
|
66 |
+
return (x * 32767.0).astype(np.int16)
|
67 |
+
|
68 |
+
|
69 |
+
# For Toy Dataset
|
70 |
+
class ToyDataset(Dataset):
|
71 |
+
def __init__(self, index_path, ipc, config, eval_mode=False):
|
72 |
+
"""Toy Dataset for testing the audioset input with text labels
|
73 |
+
Parameters
|
74 |
+
----------
|
75 |
+
index_path: str
|
76 |
+
the link to the h5 file of each audio
|
77 |
+
idc: str
|
78 |
+
the link to the npy file, the number of samples in each class
|
79 |
+
config: dict
|
80 |
+
the audio cfg file
|
81 |
+
eval_model (bool): to indicate if the dataset is a testing dataset
|
82 |
+
"""
|
83 |
+
self.audio_cfg = config["audio_cfg"]
|
84 |
+
self.text_cfg = config["text_cfg"]
|
85 |
+
self.fp = h5py.File(index_path, "r")
|
86 |
+
self.ipc = np.load(ipc, allow_pickle=True)
|
87 |
+
self.total_size = len(self.fp["audio_name"])
|
88 |
+
self.classes_num = self.audio_cfg["class_num"]
|
89 |
+
self.eval_mode = eval_mode
|
90 |
+
|
91 |
+
if not eval_mode:
|
92 |
+
self.generate_queue()
|
93 |
+
else:
|
94 |
+
self.queue = []
|
95 |
+
for i in range(self.total_size):
|
96 |
+
target = self.fp["target"][i]
|
97 |
+
if np.sum(target) > 0:
|
98 |
+
self.queue.append(i)
|
99 |
+
self.total_size = len(self.queue)
|
100 |
+
logging.info("total dataset size: %d" % (self.total_size))
|
101 |
+
logging.info("class num: %d" % (self.classes_num))
|
102 |
+
|
103 |
+
def time_shifting(self, x):
|
104 |
+
frame_num = len(x)
|
105 |
+
shift_len = random.randint(0, frame_num - 1)
|
106 |
+
new_sample = np.concatenate([x[shift_len:], x[:shift_len]], axis=0)
|
107 |
+
return new_sample
|
108 |
+
|
109 |
+
def generate_queue(self):
|
110 |
+
self.queue = []
|
111 |
+
while len(self.queue) < self.total_size:
|
112 |
+
class_set = [*range(self.classes_num)]
|
113 |
+
random.shuffle(class_set)
|
114 |
+
self.queue += [
|
115 |
+
self.ipc[d][random.randint(0, len(self.ipc[d]) - 1)] for d in class_set
|
116 |
+
]
|
117 |
+
self.queue = self.queue[: self.total_size]
|
118 |
+
|
119 |
+
logging.info("queue regenerated:%s" % (self.queue[-5:]))
|
120 |
+
|
121 |
+
def crop_wav(self, x):
|
122 |
+
crop_size = self.audio_cfg["crop_size"]
|
123 |
+
crop_pos = random.randint(0, len(x) - crop_size - 1)
|
124 |
+
return x[crop_pos : crop_pos + crop_size]
|
125 |
+
|
126 |
+
def prompt_text(self, target):
|
127 |
+
events = _AUDIOSET_MAP[np.where(target > 0)]
|
128 |
+
event_text = "The sounds of " + ", ".join(events[:-1]) + " and " + events[-1]
|
129 |
+
text = tokenize(event_text)[0]
|
130 |
+
return text
|
131 |
+
|
132 |
+
def __getitem__(self, index):
|
133 |
+
"""Load waveform, text, and target of an audio clip
|
134 |
+
|
135 |
+
Parameters
|
136 |
+
----------
|
137 |
+
index: int
|
138 |
+
the index number
|
139 |
+
Return
|
140 |
+
------
|
141 |
+
output: dict {
|
142 |
+
"hdf5_path": str,
|
143 |
+
"index_in_hdf5": int,
|
144 |
+
"audio_name": str,
|
145 |
+
"waveform": list (audio_length,),
|
146 |
+
"target": list (class_num, ),
|
147 |
+
"text": torch.tensor (context_length,)
|
148 |
+
}
|
149 |
+
the output dictionary
|
150 |
+
"""
|
151 |
+
s_index = self.queue[index]
|
152 |
+
|
153 |
+
audio_name = self.fp["audio_name"][s_index].decode()
|
154 |
+
# Hardcode here CHANGE
|
155 |
+
hdf5_path = (
|
156 |
+
self.fp["hdf5_path"][s_index]
|
157 |
+
.decode()
|
158 |
+
.replace(
|
159 |
+
"../workspace",
|
160 |
+
"/home/la/kechen/Research/ke_zsasp/workspace",
|
161 |
+
)
|
162 |
+
)
|
163 |
+
r_idx = self.fp["index_in_hdf5"][s_index]
|
164 |
+
target = self.fp["target"][s_index].astype(np.float32)
|
165 |
+
text = self.prompt_text(target)
|
166 |
+
with h5py.File(hdf5_path, "r") as f:
|
167 |
+
waveform = int16_to_float32(f["waveform"][r_idx])[
|
168 |
+
: self.audio_cfg["clip_samples"]
|
169 |
+
]
|
170 |
+
assert (
|
171 |
+
len(waveform) == self.audio_cfg["clip_samples"]
|
172 |
+
), "The sample length is not match"
|
173 |
+
# Time shift
|
174 |
+
# if (self.config.enable_time_shift) and (not self.eval_mode):
|
175 |
+
# waveform = self.time_shifting(waveform)
|
176 |
+
# # Label Enhance
|
177 |
+
# if (self.config.crop_size is not None) and (not self.eval_mode):
|
178 |
+
# waveform = self.crop_wav(waveform)
|
179 |
+
# # the label enhance rate is fixed 0.5
|
180 |
+
# if (self.config.enable_label_enhance) and (not self.eval_mode) and random.random() < 0.5:
|
181 |
+
# kidx = np.where(target)[0]
|
182 |
+
# for k in kidx:
|
183 |
+
# for add_key in self.class_map[k][1]:
|
184 |
+
# target[add_key] = 1.0
|
185 |
+
# if len(self.class_map[k][2]) > 0:
|
186 |
+
# add_key = random.choice(self.class_map[k][2])
|
187 |
+
# target[add_key] = 1.0
|
188 |
+
|
189 |
+
# missing the text input
|
190 |
+
mel_spec = get_mel(torch.from_numpy(waveform), self.audio_cfg)[None, :, :]
|
191 |
+
mel_spec = (
|
192 |
+
torch.cat(
|
193 |
+
[mel_spec, mel_spec.clone(), mel_spec.clone(), mel_spec.clone()], dim=0
|
194 |
+
)
|
195 |
+
.cpu()
|
196 |
+
.numpy()
|
197 |
+
)
|
198 |
+
longer = random.choice([True, False])
|
199 |
+
if longer == False:
|
200 |
+
mel_spec[1:, :, :] = 0.0
|
201 |
+
data_dict = {
|
202 |
+
"hdf5_path": hdf5_path,
|
203 |
+
"index_in_hdf5": r_idx,
|
204 |
+
"audio_name": audio_name,
|
205 |
+
"waveform": waveform,
|
206 |
+
"class_label": target,
|
207 |
+
"text": text,
|
208 |
+
"longer": longer,
|
209 |
+
"mel_fusion": mel_spec,
|
210 |
+
}
|
211 |
+
return data_dict
|
212 |
+
|
213 |
+
def __len__(self):
|
214 |
+
return self.total_size
|
215 |
+
|
216 |
+
|
217 |
+
class CsvDataset(Dataset):
|
218 |
+
def __init__(self, input_filename, transforms, img_key, caption_key, sep="\t"):
|
219 |
+
logging.debug(f"Loading csv data from {input_filename}.")
|
220 |
+
df = pd.read_csv(input_filename, sep=sep)
|
221 |
+
|
222 |
+
self.images = df[img_key].tolist()
|
223 |
+
self.captions = df[caption_key].tolist()
|
224 |
+
self.transforms = transforms
|
225 |
+
logging.debug("Done loading data.")
|
226 |
+
|
227 |
+
def __len__(self):
|
228 |
+
return len(self.captions)
|
229 |
+
|
230 |
+
def __getitem__(self, idx):
|
231 |
+
images = self.transforms(Image.open(str(self.images[idx])))
|
232 |
+
texts = tokenize([str(self.captions[idx])])[0]
|
233 |
+
return images, texts
|
234 |
+
|
235 |
+
|
236 |
+
@dataclass
|
237 |
+
class DataInfo:
|
238 |
+
dataloader: DataLoader
|
239 |
+
sampler: DistributedSampler
|
240 |
+
|
241 |
+
|
242 |
+
def preprocess_txt(text):
|
243 |
+
return tokenize([str(text)])[0]
|
244 |
+
|
245 |
+
|
246 |
+
# def get_dataset_size(shards, sizefilepath_=None, is_local=True):
|
247 |
+
# if isinstance(shards, list):
|
248 |
+
# size_list = []
|
249 |
+
# for s in shards:
|
250 |
+
# size_list.append(
|
251 |
+
# get_dataset_size(s, sizefilepath_=sizefilepath_, is_local=is_local)[0]
|
252 |
+
# )
|
253 |
+
# else:
|
254 |
+
# if not is_local:
|
255 |
+
# for n in dataset_split.keys():
|
256 |
+
# if n in shards.split("/"):
|
257 |
+
# break
|
258 |
+
# for s in dataset_split[n]:
|
259 |
+
# if s in shards.split("/"):
|
260 |
+
# break
|
261 |
+
# sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
|
262 |
+
# shards_list = list(braceexpand.braceexpand(shards))
|
263 |
+
# dir_path = os.path.dirname(shards)
|
264 |
+
# if sizefilepath_ is not None:
|
265 |
+
# sizes = json.load(open(sizefilepath_, "r"))
|
266 |
+
# total_size = sum(
|
267 |
+
# [
|
268 |
+
# int(sizes[os.path.basename(shard.replace(".tar -", ".tar"))])
|
269 |
+
# for shard in shards_list
|
270 |
+
# ]
|
271 |
+
# )
|
272 |
+
# else:
|
273 |
+
# sizes_filename = os.path.join(dir_path, "sizes.json")
|
274 |
+
# len_filename = os.path.join(dir_path, "__len__")
|
275 |
+
# if os.path.exists(sizes_filename):
|
276 |
+
# sizes = json.load(open(sizes_filename, "r"))
|
277 |
+
# total_size = sum(
|
278 |
+
# [int(sizes[os.path.basename(shard)]) for shard in shards_list]
|
279 |
+
# )
|
280 |
+
# elif os.path.exists(len_filename):
|
281 |
+
# # FIXME this used to be eval(open(...)) but that seemed rather unsafe
|
282 |
+
# total_size = ast.literal_eval(open(len_filename, "r").read())
|
283 |
+
# else:
|
284 |
+
# raise Exception(
|
285 |
+
# "Cannot find sizes file for dataset. Please specify the path to the file."
|
286 |
+
# )
|
287 |
+
# # total_size = None # num samples undefined
|
288 |
+
# # some common dataset sizes (at time of authors last download)
|
289 |
+
# # cc3m-train: 2905954
|
290 |
+
# # cc12m: 10968539
|
291 |
+
# # LAION-400m: 407332084
|
292 |
+
# num_shards = len(shards_list)
|
293 |
+
# if isinstance(shards, list):
|
294 |
+
# return sum(size_list), len(shards)
|
295 |
+
# else:
|
296 |
+
# return total_size, num_shards
|
297 |
+
|
298 |
+
|
299 |
+
def get_imagenet(args, preprocess_fns, split):
|
300 |
+
assert split in ["train", "val", "v2"]
|
301 |
+
is_train = split == "train"
|
302 |
+
preprocess_train, preprocess_val = preprocess_fns
|
303 |
+
|
304 |
+
if split == "v2":
|
305 |
+
from imagenetv2_pytorch import ImageNetV2Dataset
|
306 |
+
|
307 |
+
dataset = ImageNetV2Dataset(location=args.imagenet_v2, transform=preprocess_val)
|
308 |
+
else:
|
309 |
+
if is_train:
|
310 |
+
data_path = args.imagenet_train
|
311 |
+
preprocess_fn = preprocess_train
|
312 |
+
else:
|
313 |
+
data_path = args.imagenet_val
|
314 |
+
preprocess_fn = preprocess_val
|
315 |
+
assert data_path
|
316 |
+
|
317 |
+
dataset = datasets.ImageFolder(data_path, transform=preprocess_fn)
|
318 |
+
|
319 |
+
if is_train:
|
320 |
+
idxs = np.zeros(len(dataset.targets))
|
321 |
+
target_array = np.array(dataset.targets)
|
322 |
+
k = 50
|
323 |
+
for c in range(1000):
|
324 |
+
m = target_array == c
|
325 |
+
n = len(idxs[m])
|
326 |
+
arr = np.zeros(n)
|
327 |
+
arr[:k] = 1
|
328 |
+
np.random.shuffle(arr)
|
329 |
+
idxs[m] = arr
|
330 |
+
|
331 |
+
idxs = idxs.astype("int")
|
332 |
+
sampler = SubsetRandomSampler(np.where(idxs)[0])
|
333 |
+
else:
|
334 |
+
sampler = None
|
335 |
+
|
336 |
+
dataloader = torch.utils.data.DataLoader(
|
337 |
+
dataset,
|
338 |
+
batch_size=args.batch_size,
|
339 |
+
num_workers=args.workers,
|
340 |
+
sampler=sampler,
|
341 |
+
)
|
342 |
+
|
343 |
+
return DataInfo(dataloader, sampler)
|
344 |
+
|
345 |
+
|
346 |
+
def count_samples(dataloader):
|
347 |
+
os.environ["WDS_EPOCH"] = "0"
|
348 |
+
n_elements, n_batches = 0, 0
|
349 |
+
for images, texts in dataloader:
|
350 |
+
n_batches += 1
|
351 |
+
n_elements += len(images)
|
352 |
+
assert len(images) == len(texts)
|
353 |
+
return n_elements, n_batches
|
354 |
+
|
355 |
+
|
356 |
+
def filter_no_caption(sample):
|
357 |
+
return "txt" in sample
|
358 |
+
|
359 |
+
|
360 |
+
def log_and_continue(exn):
|
361 |
+
"""Call in an exception handler to ignore any exception, isssue a warning, and continue."""
|
362 |
+
logging.warning(f"Handling webdataset error ({repr(exn)}). Ignoring.")
|
363 |
+
return True
|
364 |
+
|
365 |
+
|
366 |
+
_SHARD_SHUFFLE_SIZE = 2000
|
367 |
+
_SHARD_SHUFFLE_INITIAL = 500
|
368 |
+
_SAMPLE_SHUFFLE_SIZE = 5000
|
369 |
+
_SAMPLE_SHUFFLE_INITIAL = 1000
|
370 |
+
|
371 |
+
|
372 |
+
# def sample_prop(sizefile, inputs, proportion, is_local=True):
|
373 |
+
# """
|
374 |
+
# Sample a proportion of the data.
|
375 |
+
# """
|
376 |
+
# file_path_dict = {
|
377 |
+
# os.path.split(inputs[i])[1]: os.path.split(inputs[i])[0]
|
378 |
+
# for i in range(len(inputs))
|
379 |
+
# }
|
380 |
+
# sampled_filepath_dict = {}
|
381 |
+
# sampled_size_dict = {}
|
382 |
+
# if not is_local:
|
383 |
+
# if os.path.exists("sizes.json"):
|
384 |
+
# os.remove("sizes.json")
|
385 |
+
# wget.download(sizefile, "sizes.json")
|
386 |
+
# sizefile = "sizes.json"
|
387 |
+
# with open(sizefile, "r", encoding="UTF-8") as f:
|
388 |
+
# load_dict = json.load(f)
|
389 |
+
# L = int(len(file_path_dict) * proportion)
|
390 |
+
# subkeys = random.sample(file_path_dict.keys(), L)
|
391 |
+
# for k in subkeys:
|
392 |
+
# sampled_size_dict[k] = load_dict[k]
|
393 |
+
# sampled_filepath_dict[k] = file_path_dict[k]
|
394 |
+
# return (
|
395 |
+
# sum(sampled_size_dict.values()),
|
396 |
+
# L,
|
397 |
+
# [os.path.join(v, k) for k, v in sampled_filepath_dict.items()],
|
398 |
+
# sampled_size_dict,
|
399 |
+
# )
|
400 |
+
|
401 |
+
|
402 |
+
def get_mel(audio_data, audio_cfg):
|
403 |
+
# mel shape: (n_mels, T)
|
404 |
+
mel = torchaudio.transforms.MelSpectrogram(
|
405 |
+
sample_rate=audio_cfg["sample_rate"],
|
406 |
+
n_fft=audio_cfg["window_size"],
|
407 |
+
win_length=audio_cfg["window_size"],
|
408 |
+
hop_length=audio_cfg["hop_size"],
|
409 |
+
center=True,
|
410 |
+
pad_mode="reflect",
|
411 |
+
power=2.0,
|
412 |
+
norm=None,
|
413 |
+
onesided=True,
|
414 |
+
n_mels=64,
|
415 |
+
f_min=audio_cfg["fmin"],
|
416 |
+
f_max=audio_cfg["fmax"],
|
417 |
+
).to(audio_data.device)
|
418 |
+
mel = mel(audio_data)
|
419 |
+
# we use log mel spectrogram as input
|
420 |
+
mel = torchaudio.transforms.AmplitudeToDB(top_db=None)(mel)
|
421 |
+
return mel.T # (T, n_mels)
|
422 |
+
|
423 |
+
|
424 |
+
def get_audio_features(
|
425 |
+
audio_data, mel, max_len, data_truncating, data_filling, audio_cfg
|
426 |
+
):
|
427 |
+
"""
|
428 |
+
Calculate and add audio features to sample.
|
429 |
+
Sample: a dict containing all the data of current sample.
|
430 |
+
audio_data: a tensor of shape (T) containing audio data.
|
431 |
+
max_len: the maximum length of audio data.
|
432 |
+
data_truncating: the method of truncating data.
|
433 |
+
data_filling: the method of filling data.
|
434 |
+
audio_cfg: a dict containing audio configuration. Comes from model_cfg['audio_cfg'].
|
435 |
+
"""
|
436 |
+
sample = {}
|
437 |
+
|
438 |
+
# assert audio_data.size(-1) <= max_len, str(audio_data.size())
|
439 |
+
|
440 |
+
# split to three parts
|
441 |
+
chunk_frames = (
|
442 |
+
max_len // audio_cfg["hop_size"] + 1
|
443 |
+
) # the +1 related to how the spectrogram is computed
|
444 |
+
mel = mel[:chunk_frames]
|
445 |
+
|
446 |
+
audio_data = audio_data[..., :max_len]
|
447 |
+
sample["mel_fusion"] = mel
|
448 |
+
longer = torch.tensor([True])
|
449 |
+
|
450 |
+
sample["longer"] = longer
|
451 |
+
sample["waveform"] = audio_data
|
452 |
+
|
453 |
+
return sample
|
454 |
+
|
455 |
+
|
456 |
+
def preprocess(
|
457 |
+
sample,
|
458 |
+
audio_ext,
|
459 |
+
text_ext,
|
460 |
+
max_len,
|
461 |
+
audio_cfg,
|
462 |
+
class_index_dict=None,
|
463 |
+
data_filling="pad",
|
464 |
+
data_truncating="rand_trunc",
|
465 |
+
text_augment_selection=None,
|
466 |
+
):
|
467 |
+
"""
|
468 |
+
Preprocess a single sample for wdsdataloader.
|
469 |
+
"""
|
470 |
+
audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
|
471 |
+
audio_data = int16_to_float32(float32_to_int16(audio_data))
|
472 |
+
audio_data = torch.tensor(audio_data).float()
|
473 |
+
|
474 |
+
# TODO: (yusong) to be include in the future
|
475 |
+
# # if torchaudio not installed, use soundfile to load audio
|
476 |
+
# if torchaudio is None:
|
477 |
+
# audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
|
478 |
+
# audio_data = torch.tensor(audio_data).float()
|
479 |
+
# else:
|
480 |
+
# # https://github.com/webdataset/webdataset/blob/main/webdataset/autodecode.py
|
481 |
+
# with tempfile.TemporaryDirectory() as dirname:
|
482 |
+
# os.makedirs(dirname, exist_ok=True)
|
483 |
+
# fname = os.path.join(dirname, f"file.flac")
|
484 |
+
# with open(fname, "wb") as stream:
|
485 |
+
# stream.write(sample[audio_ext])
|
486 |
+
# audio_data, orig_sr = torchaudio.load(fname)
|
487 |
+
# audio_data = audio_data[0, :].float()
|
488 |
+
|
489 |
+
sample = get_audio_features(
|
490 |
+
sample, audio_data, max_len, data_truncating, data_filling, audio_cfg
|
491 |
+
)
|
492 |
+
del sample[audio_ext]
|
493 |
+
|
494 |
+
try:
|
495 |
+
json_dict_raw = json.loads(sample[text_ext].decode("utf-8"))
|
496 |
+
except:
|
497 |
+
print("sample[__url__]:", sample["__url__"])
|
498 |
+
|
499 |
+
# For selecting augmented text from dataset
|
500 |
+
if text_augment_selection is None or text_augment_selection == "none":
|
501 |
+
texts = json_dict_raw["text"]
|
502 |
+
elif text_augment_selection == "all":
|
503 |
+
if "text_augment_all" in json_dict_raw.keys():
|
504 |
+
texts = json_dict_raw["text_augment_all"]
|
505 |
+
else:
|
506 |
+
texts = json_dict_raw["text"]
|
507 |
+
elif text_augment_selection == "augment_only":
|
508 |
+
if "text_augment_all" in json_dict_raw.keys():
|
509 |
+
if json_dict_raw["text_augment_t5"] is None:
|
510 |
+
texts = json_dict_raw["text"]
|
511 |
+
else:
|
512 |
+
texts = json_dict_raw["text_augment_t5"]
|
513 |
+
else:
|
514 |
+
texts = json_dict_raw["text"]
|
515 |
+
else:
|
516 |
+
raise NotImplementedError(
|
517 |
+
f"text_augment_selection {text_augment_selection} not implemented"
|
518 |
+
)
|
519 |
+
sample["full_text"] = texts
|
520 |
+
|
521 |
+
if isinstance(texts, list) and isinstance(texts[0], str) and len(texts) > 1:
|
522 |
+
texts = random.choice(texts)
|
523 |
+
sample["raw_text"] = texts
|
524 |
+
sample["text"] = tokenizer(texts) # text shape: [num_token]
|
525 |
+
if class_index_dict is not None:
|
526 |
+
# https://stackoverflow.com/questions/48004243/how-to-share-large-read-only-dictionary-list-across-processes-in-multiprocessing
|
527 |
+
# https://stackoverflow.com/questions/45693949/storing-strings-in-a-multiprocessing-sharedctypes-array
|
528 |
+
# key, val = class_index_dict
|
529 |
+
# key = key[:].split('\n')
|
530 |
+
# _dict = {k: v for k, v in zip(key, val)}
|
531 |
+
sample["class_label"] = np.zeros(len(class_index_dict.keys()))
|
532 |
+
for x in json_dict_raw["tag"]:
|
533 |
+
sample["class_label"][class_index_dict[x]] = 1
|
534 |
+
sample["class_label"] = torch.tensor(sample["class_label"]).float()
|
535 |
+
del sample[text_ext]
|
536 |
+
sample["audio_name"] = sample["__key__"].split("/")[-1] + "." + audio_ext
|
537 |
+
sample["text_name"] = sample["__key__"].split("/")[-1] + "." + text_ext
|
538 |
+
sample["audio_orig_sr"] = orig_sr
|
539 |
+
return sample
|
540 |
+
|
541 |
+
|
542 |
+
def collate_fn(batch):
|
543 |
+
"""
|
544 |
+
Collate function for wdsdataloader.
|
545 |
+
batch: a list of dict, each dict is a sample
|
546 |
+
"""
|
547 |
+
# concatenate values in each dictionary. if it is a tensor, concatenate. if it is a list, extend.
|
548 |
+
batch_dict = {}
|
549 |
+
for k in batch[0].keys():
|
550 |
+
if isinstance(batch[0][k], dict): # dealwith bert tokenizer output
|
551 |
+
batch_dict[k] = {}
|
552 |
+
for kk in batch[0][k].keys():
|
553 |
+
tmp = []
|
554 |
+
for i in range(len(batch)):
|
555 |
+
tmp.append(batch[i][k][kk])
|
556 |
+
batch_dict[k][kk] = torch.vstack(tmp)
|
557 |
+
elif isinstance(batch[0][k], torch.Tensor):
|
558 |
+
batch_dict[k] = torch.stack([sample[k] for sample in batch])
|
559 |
+
elif isinstance(batch[0][k], np.ndarray):
|
560 |
+
batch_dict[k] = torch.tensor(np.stack([sample[k] for sample in batch]))
|
561 |
+
else:
|
562 |
+
batch_dict[k] = [sample[k] for sample in batch]
|
563 |
+
return batch_dict
|
564 |
+
|
565 |
+
|
566 |
+
# def get_wds_dataset(
|
567 |
+
# args,
|
568 |
+
# model_cfg,
|
569 |
+
# is_train,
|
570 |
+
# audio_ext="flac",
|
571 |
+
# text_ext="json",
|
572 |
+
# max_len=480000,
|
573 |
+
# proportion=1.0,
|
574 |
+
# sizefilepath_=None,
|
575 |
+
# is_local=None,
|
576 |
+
# ):
|
577 |
+
# """
|
578 |
+
# Get a dataset for wdsdataloader.
|
579 |
+
# """
|
580 |
+
# if is_local is None and (not args.remotedata is None):
|
581 |
+
# is_local = not args.remotedata
|
582 |
+
|
583 |
+
# input_shards = args.train_data if is_train else args.val_data
|
584 |
+
# assert input_shards is not None
|
585 |
+
|
586 |
+
# if not sizefilepath_ is None:
|
587 |
+
# sizefilepath = sizefilepath_
|
588 |
+
# else:
|
589 |
+
# sizefilepath = os.path.join(os.path.dirname(input_shards[0]), "sizes.json")
|
590 |
+
|
591 |
+
# if proportion != 1.0:
|
592 |
+
# num_samples, num_shards, input_shards, _ = sample_prop(
|
593 |
+
# sizefilepath, input_shards, proportion, is_local=is_local
|
594 |
+
# )
|
595 |
+
# else:
|
596 |
+
# num_samples, num_shards = get_dataset_size(
|
597 |
+
# input_shards, sizefilepath_=sizefilepath_, is_local=is_local
|
598 |
+
# )
|
599 |
+
|
600 |
+
# if not num_samples:
|
601 |
+
# if is_train:
|
602 |
+
# num_samples = args.train_num_samples
|
603 |
+
# if not num_samples:
|
604 |
+
# raise RuntimeError(
|
605 |
+
# "Currently, number of dataset samples must be specified for training dataset. "
|
606 |
+
# "Please specify via `--train-num-samples` if no dataset length info present."
|
607 |
+
# )
|
608 |
+
# else:
|
609 |
+
# num_samples = (
|
610 |
+
# args.val_num_samples or 0
|
611 |
+
# ) # eval will just exhaust the iterator if not specified
|
612 |
+
|
613 |
+
# pipeline = [wds.SimpleShardList(input_shards)]
|
614 |
+
# # at this point we have an iterator over all the shards
|
615 |
+
# # TODO: (yusong): add a if statement of distributed. If not, we don't need to split_by_node
|
616 |
+
# if is_train or args.parallel_eval:
|
617 |
+
# pipeline.extend(
|
618 |
+
# [
|
619 |
+
# wds.detshuffle(
|
620 |
+
# bufsize=_SHARD_SHUFFLE_SIZE,
|
621 |
+
# initial=_SHARD_SHUFFLE_INITIAL,
|
622 |
+
# seed=args.seed,
|
623 |
+
# ),
|
624 |
+
# wds.split_by_node,
|
625 |
+
# wds.split_by_worker,
|
626 |
+
# # at this point, we have an iterator over the shards assigned to each worker at each node
|
627 |
+
# wds.tarfile_to_samples(handler=log_and_continue),
|
628 |
+
# wds.shuffle(
|
629 |
+
# bufsize=_SAMPLE_SHUFFLE_SIZE,
|
630 |
+
# initial=_SAMPLE_SHUFFLE_INITIAL,
|
631 |
+
# rng=random.Random(args.seed),
|
632 |
+
# ),
|
633 |
+
# # wds.repeatedly, # FIXME determine if this is beneficial
|
634 |
+
# ]
|
635 |
+
# )
|
636 |
+
# else:
|
637 |
+
# pipeline.extend(
|
638 |
+
# [
|
639 |
+
# wds.split_by_worker,
|
640 |
+
# # at this point, we have an iterator over the shards assigned to each worker
|
641 |
+
# wds.tarfile_to_samples(handler=log_and_continue),
|
642 |
+
# ]
|
643 |
+
# )
|
644 |
+
# pipeline.append(
|
645 |
+
# wds.map(
|
646 |
+
# partial(
|
647 |
+
# preprocess,
|
648 |
+
# audio_ext=audio_ext,
|
649 |
+
# text_ext=text_ext,
|
650 |
+
# max_len=max_len,
|
651 |
+
# audio_cfg=model_cfg["audio_cfg"],
|
652 |
+
# class_index_dict=copy.deepcopy(args.class_index_dict),
|
653 |
+
# data_filling=args.data_filling,
|
654 |
+
# data_truncating=args.data_truncating,
|
655 |
+
# text_augment_selection=args.text_augment_selection,
|
656 |
+
# )
|
657 |
+
# ),
|
658 |
+
# )
|
659 |
+
|
660 |
+
# pipeline.append(
|
661 |
+
# wds.batched(
|
662 |
+
# args.batch_size,
|
663 |
+
# partial=not (is_train or args.parallel_eval),
|
664 |
+
# collation_fn=collate_fn,
|
665 |
+
# )
|
666 |
+
# )
|
667 |
+
|
668 |
+
# dataset = wds.DataPipeline(*pipeline)
|
669 |
+
# if is_train or args.parallel_eval:
|
670 |
+
# # (yusong): Currently parallel evaluation will be not precise as we are repeat the last few samples.
|
671 |
+
# # (yusong): See comments below.
|
672 |
+
# # roll over and repeat a few samples to get same number of full batches on each node
|
673 |
+
# global_batch_size = args.batch_size * args.world_size
|
674 |
+
# num_batches = math.ceil(num_samples / global_batch_size)
|
675 |
+
# num_workers = max(1, args.workers)
|
676 |
+
# num_worker_batches = math.ceil(
|
677 |
+
# num_batches / num_workers
|
678 |
+
# ) # per dataloader worker
|
679 |
+
# num_batches = num_worker_batches * num_workers
|
680 |
+
# num_samples = num_batches * global_batch_size
|
681 |
+
# dataset = dataset.with_epoch(
|
682 |
+
# num_worker_batches
|
683 |
+
# ) # each worker is iterating over this
|
684 |
+
# else:
|
685 |
+
# # last batches are partial, eval is done on single (master) node
|
686 |
+
# num_batches = math.ceil(num_samples / args.batch_size)
|
687 |
+
|
688 |
+
# kwargs = {}
|
689 |
+
# if args.horovod: # multi-node training on summit
|
690 |
+
# kwargs["multiprocessing_context"] = "forkserver"
|
691 |
+
|
692 |
+
# dataloader = wds.WebLoader(
|
693 |
+
# dataset, batch_size=None, shuffle=False, num_workers=args.workers, **kwargs
|
694 |
+
# )
|
695 |
+
|
696 |
+
# # FIXME not clear which approach is better, with_epoch before vs after dataloader?
|
697 |
+
# # hoping to resolve via https://github.com/webdataset/webdataset/issues/169
|
698 |
+
# # if is_train:
|
699 |
+
# # # roll over and repeat a few samples to get same number of full batches on each node
|
700 |
+
# # global_batch_size = args.batch_size * args.world_size
|
701 |
+
# # num_batches = math.ceil(num_samples / global_batch_size)
|
702 |
+
# # num_workers = max(1, args.workers)
|
703 |
+
# # num_batches = math.ceil(num_batches / num_workers) * num_workers
|
704 |
+
# # num_samples = num_batches * global_batch_size
|
705 |
+
# # dataloader = dataloader.with_epoch(num_batches)
|
706 |
+
# # else:
|
707 |
+
# # # last batches are partial, eval is done on single (master) node
|
708 |
+
# # num_batches = math.ceil(num_samples / args.batch_size)
|
709 |
+
|
710 |
+
# # add meta-data to dataloader instance for convenience
|
711 |
+
# dataloader.num_batches = num_batches
|
712 |
+
# dataloader.num_samples = num_samples
|
713 |
+
|
714 |
+
# return DataInfo(dataloader, None)
|
715 |
+
|
716 |
+
|
717 |
+
def wds_batch_list2dict(
|
718 |
+
batch,
|
719 |
+
keys=[
|
720 |
+
"__url__",
|
721 |
+
"__key__",
|
722 |
+
"waveform",
|
723 |
+
"text",
|
724 |
+
"raw_text",
|
725 |
+
"audio_name",
|
726 |
+
"text_name",
|
727 |
+
"audio_orig_sr",
|
728 |
+
],
|
729 |
+
):
|
730 |
+
"""
|
731 |
+
Return a dictionary of the batch, with keys as the names of the fields.
|
732 |
+
"""
|
733 |
+
assert len(keys) == len(
|
734 |
+
batch
|
735 |
+
), "batch must have same number of keys as keys argument"
|
736 |
+
return {keys[i]: batch[i] for i in range(len(batch))}
|
737 |
+
|
738 |
+
|
739 |
+
def get_csv_dataset(args, preprocess_fn, is_train):
|
740 |
+
input_filename = args.train_data if is_train else args.val_data
|
741 |
+
assert input_filename
|
742 |
+
dataset = CsvDataset(
|
743 |
+
input_filename,
|
744 |
+
preprocess_fn,
|
745 |
+
img_key=args.csv_img_key,
|
746 |
+
caption_key=args.csv_caption_key,
|
747 |
+
sep=args.csv_separator,
|
748 |
+
)
|
749 |
+
num_samples = len(dataset)
|
750 |
+
sampler = DistributedSampler(dataset) if args.distributed and is_train else None
|
751 |
+
shuffle = is_train and sampler is None
|
752 |
+
|
753 |
+
dataloader = DataLoader(
|
754 |
+
dataset,
|
755 |
+
batch_size=args.batch_size,
|
756 |
+
shuffle=shuffle,
|
757 |
+
num_workers=args.workers,
|
758 |
+
pin_memory=True,
|
759 |
+
sampler=sampler,
|
760 |
+
drop_last=is_train,
|
761 |
+
)
|
762 |
+
dataloader.num_samples = num_samples
|
763 |
+
dataloader.num_batches = len(dataloader)
|
764 |
+
|
765 |
+
return DataInfo(dataloader, sampler)
|
766 |
+
|
767 |
+
|
768 |
+
def get_toy_dataset(args, model_cfg, is_train):
|
769 |
+
index_path = args.train_data if is_train else args.val_data
|
770 |
+
ipc_path = args.train_ipc if is_train else args.val_ipc
|
771 |
+
assert index_path and ipc_path
|
772 |
+
eval_mode = not is_train
|
773 |
+
dataset = ToyDataset(index_path, ipc_path, model_cfg, eval_mode=eval_mode)
|
774 |
+
|
775 |
+
num_samples = len(dataset)
|
776 |
+
sampler = (
|
777 |
+
DistributedSampler(dataset, shuffle=False)
|
778 |
+
if args.distributed and is_train
|
779 |
+
else None
|
780 |
+
)
|
781 |
+
|
782 |
+
dataloader = DataLoader(
|
783 |
+
dataset,
|
784 |
+
batch_size=args.batch_size,
|
785 |
+
shuffle=False,
|
786 |
+
num_workers=args.workers,
|
787 |
+
sampler=sampler,
|
788 |
+
drop_last=is_train,
|
789 |
+
)
|
790 |
+
dataloader.num_samples = num_samples
|
791 |
+
dataloader.num_batches = len(dataloader)
|
792 |
+
|
793 |
+
return DataInfo(dataloader, sampler)
|
794 |
+
|
795 |
+
|
796 |
+
def get_dataset_fn(data_path, dataset_type):
|
797 |
+
if dataset_type == "webdataset":
|
798 |
+
return get_wds_dataset
|
799 |
+
elif dataset_type == "csv":
|
800 |
+
return get_csv_dataset
|
801 |
+
elif dataset_type == "auto":
|
802 |
+
ext = data_path.split(".")[-1]
|
803 |
+
if ext in ["csv", "tsv"]:
|
804 |
+
return get_csv_dataset
|
805 |
+
elif ext in ["tar"]:
|
806 |
+
return get_wds_dataset
|
807 |
+
else:
|
808 |
+
raise ValueError(
|
809 |
+
f"Tried to figure out dataset type, but failed for extention {ext}."
|
810 |
+
)
|
811 |
+
elif dataset_type == "toy":
|
812 |
+
return get_toy_dataset
|
813 |
+
else:
|
814 |
+
raise ValueError(f"Unsupported dataset type: {dataset_type}")
|
815 |
+
|
816 |
+
|
817 |
+
def get_data(args, model_cfg):
|
818 |
+
data = {}
|
819 |
+
|
820 |
+
args.class_index_dict = load_class_label(args.class_label_path)
|
821 |
+
|
822 |
+
if args.datasetinfos is None:
|
823 |
+
args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
|
824 |
+
if args.dataset_type == "webdataset":
|
825 |
+
args.train_data = get_tar_path_from_dataset_name(
|
826 |
+
args.datasetnames,
|
827 |
+
args.datasetinfos,
|
828 |
+
islocal=not args.remotedata,
|
829 |
+
proportion=args.dataset_proportion,
|
830 |
+
dataset_path=args.datasetpath,
|
831 |
+
full_dataset=args.full_train_dataset,
|
832 |
+
)
|
833 |
+
|
834 |
+
if args.full_train_dataset is None:
|
835 |
+
args.full_train_dataset = []
|
836 |
+
if args.exclude_eval_dataset is None:
|
837 |
+
args.exclude_eval_dataset = []
|
838 |
+
excluded_eval_datasets = args.full_train_dataset + args.exclude_eval_dataset
|
839 |
+
|
840 |
+
val_dataset_names = (
|
841 |
+
[n for n in args.datasetnames if n not in excluded_eval_datasets]
|
842 |
+
if excluded_eval_datasets
|
843 |
+
else args.datasetnames
|
844 |
+
)
|
845 |
+
args.val_dataset_names = val_dataset_names
|
846 |
+
args.val_data = get_tar_path_from_dataset_name(
|
847 |
+
val_dataset_names,
|
848 |
+
["valid", "test", "eval"],
|
849 |
+
islocal=not args.remotedata,
|
850 |
+
proportion=1,
|
851 |
+
dataset_path=args.datasetpath,
|
852 |
+
full_dataset=None,
|
853 |
+
)
|
854 |
+
|
855 |
+
if args.train_data:
|
856 |
+
data["train"] = get_dataset_fn(args.train_data, args.dataset_type)(
|
857 |
+
args, model_cfg, is_train=True
|
858 |
+
)
|
859 |
+
|
860 |
+
if args.val_data:
|
861 |
+
data["val"] = get_dataset_fn(args.val_data, args.dataset_type)(
|
862 |
+
args, model_cfg, is_train=False
|
863 |
+
)
|
864 |
+
|
865 |
+
return data
|
audioldm2/clap/training/params.py
ADDED
@@ -0,0 +1,563 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import argparse
|
2 |
+
|
3 |
+
|
4 |
+
def get_default_params(model_name):
|
5 |
+
# Params from paper (https://arxiv.org/pdf/2103.00020.pdf)
|
6 |
+
model_name = model_name.lower()
|
7 |
+
if "vit" in model_name:
|
8 |
+
return {"lr": 5.0e-4, "beta1": 0.9, "beta2": 0.98, "eps": 1.0e-6}
|
9 |
+
else:
|
10 |
+
return {"lr": 5.0e-4, "beta1": 0.9, "beta2": 0.999, "eps": 1.0e-8}
|
11 |
+
|
12 |
+
|
13 |
+
def parse_args():
|
14 |
+
parser = argparse.ArgumentParser()
|
15 |
+
parser.add_argument(
|
16 |
+
"--train-data",
|
17 |
+
type=str,
|
18 |
+
default=None,
|
19 |
+
help="Path to h5 filewith training data",
|
20 |
+
)
|
21 |
+
parser.add_argument(
|
22 |
+
"--val-data",
|
23 |
+
type=str,
|
24 |
+
default=None,
|
25 |
+
help="Path to h5 file with validation data",
|
26 |
+
)
|
27 |
+
parser.add_argument(
|
28 |
+
"--freeze-text",
|
29 |
+
default=False,
|
30 |
+
action="store_true",
|
31 |
+
help="if you need to freeze the text encoder, make this True",
|
32 |
+
)
|
33 |
+
parser.add_argument(
|
34 |
+
"--freeze-text-after",
|
35 |
+
type=int,
|
36 |
+
default=-1,
|
37 |
+
help="if you need to freeze the text encoder after (include) epoch x, set this param to x. Set -1 to disable it",
|
38 |
+
)
|
39 |
+
parser.add_argument(
|
40 |
+
"--train-ipc",
|
41 |
+
type=str,
|
42 |
+
default=None,
|
43 |
+
help="Path to npy file of the number of instance per class in training data",
|
44 |
+
)
|
45 |
+
parser.add_argument(
|
46 |
+
"--val-ipc",
|
47 |
+
type=str,
|
48 |
+
default=None,
|
49 |
+
help="Path to npy file of the number of instance per class in validation data",
|
50 |
+
)
|
51 |
+
parser.add_argument(
|
52 |
+
"--train-num-samples",
|
53 |
+
type=int,
|
54 |
+
default=None,
|
55 |
+
help="Number of samples in dataset. Required for webdataset if not available in info file.",
|
56 |
+
)
|
57 |
+
parser.add_argument(
|
58 |
+
"--val-num-samples",
|
59 |
+
type=int,
|
60 |
+
default=None,
|
61 |
+
help="Number of samples in dataset. Useful for webdataset if not available in info file.",
|
62 |
+
)
|
63 |
+
parser.add_argument(
|
64 |
+
"--dataset-type",
|
65 |
+
choices=["webdataset", "csv", "auto", "toy"],
|
66 |
+
default="auto",
|
67 |
+
help="Which type of dataset to process.",
|
68 |
+
)
|
69 |
+
parser.add_argument(
|
70 |
+
"--csv-separator",
|
71 |
+
type=str,
|
72 |
+
default="\t",
|
73 |
+
help="For csv-like datasets, which separator to use.",
|
74 |
+
)
|
75 |
+
parser.add_argument(
|
76 |
+
"--csv-img-key",
|
77 |
+
type=str,
|
78 |
+
default="filepath",
|
79 |
+
help="For csv-like datasets, the name of the key for the image paths.",
|
80 |
+
)
|
81 |
+
parser.add_argument(
|
82 |
+
"--csv-caption-key",
|
83 |
+
type=str,
|
84 |
+
default="title",
|
85 |
+
help="For csv-like datasets, the name of the key for the captions.",
|
86 |
+
)
|
87 |
+
parser.add_argument(
|
88 |
+
"--imagenet-val",
|
89 |
+
type=str,
|
90 |
+
default=None,
|
91 |
+
help="Path to imagenet val set for conducting zero shot evaluation.",
|
92 |
+
)
|
93 |
+
parser.add_argument(
|
94 |
+
"--imagenet-v2",
|
95 |
+
type=str,
|
96 |
+
default=None,
|
97 |
+
help="Path to imagenet v2 for conducting zero shot evaluation.",
|
98 |
+
)
|
99 |
+
parser.add_argument(
|
100 |
+
"--datasetnames",
|
101 |
+
nargs="+",
|
102 |
+
default=None,
|
103 |
+
help="If loading webdataset, spedify the dataset names to load. Can be some of these: Clotho, audioset, audiocaps, BBCSoundEffects",
|
104 |
+
)
|
105 |
+
parser.add_argument(
|
106 |
+
"--full-train-dataset",
|
107 |
+
nargs="+",
|
108 |
+
default=None,
|
109 |
+
help="Which dataset will be trained with all the subsets. (train+test)",
|
110 |
+
)
|
111 |
+
parser.add_argument(
|
112 |
+
"--exclude-eval-dataset",
|
113 |
+
nargs="+",
|
114 |
+
default=None,
|
115 |
+
help="Which dataset will be excluded with evaluation",
|
116 |
+
)
|
117 |
+
parser.add_argument(
|
118 |
+
"--datasetinfos",
|
119 |
+
nargs="+",
|
120 |
+
default=None,
|
121 |
+
help="If loading webdataset, spedify the dataset types to load. Can be some of these: train, test, valid, unbalanced_train, balanced_train, eval",
|
122 |
+
)
|
123 |
+
parser.add_argument(
|
124 |
+
"--dataset-proportion",
|
125 |
+
type=float,
|
126 |
+
default=1.0,
|
127 |
+
help="How much proportion of dataset we want to train.",
|
128 |
+
)
|
129 |
+
parser.add_argument(
|
130 |
+
"--remotedata",
|
131 |
+
default=False,
|
132 |
+
action="store_true",
|
133 |
+
help="if the dataset is remote, set this flag",
|
134 |
+
)
|
135 |
+
parser.add_argument(
|
136 |
+
"--class-label-path",
|
137 |
+
type=str,
|
138 |
+
default=None,
|
139 |
+
help="The path of the class label pickle or csv.",
|
140 |
+
)
|
141 |
+
parser.add_argument(
|
142 |
+
"--datasetpath",
|
143 |
+
type=str,
|
144 |
+
default="/mnt/audio_clip/webdataset_tar",
|
145 |
+
help="The path to the dataset",
|
146 |
+
)
|
147 |
+
parser.add_argument(
|
148 |
+
"--logs",
|
149 |
+
type=str,
|
150 |
+
default="./logs/",
|
151 |
+
help="Where to store tensorboard logs. Use None to avoid storing logs.",
|
152 |
+
)
|
153 |
+
parser.add_argument(
|
154 |
+
"--log-local",
|
155 |
+
action="store_true",
|
156 |
+
default=False,
|
157 |
+
help="log files on local master, otherwise global master only.",
|
158 |
+
)
|
159 |
+
parser.add_argument(
|
160 |
+
"--name",
|
161 |
+
type=str,
|
162 |
+
default=None,
|
163 |
+
help="Optional identifier for the experiment when storing logs. Otherwise use current time.",
|
164 |
+
)
|
165 |
+
parser.add_argument(
|
166 |
+
"--workers", type=int, default=1, help="Number of workers per GPU."
|
167 |
+
)
|
168 |
+
parser.add_argument(
|
169 |
+
"--batch-size", type=int, default=64, help="Batch size per GPU."
|
170 |
+
)
|
171 |
+
parser.add_argument(
|
172 |
+
"--epochs", type=int, default=32, help="Number of epochs to train for."
|
173 |
+
)
|
174 |
+
parser.add_argument("--lr", type=float, default=None, help="Learning rate.")
|
175 |
+
parser.add_argument("--beta1", type=float, default=None, help="Adam beta 1.")
|
176 |
+
parser.add_argument("--beta2", type=float, default=None, help="Adam beta 2.")
|
177 |
+
parser.add_argument("--eps", type=float, default=None, help="Adam epsilon.")
|
178 |
+
parser.add_argument("--momentum", type=float, default=None, help="SGD epsilon.")
|
179 |
+
parser.add_argument("--wd", type=float, default=0.2, help="Weight decay.")
|
180 |
+
|
181 |
+
parser.add_argument(
|
182 |
+
"--split-opt",
|
183 |
+
action="store_true",
|
184 |
+
default=False,
|
185 |
+
help="Use this flag to skip the learning rate decay.",
|
186 |
+
)
|
187 |
+
parser.add_argument(
|
188 |
+
"--lr-pretrained", type=float, default=None, help="Learning rate for text."
|
189 |
+
)
|
190 |
+
parser.add_argument(
|
191 |
+
"--beta1-pretrained", type=float, default=None, help="Adam beta 1 for text."
|
192 |
+
)
|
193 |
+
parser.add_argument(
|
194 |
+
"--beta2-pretrained", type=float, default=None, help="Adam beta 2 for text."
|
195 |
+
)
|
196 |
+
parser.add_argument(
|
197 |
+
"--eps-pretrained", type=float, default=None, help="Adam epsilon for text."
|
198 |
+
)
|
199 |
+
parser.add_argument(
|
200 |
+
"--wd-pretrained", type=float, default=0.2, help="Weight decay for text."
|
201 |
+
)
|
202 |
+
parser.add_argument(
|
203 |
+
"--momentum-pretrained", type=float, default=0.9, help="Momentum for text."
|
204 |
+
)
|
205 |
+
parser.add_argument(
|
206 |
+
"--lr-new", type=float, default=None, help="Learning rate for audio."
|
207 |
+
)
|
208 |
+
parser.add_argument(
|
209 |
+
"--beta1-new", type=float, default=None, help="Adam beta 1 for audio."
|
210 |
+
)
|
211 |
+
parser.add_argument(
|
212 |
+
"--beta2-new", type=float, default=None, help="Adam beta 2 for audio."
|
213 |
+
)
|
214 |
+
parser.add_argument(
|
215 |
+
"--eps-new", type=float, default=None, help="Adam epsilon for audio."
|
216 |
+
)
|
217 |
+
parser.add_argument(
|
218 |
+
"--wd-new", type=float, default=0.2, help="Weight decay for audio."
|
219 |
+
)
|
220 |
+
parser.add_argument(
|
221 |
+
"--momentum-new", type=float, default=0.9, help="Momentum for audio."
|
222 |
+
)
|
223 |
+
parser.add_argument(
|
224 |
+
"--warmup", type=int, default=10000, help="Number of steps to warmup for."
|
225 |
+
)
|
226 |
+
parser.add_argument(
|
227 |
+
"--use-bn-sync",
|
228 |
+
default=False,
|
229 |
+
action="store_true",
|
230 |
+
help="Whether to use batch norm sync.",
|
231 |
+
)
|
232 |
+
parser.add_argument(
|
233 |
+
"--skip-scheduler",
|
234 |
+
action="store_true",
|
235 |
+
default=False,
|
236 |
+
help="Use this flag to skip the learning rate decay.",
|
237 |
+
)
|
238 |
+
parser.add_argument(
|
239 |
+
"--save-frequency", type=int, default=1, help="How often to save checkpoints."
|
240 |
+
)
|
241 |
+
parser.add_argument(
|
242 |
+
"--save-top-performance",
|
243 |
+
type=int,
|
244 |
+
default=0,
|
245 |
+
help="Save the top x performance weights if the value >0",
|
246 |
+
)
|
247 |
+
parser.add_argument(
|
248 |
+
"--save-most-recent",
|
249 |
+
action="store_true",
|
250 |
+
default=False,
|
251 |
+
help="Always save the most recent model trained to epoch_latest.pt.",
|
252 |
+
)
|
253 |
+
parser.add_argument(
|
254 |
+
"--zeroshot-frequency", type=int, default=2, help="How often to run zero shot."
|
255 |
+
)
|
256 |
+
parser.add_argument(
|
257 |
+
"--val-frequency",
|
258 |
+
type=int,
|
259 |
+
default=1,
|
260 |
+
help="How often to run evaluation with val data.",
|
261 |
+
)
|
262 |
+
parser.add_argument(
|
263 |
+
"--resume",
|
264 |
+
default=None,
|
265 |
+
type=str,
|
266 |
+
help="path to latest checkpoint (default: none)",
|
267 |
+
)
|
268 |
+
parser.add_argument(
|
269 |
+
"--precision",
|
270 |
+
choices=["amp", "fp16", "fp32"],
|
271 |
+
default="amp",
|
272 |
+
help="Floating point precision.",
|
273 |
+
)
|
274 |
+
parser.add_argument(
|
275 |
+
"--amodel",
|
276 |
+
type=str,
|
277 |
+
default="RN50",
|
278 |
+
help="Name of the audio backbone to use.",
|
279 |
+
)
|
280 |
+
parser.add_argument(
|
281 |
+
"--tmodel",
|
282 |
+
type=str,
|
283 |
+
default="transformer",
|
284 |
+
help="Name of the text backbone to use. Can be [transformer, bert, roberta, bart]",
|
285 |
+
)
|
286 |
+
parser.add_argument(
|
287 |
+
"--pretrained-audio",
|
288 |
+
default="",
|
289 |
+
type=str,
|
290 |
+
help="Use a pretrained audio model weights for the audio encoder of CLAP",
|
291 |
+
)
|
292 |
+
parser.add_argument(
|
293 |
+
"--pretrained-text",
|
294 |
+
default="",
|
295 |
+
type=str,
|
296 |
+
help="Use a pretrained text model weights for the text encoder of CLAP",
|
297 |
+
)
|
298 |
+
parser.add_argument(
|
299 |
+
"--pretrained",
|
300 |
+
default="",
|
301 |
+
type=str,
|
302 |
+
help="Use a pretrained CLIP model weights with the specified tag or file path.",
|
303 |
+
)
|
304 |
+
parser.add_argument(
|
305 |
+
"--pretrained-image",
|
306 |
+
default=False,
|
307 |
+
action="store_true",
|
308 |
+
help="Load imagenet pretrained weights for image tower backbone if available.",
|
309 |
+
)
|
310 |
+
parser.add_argument(
|
311 |
+
"--lock-image",
|
312 |
+
default=False,
|
313 |
+
action="store_true",
|
314 |
+
help="Lock full image tower by disabling gradients.",
|
315 |
+
)
|
316 |
+
parser.add_argument(
|
317 |
+
"--lock-image-unlocked-groups",
|
318 |
+
type=int,
|
319 |
+
default=0,
|
320 |
+
help="Leave last n image tower layer groups unlocked.",
|
321 |
+
)
|
322 |
+
parser.add_argument(
|
323 |
+
"--lock-image-freeze-bn-stats",
|
324 |
+
default=False,
|
325 |
+
action="store_true",
|
326 |
+
help="Freeze BatchNorm running stats in image tower for any locked layers.",
|
327 |
+
)
|
328 |
+
parser.add_argument(
|
329 |
+
"--local-loss",
|
330 |
+
default=False,
|
331 |
+
action="store_true",
|
332 |
+
help="calculate loss w/ local features @ global (instead of realizing full global @ global matrix)",
|
333 |
+
)
|
334 |
+
parser.add_argument(
|
335 |
+
"--gather-with-grad",
|
336 |
+
default=False,
|
337 |
+
action="store_true",
|
338 |
+
help="enable full distributed gradient for feature gather",
|
339 |
+
)
|
340 |
+
parser.add_argument(
|
341 |
+
"--force-quick-gelu",
|
342 |
+
default=False,
|
343 |
+
action="store_true",
|
344 |
+
help="Force use of QuickGELU activation for non-OpenAI transformer models.",
|
345 |
+
)
|
346 |
+
parser.add_argument(
|
347 |
+
"--torchscript",
|
348 |
+
default=False,
|
349 |
+
action="store_true",
|
350 |
+
help="torch.jit.script the model, also uses jit version of OpenAI models if pretrained=='openai'",
|
351 |
+
)
|
352 |
+
parser.add_argument(
|
353 |
+
"--trace",
|
354 |
+
default=False,
|
355 |
+
action="store_true",
|
356 |
+
help="torch.jit.trace the model for inference / eval only",
|
357 |
+
)
|
358 |
+
# arguments for distributed training
|
359 |
+
parser.add_argument(
|
360 |
+
"--dist-url",
|
361 |
+
default="env://",
|
362 |
+
type=str,
|
363 |
+
help="url used to set up distributed training",
|
364 |
+
)
|
365 |
+
parser.add_argument(
|
366 |
+
"--dist-backend", default="nccl", type=str, help="distributed backend"
|
367 |
+
)
|
368 |
+
parser.add_argument(
|
369 |
+
"--report-to",
|
370 |
+
default="",
|
371 |
+
type=str,
|
372 |
+
help="Options are ['wandb', 'tensorboard', 'wandb,tensorboard']",
|
373 |
+
)
|
374 |
+
parser.add_argument(
|
375 |
+
"--wandb-notes", default="", type=str, help="Notes if logging with wandb"
|
376 |
+
)
|
377 |
+
parser.add_argument(
|
378 |
+
"--C", type=float, default=3.16, help="inverse regularizer for logistic reg."
|
379 |
+
)
|
380 |
+
parser.add_argument(
|
381 |
+
"--debug",
|
382 |
+
default=False,
|
383 |
+
action="store_true",
|
384 |
+
help="If true, more information is logged.",
|
385 |
+
)
|
386 |
+
parser.add_argument(
|
387 |
+
"--copy-codebase",
|
388 |
+
default=False,
|
389 |
+
action="store_true",
|
390 |
+
help="If true, we copy the entire base on the log diretory, and execute from there.",
|
391 |
+
)
|
392 |
+
parser.add_argument(
|
393 |
+
"--horovod",
|
394 |
+
default=False,
|
395 |
+
action="store_true",
|
396 |
+
help="Use horovod for distributed training.",
|
397 |
+
)
|
398 |
+
parser.add_argument(
|
399 |
+
"--ddp-static-graph",
|
400 |
+
default=False,
|
401 |
+
action="store_true",
|
402 |
+
help="Enable static graph optimization for DDP in PyTorch >= 1.11.",
|
403 |
+
)
|
404 |
+
parser.add_argument(
|
405 |
+
"--no-set-device-rank",
|
406 |
+
default=False,
|
407 |
+
action="store_true",
|
408 |
+
help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc).",
|
409 |
+
)
|
410 |
+
parser.add_argument("--seed", type=int, default=4242, help="Default random seed.")
|
411 |
+
|
412 |
+
parser.add_argument(
|
413 |
+
"--top-k-checkpoint-select-dataset",
|
414 |
+
type=str,
|
415 |
+
default="all",
|
416 |
+
help="The dataset of selecting top-k checkpoint.",
|
417 |
+
)
|
418 |
+
|
419 |
+
# @R10, @R@5, @R1, mAP@10
|
420 |
+
parser.add_argument(
|
421 |
+
"--top-k-checkpoint-select-metric",
|
422 |
+
type=str,
|
423 |
+
default="_R@10",
|
424 |
+
help="The metric for selecting top-k checkpoint.",
|
425 |
+
)
|
426 |
+
parser.add_argument(
|
427 |
+
"--openai-model-cache-dir",
|
428 |
+
type=str,
|
429 |
+
default="~/.cache/clip",
|
430 |
+
help="Directory to download OpenAI models.",
|
431 |
+
)
|
432 |
+
parser.add_argument(
|
433 |
+
"--optimizer",
|
434 |
+
type=str,
|
435 |
+
default="adamw",
|
436 |
+
help="can be AdamW or SGD",
|
437 |
+
)
|
438 |
+
parser.add_argument(
|
439 |
+
"--parallel-eval",
|
440 |
+
default=False,
|
441 |
+
action="store_true",
|
442 |
+
help="Eval in parallel (multi-GPU, multi-node).",
|
443 |
+
)
|
444 |
+
|
445 |
+
parser.add_argument(
|
446 |
+
"--no-eval",
|
447 |
+
default=False,
|
448 |
+
action="store_true",
|
449 |
+
help="Training without evaluation.",
|
450 |
+
)
|
451 |
+
|
452 |
+
parser.add_argument(
|
453 |
+
"--lp-mlp",
|
454 |
+
default=False,
|
455 |
+
action="store_true",
|
456 |
+
help="Linear Probe using MLP layer or not.",
|
457 |
+
)
|
458 |
+
|
459 |
+
parser.add_argument(
|
460 |
+
"--lp-freeze",
|
461 |
+
default=False,
|
462 |
+
action="store_true",
|
463 |
+
help="Linear Probe using Freeze CLAP or not",
|
464 |
+
)
|
465 |
+
|
466 |
+
parser.add_argument(
|
467 |
+
"--lp-act",
|
468 |
+
default="None",
|
469 |
+
type=str,
|
470 |
+
help="Options are ['relu','elu','prelu','softmax','sigmoid']",
|
471 |
+
)
|
472 |
+
|
473 |
+
parser.add_argument(
|
474 |
+
"--lp-loss", type=str, default="bce", help="Loss func of Linear Probe."
|
475 |
+
)
|
476 |
+
|
477 |
+
parser.add_argument(
|
478 |
+
"--lp-metrics",
|
479 |
+
type=str,
|
480 |
+
default="map,mauc,acc",
|
481 |
+
help="Metrics of Linear Probe.",
|
482 |
+
)
|
483 |
+
|
484 |
+
parser.add_argument(
|
485 |
+
"--lp-lr", type=float, default=1e-4, help="learning rate of linear probe"
|
486 |
+
)
|
487 |
+
parser.add_argument(
|
488 |
+
"--kappa",
|
489 |
+
type=float,
|
490 |
+
default=0,
|
491 |
+
help="the kappa in the weighted contrastive loss, default is to turn off the weighted contrastive loss",
|
492 |
+
)
|
493 |
+
|
494 |
+
parser.add_argument(
|
495 |
+
"--data-filling",
|
496 |
+
type=str,
|
497 |
+
default="pad",
|
498 |
+
help="type of data filling when the audio length is shorter than the max length."
|
499 |
+
"Can be one of the following: repeat, repeatpad, pad",
|
500 |
+
)
|
501 |
+
parser.add_argument(
|
502 |
+
"--data-truncating",
|
503 |
+
type=str,
|
504 |
+
default="rand_trunc",
|
505 |
+
help="type of data truncation when the audio length is longer than the max length."
|
506 |
+
"Can be one of the following: rand_trunc, fusion",
|
507 |
+
)
|
508 |
+
|
509 |
+
parser.add_argument(
|
510 |
+
"--clap-mlploss",
|
511 |
+
default=False,
|
512 |
+
action="store_true",
|
513 |
+
help="Using MLP loss for CLAP model or not",
|
514 |
+
)
|
515 |
+
|
516 |
+
parser.add_argument(
|
517 |
+
"--wandb-id",
|
518 |
+
type=str,
|
519 |
+
default=None,
|
520 |
+
help="the id of wandb experiment to restore.",
|
521 |
+
)
|
522 |
+
|
523 |
+
parser.add_argument(
|
524 |
+
"--sleep", type=float, default=0, help="sleep n seconds before start training"
|
525 |
+
)
|
526 |
+
|
527 |
+
# variable length processing
|
528 |
+
parser.add_argument(
|
529 |
+
"--enable-fusion",
|
530 |
+
default=False,
|
531 |
+
action="store_true",
|
532 |
+
help="Enable feature funsion for variable-length data",
|
533 |
+
)
|
534 |
+
|
535 |
+
parser.add_argument(
|
536 |
+
"--fusion-type",
|
537 |
+
type=str,
|
538 |
+
default="None",
|
539 |
+
help="Type is among ['channel_map', 'daf_1d','aff_1d','iaff_1d','daf_2d','aff_2d','iaff_2d']",
|
540 |
+
)
|
541 |
+
|
542 |
+
parser.add_argument(
|
543 |
+
"--mixup",
|
544 |
+
default=False,
|
545 |
+
action="store_true",
|
546 |
+
help="Enable mixup in finetuning training.",
|
547 |
+
)
|
548 |
+
parser.add_argument(
|
549 |
+
"--text-augment-selection",
|
550 |
+
type=str,
|
551 |
+
default=None,
|
552 |
+
help="For selecting levels of augmented text. Type is among ['all', 'augment_only', 'none']",
|
553 |
+
)
|
554 |
+
|
555 |
+
args = parser.parse_args()
|
556 |
+
|
557 |
+
# If some params are not passed, we use the default values based on model name.
|
558 |
+
default_params = get_default_params(args.amodel)
|
559 |
+
for name, val in default_params.items():
|
560 |
+
if getattr(args, name) is None:
|
561 |
+
setattr(args, name, val)
|
562 |
+
|
563 |
+
return args
|