13 |
14 | 
15 |
16 | 
17 | Nankai University
9 | 10 |
11 |
12 |
13 |
14 |
26 | 27 | 1. [Requirements](#%EF%B8%8F-requirements) 28 | 2. [Datasets](#-datasets) 29 | 3. [Implementary Details](#-implementary-details) 30 | 4. [Train and Test](#%EF%B8%8F-train-and-test) 31 | 5. [Results and Models](#-results-and-models) 32 | 6. [Acknowledgments](#-acknowledgments) 33 | 7. [Citation](#-citation) 34 | 35 |
36 | 37 |
62 | 
70 | 
| 
|
|
|
|
|
101 | | 
| 
|
|
|
|
|
102 | | 
| 
|
|
|
|
|
103 | | 
| 
|
|
|
|
|
104 | |**Params/FLOPs**| 150K/8G|1518K/114G|840K/47G|43090K/2895G|8712K/495G|
105 |
106 | Results of our MAN-tiny/light/base models. Set5 validation set is used below to show the general performance. The visual results of five testsets are provided in the last column.
107 |
108 | | Methods | Params | FLOPs |PSNR/SSIM (x2)|PSNR/SSIM (x3)|PSNR/SSIM (x4)|Results|
109 | |:---------|:---------:|:--------:|:------:|:------:|:------:|:--------:|
110 | | MAN-tiny | 150K | 8.4G | 37.91/0.9603 | 34.23/0.9258 | 32.07/0.8930 | [x2](https://pan.baidu.com/s/1mYkGvAlz0bSZuCVubkpsmg?pwd=mans)/[x3](https://pan.baidu.com/s/1RP5gGu-QPXTkH1NPH7axag?pwd=mans)/[x4](https://pan.baidu.com/s/1u22su2bT4Pq_idVxAnqWdw?pwd=mans) |
111 | | MAN-light| 840K | 47.1G | 38.18/0.9612 | 34.65/0.9292 | 32.50/0.8988 | [x2](https://pan.baidu.com/s/1AVuPa7bsbb3qMQqMSM-IJQ?pwd=mans)/[x3](https://pan.baidu.com/s/1TRL7-Y23JddVOpEhH0ObEQ?pwd=mans)/[x4](https://pan.baidu.com/s/1T2bPZcjFRxAgMxGWtPv-Lw?pwd=mans) |
112 | | MAN+ | 8712K | 495G | 38.44/0.9623 | 34.97/0.9315 | 32.87/0.9030 | [x2](https://pan.baidu.com/s/1pTb3Fob_7MOxMKIdopI0hQ?pwd=mans)/[x3](https://pan.baidu.com/s/1L3HEtcraU8Y9VY-HpCZdfg?pwd=mans)/[x4](https://pan.baidu.com/s/1FCNqht9zi9HecG3ExRdeWQ?pwd=mans) |
113 |
114 | 💖 Acknowledgments
115 | ---
116 |
117 | We would thank [VAN](https://github.com/Visual-Attention-Network/VAN-Classification) and [BasicSR](https://github.com/XPixelGroup/BasicSR) for their enlightening work!
118 |
119 | 🎓 Citation
120 | ---
121 | ```
122 | @inproceedings{wang2024multi,
123 | title={Multi-scale Attention Network for Single Image Super-Resolution},
124 | author={Wang, Yan and Li, Yusen and Wang, Gang and Liu, Xiaoguang},
125 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops},
126 | year={2024}
127 | }
128 | ```
129 |
130 | or
131 |
132 | ```
133 | @article{wang2022multi,
134 | title={Multi-scale Attention Network for Single Image Super-Resolution},
135 | author={Wang, Yan and Li, Yusen and Wang, Gang and Liu, Xiaoguang},
136 | journal={arXiv preprint arXiv:2209.14145},
137 | year={2022}
138 | }
139 | ```
140 |
--------------------------------------------------------------------------------
/LAM/ModelZoo/NN/drln.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | from ..NN import drln_ops as ops
4 | import torch.nn.functional as F
5 |
6 |
7 | def make_model(args, parent=False):
8 | return DRLN(args)
9 |
10 |
11 | class CALayer(nn.Module):
12 | def __init__(self, channel, reduction=16):
13 | super(CALayer, self).__init__()
14 |
15 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
16 |
17 | self.c1 = ops.BasicBlock(channel, channel // reduction, 3, 1, 3, 3)
18 | self.c2 = ops.BasicBlock(channel, channel // reduction, 3, 1, 5, 5)
19 | self.c3 = ops.BasicBlock(channel, channel // reduction, 3, 1, 7, 7)
20 | self.c4 = ops.BasicBlockSig((channel // reduction) * 3, channel, 3, 1, 1)
21 |
22 | def forward(self, x):
23 | y = self.avg_pool(x)
24 | c1 = self.c1(y)
25 | c2 = self.c2(y)
26 | c3 = self.c3(y)
27 | c_out = torch.cat([c1, c2, c3], dim=1)
28 | y = self.c4(c_out)
29 | return x * y
30 |
31 |
32 | class Block(nn.Module):
33 | def __init__(self, in_channels, out_channels, group=1):
34 | super(Block, self).__init__()
35 |
36 | self.r1 = ops.ResidualBlock(in_channels, out_channels)
37 | self.r2 = ops.ResidualBlock(in_channels * 2, out_channels * 2)
38 | self.r3 = ops.ResidualBlock(in_channels * 4, out_channels * 4)
39 | self.g = ops.BasicBlock(in_channels * 8, out_channels, 1, 1, 0)
40 | self.ca = CALayer(in_channels)
41 |
42 | def forward(self, x):
43 | c0 = x
44 |
45 | r1 = self.r1(c0)
46 | c1 = torch.cat([c0, r1], dim=1)
47 |
48 | r2 = self.r2(c1)
49 | c2 = torch.cat([c1, r2], dim=1)
50 |
51 | r3 = self.r3(c2)
52 | c3 = torch.cat([c2, r3], dim=1)
53 |
54 | g = self.g(c3)
55 | out = self.ca(g)
56 | return out
57 |
58 |
59 | class DRLN(nn.Module):
60 | def __init__(self):
61 | super(DRLN, self).__init__()
62 |
63 | # n_resgroups = args.n_resgroups
64 | # n_resblocks = args.n_resblocks
65 | # n_feats = args.n_feats
66 | # kernel_size = 3
67 | # reduction = args.reduction
68 | # scale = args.scale[0]
69 | # act = nn.ReLU(True)
70 |
71 | self.scale = 4
72 | chs = 64
73 |
74 | self.sub_mean = ops.MeanShift((0.4488, 0.4371, 0.4040), sub=True)
75 | self.add_mean = ops.MeanShift((0.4488, 0.4371, 0.4040), sub=False)
76 |
77 | self.head = nn.Conv2d(3, chs, 3, 1, 1)
78 |
79 | self.b1 = Block(chs, chs)
80 | self.b2 = Block(chs, chs)
81 | self.b3 = Block(chs, chs)
82 | self.b4 = Block(chs, chs)
83 | self.b5 = Block(chs, chs)
84 | self.b6 = Block(chs, chs)
85 | self.b7 = Block(chs, chs)
86 | self.b8 = Block(chs, chs)
87 | self.b9 = Block(chs, chs)
88 | self.b10 = Block(chs, chs)
89 | self.b11 = Block(chs, chs)
90 | self.b12 = Block(chs, chs)
91 | self.b13 = Block(chs, chs)
92 | self.b14 = Block(chs, chs)
93 | self.b15 = Block(chs, chs)
94 | self.b16 = Block(chs, chs)
95 | self.b17 = Block(chs, chs)
96 | self.b18 = Block(chs, chs)
97 | self.b19 = Block(chs, chs)
98 | self.b20 = Block(chs, chs)
99 |
100 | self.c1 = ops.BasicBlock(chs * 2, chs, 3, 1, 1)
101 | self.c2 = ops.BasicBlock(chs * 3, chs, 3, 1, 1)
102 | self.c3 = ops.BasicBlock(chs * 4, chs, 3, 1, 1)
103 | self.c4 = ops.BasicBlock(chs * 2, chs, 3, 1, 1)
104 | self.c5 = ops.BasicBlock(chs * 3, chs, 3, 1, 1)
105 | self.c6 = ops.BasicBlock(chs * 4, chs, 3, 1, 1)
106 | self.c7 = ops.BasicBlock(chs * 2, chs, 3, 1, 1)
107 | self.c8 = ops.BasicBlock(chs * 3, chs, 3, 1, 1)
108 | self.c9 = ops.BasicBlock(chs * 4, chs, 3, 1, 1)
109 | self.c10 = ops.BasicBlock(chs * 2, chs, 3, 1, 1)
110 | self.c11 = ops.BasicBlock(chs * 3, chs, 3, 1, 1)
111 | self.c12 = ops.BasicBlock(chs * 4, chs, 3, 1, 1)
112 | self.c13 = ops.BasicBlock(chs * 2, chs, 3, 1, 1)
113 | self.c14 = ops.BasicBlock(chs * 3, chs, 3, 1, 1)
114 | self.c15 = ops.BasicBlock(chs * 4, chs, 3, 1, 1)
115 | self.c16 = ops.BasicBlock(chs * 5, chs, 3, 1, 1)
116 | self.c17 = ops.BasicBlock(chs * 2, chs, 3, 1, 1)
117 | self.c18 = ops.BasicBlock(chs * 3, chs, 3, 1, 1)
118 | self.c19 = ops.BasicBlock(chs * 4, chs, 3, 1, 1)
119 | self.c20 = ops.BasicBlock(chs * 5, chs, 3, 1, 1)
120 |
121 | self.upsample = ops.UpsampleBlock(chs, self.scale, multi_scale=False)
122 | # self.convert = ops.ConvertBlock(chs, chs, 20)
123 | self.tail = nn.Conv2d(chs, 3, 3, 1, 1)
124 |
125 | def forward(self, x):
126 | x = self.sub_mean(x * 255.)
127 | x = self.head(x)
128 | c0 = o0 = x
129 |
130 | b1 = self.b1(o0)
131 | c1 = torch.cat([c0, b1], dim=1)
132 | o1 = self.c1(c1)
133 |
134 | b2 = self.b2(o1)
135 | c2 = torch.cat([c1, b2], dim=1)
136 | o2 = self.c2(c2)
137 |
138 | b3 = self.b3(o2)
139 | c3 = torch.cat([c2, b3], dim=1)
140 | o3 = self.c3(c3)
141 | a1 = o3 + c0
142 |
143 | b4 = self.b4(a1)
144 | c4 = torch.cat([o3, b4], dim=1)
145 | o4 = self.c4(c4)
146 |
147 | b5 = self.b5(a1)
148 | c5 = torch.cat([c4, b5], dim=1)
149 | o5 = self.c5(c5)
150 |
151 | b6 = self.b6(o5)
152 | c6 = torch.cat([c5, b6], dim=1)
153 | o6 = self.c6(c6)
154 | a2 = o6 + a1
155 |
156 | b7 = self.b7(a2)
157 | c7 = torch.cat([o6, b7], dim=1)
158 | o7 = self.c7(c7)
159 |
160 | b8 = self.b8(o7)
161 | c8 = torch.cat([c7, b8], dim=1)
162 | o8 = self.c8(c8)
163 |
164 | b9 = self.b9(o8)
165 | c9 = torch.cat([c8, b9], dim=1)
166 | o9 = self.c9(c9)
167 | a3 = o9 + a2
168 |
169 | b10 = self.b10(a3)
170 | c10 = torch.cat([o9, b10], dim=1)
171 | o10 = self.c10(c10)
172 |
173 | b11 = self.b11(o10)
174 | c11 = torch.cat([c10, b11], dim=1)
175 | o11 = self.c11(c11)
176 |
177 | b12 = self.b12(o11)
178 | c12 = torch.cat([c11, b12], dim=1)
179 | o12 = self.c12(c12)
180 | a4 = o12 + a3
181 |
182 | b13 = self.b13(a4)
183 | c13 = torch.cat([o12, b13], dim=1)
184 | o13 = self.c13(c13)
185 |
186 | b14 = self.b14(o13)
187 | c14 = torch.cat([c13, b14], dim=1)
188 | o14 = self.c14(c14)
189 |
190 | b15 = self.b15(o14)
191 | c15 = torch.cat([c14, b15], dim=1)
192 | o15 = self.c15(c15)
193 |
194 | b16 = self.b16(o15)
195 | c16 = torch.cat([c15, b16], dim=1)
196 | o16 = self.c16(c16)
197 | a5 = o16 + a4
198 |
199 | b17 = self.b17(a5)
200 | c17 = torch.cat([o16, b17], dim=1)
201 | o17 = self.c17(c17)
202 |
203 | b18 = self.b18(o17)
204 | c18 = torch.cat([c17, b18], dim=1)
205 | o18 = self.c18(c18)
206 |
207 | b19 = self.b19(o18)
208 | c19 = torch.cat([c18, b19], dim=1)
209 | o19 = self.c19(c19)
210 |
211 | b20 = self.b20(o19)
212 | c20 = torch.cat([c19, b20], dim=1)
213 | o20 = self.c20(c20)
214 | a6 = o20 + a5
215 |
216 | # c_out = torch.cat([b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14, b15, b16, b17, b18, b19, b20], dim=1)
217 |
218 | # b = self.convert(c_out)
219 | b_out = a6 + x
220 | out = self.upsample(b_out, scale=self.scale)
221 |
222 | out = self.tail(out)
223 | f_out = self.add_mean(out)
224 |
225 | return f_out / 255.
226 |
227 | def load_state_dict(self, state_dict, strict=False):
228 | own_state = self.state_dict()
229 | for name, param in state_dict.items():
230 | if name in own_state:
231 | if isinstance(param, nn.Parameter):
232 | param = param.data
233 | try:
234 | own_state[name].copy_(param)
235 | except Exception:
236 | if name.find('tail') >= 0 or name.find('upsample') >= 0:
237 | print('Replace pre-trained upsampler to new one...')
238 | else:
239 | raise RuntimeError('While copying the parameter named {}, '
240 | 'whose dimensions in the model are {} and '
241 | 'whose dimensions in the checkpoint are {}.'
242 | .format(name, own_state[name].size(), param.size()))
243 | elif strict:
244 | if name.find('tail') == -1:
245 | raise KeyError('unexpected key "{}" in state_dict'
246 | .format(name))
247 |
248 | if strict:
249 | missing = set(own_state.keys()) - set(state_dict.keys())
250 | if len(missing) > 0:
251 | raise KeyError('missing keys in state_dict: "{}"'.format(missing))
252 |
253 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | Apache License
2 | Version 2.0, January 2004
3 | http://www.apache.org/licenses/
4 |
5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6 |
7 | 1. Definitions.
8 |
9 | "License" shall mean the terms and conditions for use, reproduction,
10 | and distribution as defined by Sections 1 through 9 of this document.
11 |
12 | "Licensor" shall mean the copyright owner or entity authorized by
13 | the copyright owner that is granting the License.
14 |
15 | "Legal Entity" shall mean the union of the acting entity and all
16 | other entities that control, are controlled by, or are under common
17 | control with that entity. For the purposes of this definition,
18 | "control" means (i) the power, direct or indirect, to cause the
19 | direction or management of such entity, whether by contract or
20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the
21 | outstanding shares, or (iii) beneficial ownership of such entity.
22 |
23 | "You" (or "Your") shall mean an individual or Legal Entity
24 | exercising permissions granted by this License.
25 |
26 | "Source" form shall mean the preferred form for making modifications,
27 | including but not limited to software source code, documentation
28 | source, and configuration files.
29 |
30 | "Object" form shall mean any form resulting from mechanical
31 | transformation or translation of a Source form, including but
32 | not limited to compiled object code, generated documentation,
33 | and conversions to other media types.
34 |
35 | "Work" shall mean the work of authorship, whether in Source or
36 | Object form, made available under the License, as indicated by a
37 | copyright notice that is included in or attached to the work
38 | (an example is provided in the Appendix below).
39 |
40 | "Derivative Works" shall mean any work, whether in Source or Object
41 | form, that is based on (or derived from) the Work and for which the
42 | editorial revisions, annotations, elaborations, or other modifications
43 | represent, as a whole, an original work of authorship. For the purposes
44 | of this License, Derivative Works shall not include works that remain
45 | separable from, or merely link (or bind by name) to the interfaces of,
46 | the Work and Derivative Works thereof.
47 |
48 | "Contribution" shall mean any work of authorship, including
49 | the original version of the Work and any modifications or additions
50 | to that Work or Derivative Works thereof, that is intentionally
51 | submitted to Licensor for inclusion in the Work by the copyright owner
52 | or by an individual or Legal Entity authorized to submit on behalf of
53 | the copyright owner. For the purposes of this definition, "submitted"
54 | means any form of electronic, verbal, or written communication sent
55 | to the Licensor or its representatives, including but not limited to
56 | communication on electronic mailing lists, source code control systems,
57 | and issue tracking systems that are managed by, or on behalf of, the
58 | Licensor for the purpose of discussing and improving the Work, but
59 | excluding communication that is conspicuously marked or otherwise
60 | designated in writing by the copyright owner as "Not a Contribution."
61 |
62 | "Contributor" shall mean Licensor and any individual or Legal Entity
63 | on behalf of whom a Contribution has been received by Licensor and
64 | subsequently incorporated within the Work.
65 |
66 | 2. Grant of Copyright License. Subject to the terms and conditions of
67 | this License, each Contributor hereby grants to You a perpetual,
68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69 | copyright license to reproduce, prepare Derivative Works of,
70 | publicly display, publicly perform, sublicense, and distribute the
71 | Work and such Derivative Works in Source or Object form.
72 |
73 | 3. Grant of Patent License. Subject to the terms and conditions of
74 | this License, each Contributor hereby grants to You a perpetual,
75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76 | (except as stated in this section) patent license to make, have made,
77 | use, offer to sell, sell, import, and otherwise transfer the Work,
78 | where such license applies only to those patent claims licensable
79 | by such Contributor that are necessarily infringed by their
80 | Contribution(s) alone or by combination of their Contribution(s)
81 | with the Work to which such Contribution(s) was submitted. If You
82 | institute patent litigation against any entity (including a
83 | cross-claim or counterclaim in a lawsuit) alleging that the Work
84 | or a Contribution incorporated within the Work constitutes direct
85 | or contributory patent infringement, then any patent licenses
86 | granted to You under this License for that Work shall terminate
87 | as of the date such litigation is filed.
88 |
89 | 4. Redistribution. You may reproduce and distribute copies of the
90 | Work or Derivative Works thereof in any medium, with or without
91 | modifications, and in Source or Object form, provided that You
92 | meet the following conditions:
93 |
94 | (a) You must give any other recipients of the Work or
95 | Derivative Works a copy of this License; and
96 |
97 | (b) You must cause any modified files to carry prominent notices
98 | stating that You changed the files; and
99 |
100 | (c) You must retain, in the Source form of any Derivative Works
101 | that You distribute, all copyright, patent, trademark, and
102 | attribution notices from the Source form of the Work,
103 | excluding those notices that do not pertain to any part of
104 | the Derivative Works; and
105 |
106 | (d) If the Work includes a "NOTICE" text file as part of its
107 | distribution, then any Derivative Works that You distribute must
108 | include a readable copy of the attribution notices contained
109 | within such NOTICE file, excluding those notices that do not
110 | pertain to any part of the Derivative Works, in at least one
111 | of the following places: within a NOTICE text file distributed
112 | as part of the Derivative Works; within the Source form or
113 | documentation, if provided along with the Derivative Works; or,
114 | within a display generated by the Derivative Works, if and
115 | wherever such third-party notices normally appear. The contents
116 | of the NOTICE file are for informational purposes only and
117 | do not modify the License. You may add Your own attribution
118 | notices within Derivative Works that You distribute, alongside
119 | or as an addendum to the NOTICE text from the Work, provided
120 | that such additional attribution notices cannot be construed
121 | as modifying the License.
122 |
123 | You may add Your own copyright statement to Your modifications and
124 | may provide additional or different license terms and conditions
125 | for use, reproduction, or distribution of Your modifications, or
126 | for any such Derivative Works as a whole, provided Your use,
127 | reproduction, and distribution of the Work otherwise complies with
128 | the conditions stated in this License.
129 |
130 | 5. Submission of Contributions. Unless You explicitly state otherwise,
131 | any Contribution intentionally submitted for inclusion in the Work
132 | by You to the Licensor shall be under the terms and conditions of
133 | this License, without any additional terms or conditions.
134 | Notwithstanding the above, nothing herein shall supersede or modify
135 | the terms of any separate license agreement you may have executed
136 | with Licensor regarding such Contributions.
137 |
138 | 6. Trademarks. This License does not grant permission to use the trade
139 | names, trademarks, service marks, or product names of the Licensor,
140 | except as required for reasonable and customary use in describing the
141 | origin of the Work and reproducing the content of the NOTICE file.
142 |
143 | 7. Disclaimer of Warranty. Unless required by applicable law or
144 | agreed to in writing, Licensor provides the Work (and each
145 | Contributor provides its Contributions) on an "AS IS" BASIS,
146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 | implied, including, without limitation, any warranties or conditions
148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 | PARTICULAR PURPOSE. You are solely responsible for determining the
150 | appropriateness of using or redistributing the Work and assume any
151 | risks associated with Your exercise of permissions under this License.
152 |
153 | 8. Limitation of Liability. In no event and under no legal theory,
154 | whether in tort (including negligence), contract, or otherwise,
155 | unless required by applicable law (such as deliberate and grossly
156 | negligent acts) or agreed to in writing, shall any Contributor be
157 | liable to You for damages, including any direct, indirect, special,
158 | incidental, or consequential damages of any character arising as a
159 | result of this License or out of the use or inability to use the
160 | Work (including but not limited to damages for loss of goodwill,
161 | work stoppage, computer failure or malfunction, or any and all
162 | other commercial damages or losses), even if such Contributor
163 | has been advised of the possibility of such damages.
164 |
165 | 9. Accepting Warranty or Additional Liability. While redistributing
166 | the Work or Derivative Works thereof, You may choose to offer,
167 | and charge a fee for, acceptance of support, warranty, indemnity,
168 | or other liability obligations and/or rights consistent with this
169 | License. However, in accepting such obligations, You may act only
170 | on Your own behalf and on Your sole responsibility, not on behalf
171 | of any other Contributor, and only if You agree to indemnify,
172 | defend, and hold each Contributor harmless for any liability
173 | incurred by, or claims asserted against, such Contributor by reason
174 | of your accepting any such warranty or additional liability.
175 |
176 | END OF TERMS AND CONDITIONS
177 |
178 | APPENDIX: How to apply the Apache License to your work.
179 |
180 | To apply the Apache License to your work, attach the following
181 | boilerplate notice, with the fields enclosed by brackets "[]"
182 | replaced with your own identifying information. (Don't include
183 | the brackets!) The text should be enclosed in the appropriate
184 | comment syntax for the file format. We also recommend that a
185 | file or class name and description of purpose be included on the
186 | same "printed page" as the copyright notice for easier
187 | identification within third-party archives.
188 |
189 | Copyright MAN Authors
190 |
191 | Licensed under the Apache License, Version 2.0 (the "License");
192 | you may not use this file except in compliance with the License.
193 | You may obtain a copy of the License at
194 |
195 | http://www.apache.org/licenses/LICENSE-2.0
196 |
197 | Unless required by applicable law or agreed to in writing, software
198 | distributed under the License is distributed on an "AS IS" BASIS,
199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 | See the License for the specific language governing permissions and
201 | limitations under the License.
202 |
--------------------------------------------------------------------------------
/LAM/ModelZoo/NN/common.py:
--------------------------------------------------------------------------------
1 | import math
2 |
3 | import numpy as np
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 |
8 | from torch.autograd import Variable
9 |
10 |
11 | def default_conv(in_channels, out_channels, kernel_size, bias=True):
12 | return nn.Conv2d(
13 | in_channels, out_channels, kernel_size,
14 | padding=(kernel_size//2), bias=bias)
15 |
16 |
17 | class MeanShift(nn.Conv2d):
18 | def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
19 | super(MeanShift, self).__init__(3, 3, kernel_size=1)
20 | std = torch.Tensor(rgb_std)
21 | self.weight.data = torch.eye(3).view(3, 3, 1, 1)
22 | self.weight.data.div_(std.view(3, 1, 1, 1))
23 | self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean)
24 | self.bias.data.div_(std)
25 | self.requires_grad = False
26 |
27 |
28 | class DownBlock(nn.Module):
29 | def __init__(self, opt, scale, nFeat=None, in_channels=None, out_channels=None):
30 | super(DownBlock, self).__init__()
31 | negval = opt.negval
32 |
33 | if nFeat is None:
34 | nFeat = opt.n_feats
35 |
36 | if in_channels is None:
37 | in_channels = opt.n_colors
38 |
39 | if out_channels is None:
40 | out_channels = opt.n_colors
41 |
42 | dual_block = [
43 | nn.Sequential(
44 | nn.Conv2d(in_channels, nFeat, kernel_size=3, stride=2, padding=1, bias=False),
45 | nn.LeakyReLU(negative_slope=negval, inplace=True)
46 | )
47 | ]
48 |
49 | for _ in range(1, int(np.log2(scale))):
50 | dual_block.append(
51 | nn.Sequential(
52 | nn.Conv2d(nFeat, nFeat, kernel_size=3, stride=2, padding=1, bias=False),
53 | nn.LeakyReLU(negative_slope=negval, inplace=True)
54 | )
55 | )
56 |
57 | dual_block.append(nn.Conv2d(nFeat, out_channels, kernel_size=3, stride=1, padding=1, bias=False))
58 |
59 | self.dual_module = nn.Sequential(*dual_block)
60 |
61 | def forward(self, x):
62 | x = self.dual_module(x)
63 | return x
64 |
65 |
66 | class BasicBlock(nn.Sequential):
67 | def __init__(
68 | self, in_channels, out_channels, kernel_size, stride=1, bias=False,
69 | bn=True, act=nn.ReLU(True)):
70 |
71 | m = [nn.Conv2d(
72 | in_channels, out_channels, kernel_size,
73 | padding=(kernel_size//2), stride=stride, bias=bias)
74 | ]
75 | if bn: m.append(nn.BatchNorm2d(out_channels))
76 | if act is not None: m.append(act)
77 | super(BasicBlock, self).__init__(*m)
78 |
79 |
80 | class ResBlock(nn.Module):
81 | def __init__(
82 | self, conv, n_feat, kernel_size,
83 | bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
84 |
85 | super(ResBlock, self).__init__()
86 | m = []
87 | for i in range(2):
88 | m.append(conv(n_feat, n_feat, kernel_size, bias=bias))
89 | if bn: m.append(nn.BatchNorm2d(n_feat))
90 | if i == 0: m.append(act)
91 |
92 | self.body = nn.Sequential(*m)
93 | self.res_scale = res_scale
94 |
95 | def forward(self, x):
96 | res = self.body(x).mul(self.res_scale)
97 | res += x
98 |
99 | return res
100 |
101 |
102 | class Upsampler(nn.Sequential):
103 | def __init__(self, conv, scale, n_feat, bn=False, act=False, bias=True):
104 |
105 | m = []
106 | if (scale & (scale - 1)) == 0: # Is scale = 2^n?
107 | for _ in range(int(math.log(scale, 2))):
108 | m.append(conv(n_feat, 4 * n_feat, 3, bias))
109 | m.append(nn.PixelShuffle(2))
110 | if bn: m.append(nn.BatchNorm2d(n_feat))
111 | if act: m.append(act())
112 | elif scale == 3:
113 | m.append(conv(n_feat, 9 * n_feat, 3, bias))
114 | m.append(nn.PixelShuffle(3))
115 | if bn: m.append(nn.BatchNorm2d(n_feat))
116 | if act: m.append(act())
117 | else:
118 | raise NotImplementedError
119 |
120 | super(Upsampler, self).__init__(*m)
121 |
122 |
123 | ## add SELayer
124 | class SELayer(nn.Module):
125 | def __init__(self, channel, reduction=16):
126 | super(SELayer, self).__init__()
127 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
128 | self.conv_du = nn.Sequential(
129 | nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
130 | nn.ReLU(inplace=True),
131 | nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
132 | nn.Sigmoid()
133 | )
134 |
135 | def forward(self, x):
136 | y = self.avg_pool(x)
137 | y = self.conv_du(y)
138 | return x * y
139 |
140 |
141 | ## add SEResBlock
142 | class SEResBlock(nn.Module):
143 | def __init__(
144 | self, conv, n_feat, kernel_size, reduction,
145 | bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
146 |
147 | super(SEResBlock, self).__init__()
148 | modules_body = []
149 | for i in range(2):
150 | modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
151 | if bn: modules_body.append(nn.BatchNorm2d(n_feat))
152 | if i == 0: modules_body.append(act)
153 | modules_body.append(SELayer(n_feat, reduction))
154 | self.body = nn.Sequential(*modules_body)
155 | self.res_scale = res_scale
156 |
157 | def forward(self, x):
158 | res = self.body(x)
159 | #res = self.body(x).mul(self.res_scale)
160 | res += x
161 |
162 | return res
163 |
164 |
165 | ## Channel Attention (CA) Layer
166 | class CALayer(nn.Module):
167 | def __init__(self, channel, reduction=16):
168 | super(CALayer, self).__init__()
169 | # global average pooling: feature --> point
170 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
171 | # feature channel downscale and upscale --> channel weight
172 | self.conv_du = nn.Sequential(
173 | nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
174 | nn.ReLU(inplace=True),
175 | nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
176 | nn.Sigmoid()
177 | )
178 |
179 | def forward(self, x):
180 | y = self.avg_pool(x)
181 | y = self.conv_du(y)
182 | return x * y
183 |
184 |
185 | ## Residual Channel Attention Block (RCAB)
186 | class RCAB(nn.Module):
187 | def __init__(self, conv, n_feat, kernel_size, reduction=16, bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
188 | super(RCAB, self).__init__()
189 | modules_body = []
190 | for i in range(2):
191 | modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
192 | if bn: modules_body.append(nn.BatchNorm2d(n_feat))
193 | if i == 0: modules_body.append(act)
194 | modules_body.append(CALayer(n_feat, reduction))
195 | self.body = nn.Sequential(*modules_body)
196 | self.res_scale = res_scale
197 |
198 | def forward(self, x):
199 | res = self.body(x)
200 | res += x
201 | return res
202 |
203 |
204 | # add NonLocalBlock2D
205 | # reference: https://github.com/AlexHex7/Non-local_pytorch/blob/master/lib/non_local_simple_version.py
206 | class NonLocalBlock2D(nn.Module):
207 | def __init__(self, in_channels, inter_channels):
208 | super(NonLocalBlock2D, self).__init__()
209 |
210 | self.in_channels = in_channels
211 | self.inter_channels = inter_channels
212 |
213 | self.g = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1,
214 | padding=0)
215 |
216 | self.W = nn.Conv2d(in_channels=self.inter_channels, out_channels=self.in_channels, kernel_size=1, stride=1,
217 | padding=0)
218 | # for pytorch 0.3.1
219 | # nn.init.constant(self.W.weight, 0)
220 | # nn.init.constant(self.W.bias, 0)
221 | # for pytorch 0.4.0
222 | nn.init.constant_(self.W.weight, 0)
223 | nn.init.constant_(self.W.bias, 0)
224 | self.theta = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1,
225 | padding=0)
226 |
227 | self.phi = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1,
228 | padding=0)
229 |
230 | def forward(self, x):
231 | batch_size = x.size(0)
232 |
233 | g_x = self.g(x).view(batch_size, self.inter_channels, -1)
234 |
235 | g_x = g_x.permute(0, 2, 1)
236 |
237 | theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
238 |
239 | theta_x = theta_x.permute(0, 2, 1)
240 |
241 | phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
242 |
243 | f = torch.matmul(theta_x, phi_x)
244 |
245 | f_div_C = F.softmax(f, dim=1)
246 |
247 | y = torch.matmul(f_div_C, g_x)
248 |
249 | y = y.permute(0, 2, 1).contiguous()
250 |
251 | y = y.view(batch_size, self.inter_channels, *x.size()[2:])
252 | W_y = self.W(y)
253 | z = W_y + x
254 |
255 | return z
256 |
257 |
258 | ## define trunk branch
259 | class TrunkBranch(nn.Module):
260 | def __init__(
261 | self, conv, n_feat, kernel_size,
262 | bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
263 | super(TrunkBranch, self).__init__()
264 | modules_body = []
265 | for i in range(2):
266 | modules_body.append(
267 | ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
268 | self.body = nn.Sequential(*modules_body)
269 |
270 | def forward(self, x):
271 | tx = self.body(x)
272 |
273 | return tx
274 |
275 |
276 | ## define mask branch
277 | class MaskBranchDownUp(nn.Module):
278 | def __init__(
279 | self, conv, n_feat, kernel_size,
280 | bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
281 | super(MaskBranchDownUp, self).__init__()
282 |
283 | MB_RB1 = []
284 | MB_RB1.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
285 |
286 | MB_Down = []
287 | MB_Down.append(nn.Conv2d(n_feat, n_feat, 3, stride=2, padding=1))
288 |
289 | MB_RB2 = []
290 | for i in range(2):
291 | MB_RB2.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
292 |
293 | MB_Up = []
294 | MB_Up.append(nn.ConvTranspose2d(n_feat, n_feat, 6, stride=2, padding=2))
295 |
296 | MB_RB3 = []
297 | MB_RB3.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
298 |
299 | MB_1x1conv = []
300 | MB_1x1conv.append(nn.Conv2d(n_feat, n_feat, 1, padding=0, bias=True))
301 |
302 | MB_sigmoid = []
303 | MB_sigmoid.append(nn.Sigmoid())
304 |
305 | self.MB_RB1 = nn.Sequential(*MB_RB1)
306 | self.MB_Down = nn.Sequential(*MB_Down)
307 | self.MB_RB2 = nn.Sequential(*MB_RB2)
308 | self.MB_Up = nn.Sequential(*MB_Up)
309 | self.MB_RB3 = nn.Sequential(*MB_RB3)
310 | self.MB_1x1conv = nn.Sequential(*MB_1x1conv)
311 | self.MB_sigmoid = nn.Sequential(*MB_sigmoid)
312 |
313 | def forward(self, x):
314 | x_RB1 = self.MB_RB1(x)
315 | x_Down = self.MB_Down(x_RB1)
316 | x_RB2 = self.MB_RB2(x_Down)
317 | x_Up = self.MB_Up(x_RB2)
318 | x_preRB3 = x_RB1 + x_Up
319 | x_RB3 = self.MB_RB3(x_preRB3)
320 | x_1x1 = self.MB_1x1conv(x_RB3)
321 | mx = self.MB_sigmoid(x_1x1)
322 |
323 | return mx
324 |
325 |
326 | ## define nonlocal mask branch
327 | class NLMaskBranchDownUp(nn.Module):
328 | def __init__(
329 | self, conv, n_feat, kernel_size,
330 | bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
331 | super(NLMaskBranchDownUp, self).__init__()
332 |
333 | MB_RB1 = []
334 | MB_RB1.append(NonLocalBlock2D(n_feat, n_feat // 2))
335 | MB_RB1.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
336 |
337 | MB_Down = []
338 | MB_Down.append(nn.Conv2d(n_feat, n_feat, 3, stride=2, padding=1))
339 |
340 | MB_RB2 = []
341 | for i in range(2):
342 | MB_RB2.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
343 |
344 | MB_Up = []
345 | MB_Up.append(nn.ConvTranspose2d(n_feat, n_feat, 6, stride=2, padding=2))
346 |
347 | MB_RB3 = []
348 | MB_RB3.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
349 |
350 | MB_1x1conv = []
351 | MB_1x1conv.append(nn.Conv2d(n_feat, n_feat, 1, padding=0, bias=True))
352 |
353 | MB_sigmoid = []
354 | MB_sigmoid.append(nn.Sigmoid())
355 |
356 | self.MB_RB1 = nn.Sequential(*MB_RB1)
357 | self.MB_Down = nn.Sequential(*MB_Down)
358 | self.MB_RB2 = nn.Sequential(*MB_RB2)
359 | self.MB_Up = nn.Sequential(*MB_Up)
360 | self.MB_RB3 = nn.Sequential(*MB_RB3)
361 | self.MB_1x1conv = nn.Sequential(*MB_1x1conv)
362 | self.MB_sigmoid = nn.Sequential(*MB_sigmoid)
363 |
364 | def forward(self, x):
365 | x_RB1 = self.MB_RB1(x)
366 | x_Down = self.MB_Down(x_RB1)
367 | x_RB2 = self.MB_RB2(x_Down)
368 | x_Up = self.MB_Up(x_RB2)
369 | x_preRB3 = x_RB1 + x_Up
370 | x_RB3 = self.MB_RB3(x_preRB3)
371 | x_1x1 = self.MB_1x1conv(x_RB3)
372 | mx = self.MB_sigmoid(x_1x1)
373 |
374 | return mx
375 |
376 |
377 | ## define residual attention module
378 | class ResAttModuleDownUpPlus(nn.Module):
379 | def __init__(
380 | self, conv, n_feat, kernel_size,
381 | bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
382 | super(ResAttModuleDownUpPlus, self).__init__()
383 | RA_RB1 = []
384 | RA_RB1.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
385 | RA_TB = []
386 | RA_TB.append(TrunkBranch(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
387 | RA_MB = []
388 | RA_MB.append(MaskBranchDownUp(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
389 | RA_tail = []
390 | for i in range(2):
391 | RA_tail.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
392 |
393 | self.RA_RB1 = nn.Sequential(*RA_RB1)
394 | self.RA_TB = nn.Sequential(*RA_TB)
395 | self.RA_MB = nn.Sequential(*RA_MB)
396 | self.RA_tail = nn.Sequential(*RA_tail)
397 |
398 | def forward(self, input):
399 | RA_RB1_x = self.RA_RB1(input)
400 | tx = self.RA_TB(RA_RB1_x)
401 | mx = self.RA_MB(RA_RB1_x)
402 | txmx = tx * mx
403 | hx = txmx + RA_RB1_x
404 | hx = self.RA_tail(hx)
405 |
406 | return hx
407 |
408 |
409 | ## define nonlocal residual attention module
410 | class NLResAttModuleDownUpPlus(nn.Module):
411 | def __init__(
412 | self, conv, n_feat, kernel_size,
413 | bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
414 | super(NLResAttModuleDownUpPlus, self).__init__()
415 | RA_RB1 = []
416 | RA_RB1.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
417 | RA_TB = []
418 | RA_TB.append(TrunkBranch(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
419 | RA_MB = []
420 | RA_MB.append(NLMaskBranchDownUp(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
421 | RA_tail = []
422 | for i in range(2):
423 | RA_tail.append(ResBlock(conv, n_feat, kernel_size, bias=True, bn=False, act=nn.ReLU(True), res_scale=1))
424 |
425 | self.RA_RB1 = nn.Sequential(*RA_RB1)
426 | self.RA_TB = nn.Sequential(*RA_TB)
427 | self.RA_MB = nn.Sequential(*RA_MB)
428 | self.RA_tail = nn.Sequential(*RA_tail)
429 |
430 | def forward(self, input):
431 | RA_RB1_x = self.RA_RB1(input)
432 | tx = self.RA_TB(RA_RB1_x)
433 | mx = self.RA_MB(RA_RB1_x)
434 | txmx = tx * mx
435 | hx = txmx + RA_RB1_x
436 | hx = self.RA_tail(hx)
437 |
438 | return hx
--------------------------------------------------------------------------------
/LAM/ModelZoo/NN/man.py:
--------------------------------------------------------------------------------
1 | # -*- coding: utf-8 -*-
2 | import math
3 |
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 |
8 | from basicsr.utils.registry import ARCH_REGISTRY
9 |
10 | #LKA from VAN (https://github.com/Visual-Attention-Network)
11 | class LKA(nn.Module):
12 | def __init__(self, dim):
13 | super().__init__()
14 | self.conv0 = nn.Conv2d(dim, dim, 7, padding=7//2, groups=dim)
15 | self.conv_spatial = nn.Conv2d(dim, dim, 9, stride=1, padding=((9//2)*4), groups=dim, dilation=4)
16 | self.conv1 = nn.Conv2d(dim, dim, 1)
17 |
18 | def forward(self, x):
19 | u = x.clone()
20 | attn = self.conv0(x)
21 | attn = self.conv_spatial(attn)
22 | attn = self.conv1(attn)
23 |
24 | return u * attn
25 |
26 | class Attention(nn.Module):
27 | def __init__(self, n_feats):
28 | super().__init__()
29 |
30 | self.norm = LayerNorm(n_feats, data_format='channels_first')
31 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
32 |
33 | self.proj_1 = nn.Conv2d(n_feats, n_feats, 1)
34 | self.spatial_gating_unit = LKA(n_feats)
35 | self.proj_2 = nn.Conv2d(n_feats, n_feats, 1)
36 |
37 | def forward(self, x):
38 | shorcut = x.clone()
39 | x = self.proj_1(self.norm(x))
40 | x = self.spatial_gating_unit(x)
41 | x = self.proj_2(x)
42 | x = x*self.scale + shorcut
43 | return x
44 | #----------------------------------------------------------------------------------------------------------------
45 |
46 | class MLP(nn.Module):
47 | def __init__(self, n_feats):
48 | super().__init__()
49 |
50 | self.norm = LayerNorm(n_feats, data_format='channels_first')
51 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
52 |
53 | i_feats = 2*n_feats
54 |
55 | self.fc1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)
56 | self.act = nn.GELU()
57 | self.fc2 = nn.Conv2d(i_feats, n_feats, 1, 1, 0)
58 |
59 | def forward(self, x):
60 | shortcut = x.clone()
61 | x = self.norm(x)
62 | x = self.fc1(x)
63 | x = self.act(x)
64 | x = self.fc2(x)
65 |
66 | return x*self.scale + shortcut
67 |
68 |
69 | class CFF(nn.Module):
70 | def __init__(self, n_feats, drop=0.0, k=2, squeeze_factor= 15, attn ='GLKA'):
71 | super().__init__()
72 | i_feats =n_feats*2
73 |
74 | self.Conv1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)
75 | self.DWConv1 = nn.Sequential(
76 | nn.Conv2d(i_feats, i_feats, 7, 1, 7//2, groups= n_feats),
77 | nn.GELU())
78 | self.Conv2 = nn.Conv2d(i_feats, n_feats, 1, 1, 0)
79 |
80 | self.norm = LayerNorm(n_feats, data_format='channels_first')
81 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
82 |
83 | def forward(self, x):
84 | shortcut = x.clone()
85 |
86 | #Ghost Expand
87 | x = self.Conv1(self.norm(x))
88 | x = self.DWConv1(x)
89 | x = self.Conv2(x)
90 |
91 | return x*self.scale + shortcut
92 |
93 |
94 | class SimpleGate(nn.Module):
95 | def __init__(self, n_feats):
96 | super().__init__()
97 | i_feats =n_feats*2
98 |
99 | self.Conv1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)
100 | #self.DWConv1 = nn.Conv2d(n_feats, n_feats, 7, 1, 7//2, groups= n_feats)
101 | self.Conv2 = nn.Conv2d(n_feats, n_feats, 1, 1, 0)
102 |
103 | self.norm = LayerNorm(n_feats, data_format='channels_first')
104 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
105 |
106 | def forward(self, x):
107 | shortcut = x.clone()
108 |
109 | #Ghost Expand
110 | x = self.Conv1(self.norm(x))
111 | a, x = torch.chunk(x, 2, dim=1)
112 | x = x*a #self.DWConv1(a)
113 | x = self.Conv2(x)
114 |
115 | return x*self.scale + shortcut
116 | #-----------------------------------------------------------------------------------------------------------------
117 | #RCAN-style
118 | class RCBv6(nn.Module):
119 | def __init__(
120 | self, n_feats, k, lk=7, res_scale=1.0, style ='X', act = nn.SiLU(), deploy=False):
121 | super().__init__()
122 | self.LKA = nn.Sequential(
123 | nn.Conv2d(n_feats, n_feats, 5, 1, lk//2, groups= n_feats),
124 | nn.Conv2d(n_feats, n_feats, 7, stride=1, padding=9, groups=n_feats, dilation=3),
125 | nn.Conv2d(n_feats, n_feats, 1, 1, 0),
126 | nn.Sigmoid())
127 |
128 | #self.LFE2 = LFEv3(n_feats, attn ='CA')
129 |
130 | self.LFE = nn.Sequential(
131 | nn.Conv2d(n_feats,n_feats,3,1,1),
132 | nn.GELU(),
133 | nn.Conv2d(n_feats,n_feats,3,1,1))
134 |
135 |
136 | def forward(self, x, pre_attn=None, RAA=None):
137 | shortcut = x.clone()
138 | x = self.LFE(x)
139 |
140 | x = self.LKA(x)*x
141 |
142 | return x + shortcut
143 |
144 | #-----------------------------------------------------------------------------------------------------------------
145 |
146 |
147 | class MLKA_Ablation(nn.Module):
148 | def __init__(self, n_feats, k=2, squeeze_factor=15):
149 | super().__init__()
150 | i_feats = 2*n_feats
151 |
152 | self.n_feats= n_feats
153 | self.i_feats = i_feats
154 |
155 | self.norm = LayerNorm(n_feats, data_format='channels_first')
156 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
157 |
158 | k = 2
159 |
160 | #Multiscale Large Kernel Attention
161 | self.LKA7 = nn.Sequential(
162 | nn.Conv2d(n_feats//k, n_feats//k, 7, 1, 7//2, groups= n_feats//k),
163 | nn.Conv2d(n_feats//k, n_feats//k, 9, stride=1, padding=(9//2)*4, groups=n_feats//k, dilation=4),
164 | nn.Conv2d(n_feats//k, n_feats//k, 1, 1, 0))
165 | self.LKA5 = nn.Sequential(
166 | nn.Conv2d(n_feats//k, n_feats//k, 5, 1, 5//2, groups= n_feats//k),
167 | nn.Conv2d(n_feats//k, n_feats//k, 7, stride=1, padding=(7//2)*3, groups=n_feats//k, dilation=3),
168 | nn.Conv2d(n_feats//k, n_feats//k, 1, 1, 0))
169 | '''self.LKA3 = nn.Sequential(
170 | nn.Conv2d(n_feats//k, n_feats//k, 3, 1, 1, groups= n_feats//k),
171 | nn.Conv2d(n_feats//k, n_feats//k, 5, stride=1, padding=(5//2)*2, groups=n_feats//k, dilation=2),
172 | nn.Conv2d(n_feats//k, n_feats//k, 1, 1, 0))'''
173 |
174 | #self.X3 = nn.Conv2d(n_feats//k, n_feats//k, 3, 1, 1, groups= n_feats//k)
175 | self.X5 = nn.Conv2d(n_feats//k, n_feats//k, 5, 1, 5//2, groups= n_feats//k)
176 | self.X7 = nn.Conv2d(n_feats//k, n_feats//k, 7, 1, 7//2, groups= n_feats//k)
177 |
178 | self.proj_first = nn.Sequential(
179 | nn.Conv2d(n_feats, i_feats, 1, 1, 0))
180 |
181 | self.proj_last = nn.Sequential(
182 | nn.Conv2d(n_feats, n_feats, 1, 1, 0))
183 |
184 |
185 | def forward(self, x, pre_attn=None, RAA=None):
186 | shortcut = x.clone()
187 |
188 | x = self.norm(x)
189 |
190 | x = self.proj_first(x)
191 |
192 | a, x = torch.chunk(x, 2, dim=1)
193 |
194 | #u_1, u_2, u_3= torch.chunk(u, 3, dim=1)
195 | a_1, a_2 = torch.chunk(a, 2, dim=1)
196 |
197 | a = torch.cat([self.LKA7(a_1)*self.X7(a_1), self.LKA5(a_2)*self.X5(a_2)], dim=1)
198 |
199 | x = self.proj_last(x*a)*self.scale + shortcut
200 |
201 | return x
202 | #-----------------------------------------------------------------------------------------------------------------
203 |
204 |
205 | class LayerNorm(nn.Module):
206 | r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
207 | The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
208 | shape (batch_size, height, width, channels) while channels_first corresponds to inputs
209 | with shape (batch_size, channels, height, width).
210 | """
211 | def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
212 | super().__init__()
213 | self.weight = nn.Parameter(torch.ones(normalized_shape))
214 | self.bias = nn.Parameter(torch.zeros(normalized_shape))
215 | self.eps = eps
216 | self.data_format = data_format
217 | if self.data_format not in ["channels_last", "channels_first"]:
218 | raise NotImplementedError
219 | self.normalized_shape = (normalized_shape, )
220 |
221 | def forward(self, x):
222 | if self.data_format == "channels_last":
223 | return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
224 | elif self.data_format == "channels_first":
225 | u = x.mean(1, keepdim=True)
226 | s = (x - u).pow(2).mean(1, keepdim=True)
227 | x = (x - u) / torch.sqrt(s + self.eps)
228 | x = self.weight[:, None, None] * x + self.bias[:, None, None]
229 | return x
230 |
231 |
232 |
233 | class SGAB(nn.Module):
234 | def __init__(self, n_feats, drop=0.0, k=2, squeeze_factor= 15, attn ='GLKA'):
235 | super().__init__()
236 | i_feats =n_feats*2
237 |
238 | self.Conv1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)
239 | self.DWConv1 = nn.Conv2d(n_feats, n_feats, 7, 1, 7//2, groups= n_feats)
240 | self.Conv2 = nn.Conv2d(n_feats, n_feats, 1, 1, 0)
241 |
242 | self.norm = LayerNorm(n_feats, data_format='channels_first')
243 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
244 |
245 | def forward(self, x):
246 | shortcut = x.clone()
247 |
248 | #Ghost Expand
249 | x = self.Conv1(self.norm(x))
250 | a, x = torch.chunk(x, 2, dim=1)
251 | x = x*self.DWConv1(a)
252 | x = self.Conv2(x)
253 |
254 | return x*self.scale + shortcut
255 |
256 | class GroupGLKA(nn.Module):
257 | def __init__(self, n_feats, k=2, squeeze_factor=15):
258 | super().__init__()
259 | i_feats = 2*n_feats
260 |
261 | self.n_feats= n_feats
262 | self.i_feats = i_feats
263 |
264 | self.norm = LayerNorm(n_feats, data_format='channels_first')
265 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
266 |
267 | #Multiscale Large Kernel Attention
268 | self.LKA7 = nn.Sequential(
269 | nn.Conv2d(n_feats//3, n_feats//3, 7, 1, 7//2, groups= n_feats//3),
270 | nn.Conv2d(n_feats//3, n_feats//3, 9, stride=1, padding=(9//2)*4, groups=n_feats//3, dilation=4),
271 | nn.Conv2d(n_feats//3, n_feats//3, 1, 1, 0))
272 | self.LKA5 = nn.Sequential(
273 | nn.Conv2d(n_feats//3, n_feats//3, 5, 1, 5//2, groups= n_feats//3),
274 | nn.Conv2d(n_feats//3, n_feats//3, 7, stride=1, padding=(7//2)*3, groups=n_feats//3, dilation=3),
275 | nn.Conv2d(n_feats//3, n_feats//3, 1, 1, 0))
276 | self.LKA3 = nn.Sequential(
277 | nn.Conv2d(n_feats//3, n_feats//3, 3, 1, 1, groups= n_feats//3),
278 | nn.Conv2d(n_feats//3, n_feats//3, 5, stride=1, padding=(5//2)*2, groups=n_feats//3, dilation=2),
279 | nn.Conv2d(n_feats//3, n_feats//3, 1, 1, 0))
280 |
281 | self.X3 = nn.Conv2d(n_feats//3, n_feats//3, 3, 1, 1, groups= n_feats//3)
282 | self.X5 = nn.Conv2d(n_feats//3, n_feats//3, 5, 1, 5//2, groups= n_feats//3)
283 | self.X7 = nn.Conv2d(n_feats//3, n_feats//3, 7, 1, 7//2, groups= n_feats//3)
284 |
285 | self.proj_first = nn.Sequential(
286 | nn.Conv2d(n_feats, i_feats, 1, 1, 0))
287 |
288 | self.proj_last = nn.Sequential(
289 | nn.Conv2d(n_feats, n_feats, 1, 1, 0))
290 |
291 |
292 | def forward(self, x, pre_attn=None, RAA=None):
293 | shortcut = x.clone()
294 |
295 | x = self.norm(x)
296 |
297 | x = self.proj_first(x)
298 |
299 | a, x = torch.chunk(x, 2, dim=1)
300 |
301 | a_1, a_2, a_3= torch.chunk(a, 3, dim=1)
302 |
303 | a = torch.cat([self.LKA3(a_1)*self.X3(a_1), self.LKA5(a_2)*self.X5(a_2), self.LKA7(a_3)*self.X7(a_3)], dim=1)
304 |
305 | x = self.proj_last(x*a)*self.scale + shortcut
306 |
307 | return x
308 |
309 |
310 | # MAB
311 | class MAB(nn.Module):
312 | def __init__(
313 | self, n_feats):
314 | super().__init__()
315 |
316 | self.LKA = GroupGLKA(n_feats)
317 |
318 | self.LFE = SGAB(n_feats)
319 |
320 | def forward(self, x, pre_attn=None, RAA=None):
321 | #large kernel attention
322 | x = self.LKA(x)
323 |
324 | #local feature extraction
325 | x = self.LFE(x)
326 |
327 | return x
328 |
329 | class LKAT(nn.Module):
330 | def __init__(self, n_feats):
331 | super().__init__()
332 |
333 | #self.norm = LayerNorm(n_feats, data_format='channels_first')
334 | #self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
335 |
336 | self.conv0 = nn.Sequential(
337 | nn.Conv2d(n_feats, n_feats, 1, 1, 0),
338 | nn.GELU())
339 |
340 | self.att = nn.Sequential(
341 | nn.Conv2d(n_feats, n_feats, 7, 1, 7//2, groups= n_feats),
342 | nn.Conv2d(n_feats, n_feats, 9, stride=1, padding=(9//2)*3, groups=n_feats, dilation=3),
343 | nn.Conv2d(n_feats, n_feats, 1, 1, 0))
344 |
345 | self.conv1 = nn.Conv2d(n_feats, n_feats, 1, 1, 0)
346 |
347 | def forward(self, x):
348 | x = self.conv0(x)
349 | x = x*self.att(x)
350 | x = self.conv1(x)
351 | return x
352 |
353 | class ResGroup(nn.Module):
354 | def __init__(self, n_resblocks, n_feats, res_scale=1.0):
355 | super(ResGroup, self).__init__()
356 | self.body = nn.ModuleList([
357 | MAB(n_feats) \
358 | for _ in range(n_resblocks)])
359 |
360 | self.body_t = LKAT(n_feats)
361 |
362 | def forward(self, x):
363 | res = x.clone()
364 |
365 | for i, block in enumerate(self.body):
366 | res = block(res)
367 |
368 | x = self.body_t(res) + x
369 |
370 | return x
371 |
372 | class MeanShift(nn.Conv2d):
373 | def __init__(
374 | self, rgb_range,
375 | rgb_mean=(0.4488, 0.4371, 0.4040), rgb_std=(1.0, 1.0, 1.0), sign=-1):
376 |
377 | super(MeanShift, self).__init__(3, 3, kernel_size=1)
378 | std = torch.Tensor(rgb_std)
379 | self.weight.data = torch.eye(3).view(3, 3, 1, 1) / std.view(3, 1, 1, 1)
380 | self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean) / std
381 | for p in self.parameters():
382 | p.requires_grad = False
383 |
384 |
385 | class MAN(nn.Module):
386 | def __init__(self, n_resblocks=36, n_resgroups=1, n_colors=3, n_feats=180, scale=2, res_scale = 1.0):
387 | super(MAN, self).__init__()
388 |
389 | #res_scale = res_scale
390 | self.n_resgroups = n_resgroups
391 |
392 | self.sub_mean = MeanShift(1.0)
393 | self.head = nn.Conv2d(n_colors, n_feats, 3, 1, 1)
394 |
395 | # define body module
396 | self.body = nn.ModuleList([
397 | ResGroup(
398 | n_resblocks, n_feats, res_scale=res_scale)
399 | for i in range(n_resgroups)])
400 |
401 | if self.n_resgroups > 1:
402 | self.body_t = nn.Conv2d(n_feats, n_feats, 3, 1, 1)
403 |
404 | # define tail module
405 | self.tail = nn.Sequential(
406 | nn.Conv2d(n_feats, n_colors*(scale**2), 3, 1, 1),
407 | nn.PixelShuffle(scale)
408 | )
409 | self.add_mean = MeanShift(1.0, sign=1)
410 |
411 | def forward(self, x):
412 | x = self.sub_mean(x)
413 | x = self.head(x)
414 | res = x
415 | for i in self.body:
416 | res = i(res)
417 | if self.n_resgroups>1:
418 | res = self.body_t(res) + x
419 | x = self.tail(res)
420 | x = self.add_mean(x)
421 | return x
422 |
423 | def visual_feature(self, x):
424 | fea = []
425 | x = self.head(x)
426 | res = x
427 |
428 | for i in self.body:
429 | temp = res
430 | res = i(res)
431 | fea.append(res)
432 |
433 | res = self.body_t(res) + x
434 |
435 | x = self.tail(res)
436 | return x, fea
437 |
438 | def load_state_dict(self, state_dict, strict=False):
439 | own_state = self.state_dict()
440 | for name, param in state_dict.items():
441 | if name in own_state:
442 | if isinstance(param, nn.Parameter):
443 | param = param.data
444 | try:
445 | own_state[name].copy_(param)
446 | except Exception:
447 | if name.find('tail') >= 0:
448 | print('Replace pre-trained upsampler to new one...')
449 | else:
450 | raise RuntimeError('While copying the parameter named {}, '
451 | 'whose dimensions in the model are {} and '
452 | 'whose dimensions in the checkpoint are {}.'
453 | .format(name, own_state[name].size(), param.size()))
454 | elif strict:
455 | if name.find('tail') == -1:
456 | raise KeyError('unexpected key "{}" in state_dict'
457 | .format(name))
458 |
459 | if strict:
460 | missing = set(own_state.keys()) - set(state_dict.keys())
461 | if len(missing) > 0:
462 | raise KeyError('missing keys in state_dict: "{}"'.format(missing))
--------------------------------------------------------------------------------
/archs/MAN_arch.py:
--------------------------------------------------------------------------------
1 | # -*- coding: utf-8 -*-
2 | import math
3 |
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 |
8 | from basicsr.utils.registry import ARCH_REGISTRY
9 |
10 | #LKA from VAN (https://github.com/Visual-Attention-Network)
11 | class LKA(nn.Module):
12 | def __init__(self, dim):
13 | super().__init__()
14 | self.conv0 = nn.Conv2d(dim, dim, 7, padding=7//2, groups=dim)
15 | self.conv_spatial = nn.Conv2d(dim, dim, 9, stride=1, padding=((9//2)*4), groups=dim, dilation=4)
16 | self.conv1 = nn.Conv2d(dim, dim, 1)
17 |
18 | def forward(self, x):
19 | u = x.clone()
20 | attn = self.conv0(x)
21 | attn = self.conv_spatial(attn)
22 | attn = self.conv1(attn)
23 |
24 | return u * attn
25 |
26 | class Attention(nn.Module):
27 | def __init__(self, n_feats):
28 | super().__init__()
29 |
30 | self.norm = LayerNorm(n_feats, data_format='channels_first')
31 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
32 |
33 | self.proj_1 = nn.Conv2d(n_feats, n_feats, 1)
34 | self.spatial_gating_unit = LKA(n_feats)
35 | self.proj_2 = nn.Conv2d(n_feats, n_feats, 1)
36 |
37 | def forward(self, x):
38 | shorcut = x.clone()
39 | x = self.proj_1(self.norm(x))
40 | x = self.spatial_gating_unit(x)
41 | x = self.proj_2(x)
42 | x = x*self.scale + shorcut
43 | return x
44 | #----------------------------------------------------------------------------------------------------------------
45 |
46 | class MLP(nn.Module):
47 | def __init__(self, n_feats):
48 | super().__init__()
49 |
50 | self.norm = LayerNorm(n_feats, data_format='channels_first')
51 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
52 |
53 | i_feats = 2*n_feats
54 |
55 | self.fc1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)
56 | self.act = nn.GELU()
57 | self.fc2 = nn.Conv2d(i_feats, n_feats, 1, 1, 0)
58 |
59 | def forward(self, x):
60 | shortcut = x.clone()
61 | x = self.norm(x)
62 | x = self.fc1(x)
63 | x = self.act(x)
64 | x = self.fc2(x)
65 |
66 | return x*self.scale + shortcut
67 |
68 |
69 | class CFF(nn.Module):
70 | def __init__(self, n_feats, drop=0.0, k=2, squeeze_factor= 15, attn ='GLKA'):
71 | super().__init__()
72 | i_feats =n_feats*2
73 |
74 | self.Conv1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)
75 | self.DWConv1 = nn.Sequential(
76 | nn.Conv2d(i_feats, i_feats, 7, 1, 7//2, groups= n_feats),
77 | nn.GELU())
78 | self.Conv2 = nn.Conv2d(i_feats, n_feats, 1, 1, 0)
79 |
80 | self.norm = LayerNorm(n_feats, data_format='channels_first')
81 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
82 |
83 | def forward(self, x):
84 | shortcut = x.clone()
85 |
86 | #Ghost Expand
87 | x = self.Conv1(self.norm(x))
88 | x = self.DWConv1(x)
89 | x = self.Conv2(x)
90 |
91 | return x*self.scale + shortcut
92 |
93 |
94 | class SimpleGate(nn.Module):
95 | def __init__(self, n_feats):
96 | super().__init__()
97 | i_feats =n_feats*2
98 |
99 | self.Conv1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)
100 | #self.DWConv1 = nn.Conv2d(n_feats, n_feats, 7, 1, 7//2, groups= n_feats)
101 | self.Conv2 = nn.Conv2d(n_feats, n_feats, 1, 1, 0)
102 |
103 | self.norm = LayerNorm(n_feats, data_format='channels_first')
104 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
105 |
106 | def forward(self, x):
107 | shortcut = x.clone()
108 |
109 | #Ghost Expand
110 | x = self.Conv1(self.norm(x))
111 | a, x = torch.chunk(x, 2, dim=1)
112 | x = x*a #self.DWConv1(a)
113 | x = self.Conv2(x)
114 |
115 | return x*self.scale + shortcut
116 | #-----------------------------------------------------------------------------------------------------------------
117 | #RCAN-style
118 | class RCBv6(nn.Module):
119 | def __init__(
120 | self, n_feats, k, lk=7, res_scale=1.0, style ='X', act = nn.SiLU(), deploy=False):
121 | super().__init__()
122 | self.LKA = nn.Sequential(
123 | nn.Conv2d(n_feats, n_feats, 5, 1, lk//2, groups= n_feats),
124 | nn.Conv2d(n_feats, n_feats, 7, stride=1, padding=9, groups=n_feats, dilation=3),
125 | nn.Conv2d(n_feats, n_feats, 1, 1, 0),
126 | nn.Sigmoid())
127 |
128 | #self.LFE2 = LFEv3(n_feats, attn ='CA')
129 |
130 | self.LFE = nn.Sequential(
131 | nn.Conv2d(n_feats,n_feats,3,1,1),
132 | nn.GELU(),
133 | nn.Conv2d(n_feats,n_feats,3,1,1))
134 |
135 |
136 | def forward(self, x, pre_attn=None, RAA=None):
137 | shortcut = x.clone()
138 | x = self.LFE(x)
139 |
140 | x = self.LKA(x)*x
141 |
142 | return x + shortcut
143 |
144 | #-----------------------------------------------------------------------------------------------------------------
145 |
146 |
147 | class MLKA_Ablation(nn.Module):
148 | def __init__(self, n_feats, k=2, squeeze_factor=15):
149 | super().__init__()
150 | i_feats = 2*n_feats
151 |
152 | self.n_feats= n_feats
153 | self.i_feats = i_feats
154 |
155 | self.norm = LayerNorm(n_feats, data_format='channels_first')
156 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
157 |
158 | k = 2
159 |
160 | #Multiscale Large Kernel Attention
161 | self.LKA7 = nn.Sequential(
162 | nn.Conv2d(n_feats//k, n_feats//k, 7, 1, 7//2, groups= n_feats//k),
163 | nn.Conv2d(n_feats//k, n_feats//k, 9, stride=1, padding=(9//2)*4, groups=n_feats//k, dilation=4),
164 | nn.Conv2d(n_feats//k, n_feats//k, 1, 1, 0))
165 | self.LKA5 = nn.Sequential(
166 | nn.Conv2d(n_feats//k, n_feats//k, 5, 1, 5//2, groups= n_feats//k),
167 | nn.Conv2d(n_feats//k, n_feats//k, 7, stride=1, padding=(7//2)*3, groups=n_feats//k, dilation=3),
168 | nn.Conv2d(n_feats//k, n_feats//k, 1, 1, 0))
169 | '''self.LKA3 = nn.Sequential(
170 | nn.Conv2d(n_feats//k, n_feats//k, 3, 1, 1, groups= n_feats//k),
171 | nn.Conv2d(n_feats//k, n_feats//k, 5, stride=1, padding=(5//2)*2, groups=n_feats//k, dilation=2),
172 | nn.Conv2d(n_feats//k, n_feats//k, 1, 1, 0))'''
173 |
174 | #self.X3 = nn.Conv2d(n_feats//k, n_feats//k, 3, 1, 1, groups= n_feats//k)
175 | self.X5 = nn.Conv2d(n_feats//k, n_feats//k, 5, 1, 5//2, groups= n_feats//k)
176 | self.X7 = nn.Conv2d(n_feats//k, n_feats//k, 7, 1, 7//2, groups= n_feats//k)
177 |
178 | self.proj_first = nn.Sequential(
179 | nn.Conv2d(n_feats, i_feats, 1, 1, 0))
180 |
181 | self.proj_last = nn.Sequential(
182 | nn.Conv2d(n_feats, n_feats, 1, 1, 0))
183 |
184 |
185 | def forward(self, x, pre_attn=None, RAA=None):
186 | shortcut = x.clone()
187 |
188 | x = self.norm(x)
189 |
190 | x = self.proj_first(x)
191 |
192 | a, x = torch.chunk(x, 2, dim=1)
193 |
194 | #u_1, u_2, u_3= torch.chunk(u, 3, dim=1)
195 | a_1, a_2 = torch.chunk(a, 2, dim=1)
196 |
197 | a = torch.cat([self.LKA7(a_1)*self.X7(a_1), self.LKA5(a_2)*self.X5(a_2)], dim=1)
198 |
199 | x = self.proj_last(x*a)*self.scale + shortcut
200 |
201 | return x
202 | #-----------------------------------------------------------------------------------------------------------------
203 |
204 |
205 | class LayerNorm(nn.Module):
206 | r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
207 | The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
208 | shape (batch_size, height, width, channels) while channels_first corresponds to inputs
209 | with shape (batch_size, channels, height, width).
210 | """
211 | def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
212 | super().__init__()
213 | self.weight = nn.Parameter(torch.ones(normalized_shape))
214 | self.bias = nn.Parameter(torch.zeros(normalized_shape))
215 | self.eps = eps
216 | self.data_format = data_format
217 | if self.data_format not in ["channels_last", "channels_first"]:
218 | raise NotImplementedError
219 | self.normalized_shape = (normalized_shape, )
220 |
221 | def forward(self, x):
222 | if self.data_format == "channels_last":
223 | return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
224 | elif self.data_format == "channels_first":
225 | u = x.mean(1, keepdim=True)
226 | s = (x - u).pow(2).mean(1, keepdim=True)
227 | x = (x - u) / torch.sqrt(s + self.eps)
228 | x = self.weight[:, None, None] * x + self.bias[:, None, None]
229 | return x
230 |
231 |
232 |
233 | class SGAB(nn.Module):
234 | def __init__(self, n_feats, drop=0.0, k=2, squeeze_factor= 15, attn ='GLKA'):
235 | super().__init__()
236 | i_feats =n_feats*2
237 |
238 | self.Conv1 = nn.Conv2d(n_feats, i_feats, 1, 1, 0)
239 | self.DWConv1 = nn.Conv2d(n_feats, n_feats, 7, 1, 7//2, groups= n_feats)
240 | self.Conv2 = nn.Conv2d(n_feats, n_feats, 1, 1, 0)
241 |
242 | self.norm = LayerNorm(n_feats, data_format='channels_first')
243 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
244 |
245 | def forward(self, x):
246 | shortcut = x.clone()
247 |
248 | #Ghost Expand
249 | x = self.Conv1(self.norm(x))
250 | a, x = torch.chunk(x, 2, dim=1)
251 | x = x*self.DWConv1(a)
252 | x = self.Conv2(x)
253 |
254 | return x*self.scale + shortcut
255 |
256 | class GroupGLKA(nn.Module):
257 | def __init__(self, n_feats, k=2, squeeze_factor=15):
258 | super().__init__()
259 | i_feats = 2*n_feats
260 |
261 | self.n_feats= n_feats
262 | self.i_feats = i_feats
263 |
264 | self.norm = LayerNorm(n_feats, data_format='channels_first')
265 | self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
266 |
267 | #Multiscale Large Kernel Attention
268 | self.LKA7 = nn.Sequential(
269 | nn.Conv2d(n_feats//3, n_feats//3, 7, 1, 7//2, groups= n_feats//3),
270 | nn.Conv2d(n_feats//3, n_feats//3, 9, stride=1, padding=(9//2)*4, groups=n_feats//3, dilation=4),
271 | nn.Conv2d(n_feats//3, n_feats//3, 1, 1, 0))
272 | self.LKA5 = nn.Sequential(
273 | nn.Conv2d(n_feats//3, n_feats//3, 5, 1, 5//2, groups= n_feats//3),
274 | nn.Conv2d(n_feats//3, n_feats//3, 7, stride=1, padding=(7//2)*3, groups=n_feats//3, dilation=3),
275 | nn.Conv2d(n_feats//3, n_feats//3, 1, 1, 0))
276 | self.LKA3 = nn.Sequential(
277 | nn.Conv2d(n_feats//3, n_feats//3, 3, 1, 1, groups= n_feats//3),
278 | nn.Conv2d(n_feats//3, n_feats//3, 5, stride=1, padding=(5//2)*2, groups=n_feats//3, dilation=2),
279 | nn.Conv2d(n_feats//3, n_feats//3, 1, 1, 0))
280 |
281 | self.X3 = nn.Conv2d(n_feats//3, n_feats//3, 3, 1, 1, groups= n_feats//3)
282 | self.X5 = nn.Conv2d(n_feats//3, n_feats//3, 5, 1, 5//2, groups= n_feats//3)
283 | self.X7 = nn.Conv2d(n_feats//3, n_feats//3, 7, 1, 7//2, groups= n_feats//3)
284 |
285 | self.proj_first = nn.Sequential(
286 | nn.Conv2d(n_feats, i_feats, 1, 1, 0))
287 |
288 | self.proj_last = nn.Sequential(
289 | nn.Conv2d(n_feats, n_feats, 1, 1, 0))
290 |
291 |
292 | def forward(self, x, pre_attn=None, RAA=None):
293 | shortcut = x.clone()
294 |
295 | x = self.norm(x)
296 |
297 | x = self.proj_first(x)
298 |
299 | a, x = torch.chunk(x, 2, dim=1)
300 |
301 | a_1, a_2, a_3= torch.chunk(a, 3, dim=1)
302 |
303 | a = torch.cat([self.LKA3(a_1)*self.X3(a_1), self.LKA5(a_2)*self.X5(a_2), self.LKA7(a_3)*self.X7(a_3)], dim=1)
304 |
305 | x = self.proj_last(x*a)*self.scale + shortcut
306 |
307 | return x
308 |
309 |
310 | # MAB
311 | class MAB(nn.Module):
312 | def __init__(
313 | self, n_feats):
314 | super().__init__()
315 |
316 | self.LKA = GroupGLKA(n_feats)
317 |
318 | self.LFE = SGAB(n_feats)
319 |
320 | def forward(self, x, pre_attn=None, RAA=None):
321 | #large kernel attention
322 | x = self.LKA(x)
323 |
324 | #local feature extraction
325 | x = self.LFE(x)
326 |
327 | return x
328 |
329 | class LKAT(nn.Module):
330 | def __init__(self, n_feats):
331 | super().__init__()
332 |
333 | #self.norm = LayerNorm(n_feats, data_format='channels_first')
334 | #self.scale = nn.Parameter(torch.zeros((1, n_feats, 1, 1)), requires_grad=True)
335 |
336 | self.conv0 = nn.Sequential(
337 | nn.Conv2d(n_feats, n_feats, 1, 1, 0),
338 | nn.GELU())
339 |
340 | self.att = nn.Sequential(
341 | nn.Conv2d(n_feats, n_feats, 7, 1, 7//2, groups= n_feats),
342 | nn.Conv2d(n_feats, n_feats, 9, stride=1, padding=(9//2)*3, groups=n_feats, dilation=3),
343 | nn.Conv2d(n_feats, n_feats, 1, 1, 0))
344 |
345 | self.conv1 = nn.Conv2d(n_feats, n_feats, 1, 1, 0)
346 |
347 | def forward(self, x):
348 | x = self.conv0(x)
349 | x = x*self.att(x)
350 | x = self.conv1(x)
351 | return x
352 |
353 | class ResGroup(nn.Module):
354 | def __init__(self, n_resblocks, n_feats, res_scale=1.0):
355 | super(ResGroup, self).__init__()
356 | self.body = nn.ModuleList([
357 | MAB(n_feats) \
358 | for _ in range(n_resblocks)])
359 |
360 | self.body_t = LKAT(n_feats)
361 |
362 | def forward(self, x):
363 | res = x.clone()
364 |
365 | for i, block in enumerate(self.body):
366 | res = block(res)
367 |
368 | x = self.body_t(res) + x
369 |
370 | return x
371 |
372 | class MeanShift(nn.Conv2d):
373 | def __init__(
374 | self, rgb_range,
375 | rgb_mean=(0.4488, 0.4371, 0.4040), rgb_std=(1.0, 1.0, 1.0), sign=-1):
376 |
377 | super(MeanShift, self).__init__(3, 3, kernel_size=1)
378 | std = torch.Tensor(rgb_std)
379 | self.weight.data = torch.eye(3).view(3, 3, 1, 1) / std.view(3, 1, 1, 1)
380 | self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean) / std
381 | for p in self.parameters():
382 | p.requires_grad = False
383 |
384 |
385 | @ARCH_REGISTRY.register()
386 | class MAN(nn.Module):
387 | def __init__(self, n_resblocks=36, n_resgroups=1, n_colors=3, n_feats=180, scale=2, res_scale = 1.0):
388 | super(MAN, self).__init__()
389 |
390 | #res_scale = res_scale
391 | self.n_resgroups = n_resgroups
392 |
393 | self.sub_mean = MeanShift(1.0)
394 | self.head = nn.Conv2d(n_colors, n_feats, 3, 1, 1)
395 |
396 | # define body module
397 | self.body = nn.ModuleList([
398 | ResGroup(
399 | n_resblocks, n_feats, res_scale=res_scale)
400 | for i in range(n_resgroups)])
401 |
402 | if self.n_resgroups > 1:
403 | self.body_t = nn.Conv2d(n_feats, n_feats, 3, 1, 1)
404 |
405 | # define tail module
406 | self.tail = nn.Sequential(
407 | nn.Conv2d(n_feats, n_colors*(scale**2), 3, 1, 1),
408 | nn.PixelShuffle(scale)
409 | )
410 | self.add_mean = MeanShift(1.0, sign=1)
411 |
412 | def forward(self, x):
413 | x = self.sub_mean(x)
414 | x = self.head(x)
415 | res = x
416 | for i in self.body:
417 | res = i(res)
418 | if self.n_resgroups>1:
419 | res = self.body_t(res) + x
420 | x = self.tail(res)
421 | x = self.add_mean(x)
422 | return x
423 |
424 | def visual_feature(self, x):
425 | fea = []
426 | x = self.head(x)
427 | res = x
428 |
429 | for i in self.body:
430 | temp = res
431 | res = i(res)
432 | fea.append(res)
433 |
434 | res = self.body_t(res) + x
435 |
436 | x = self.tail(res)
437 | return x, fea
438 |
439 | def load_state_dict(self, state_dict, strict=False):
440 | own_state = self.state_dict()
441 | for name, param in state_dict.items():
442 | if name in own_state:
443 | if isinstance(param, nn.Parameter):
444 | param = param.data
445 | try:
446 | own_state[name].copy_(param)
447 | except Exception:
448 | if name.find('tail') >= 0:
449 | print('Replace pre-trained upsampler to new one...')
450 | else:
451 | raise RuntimeError('While copying the parameter named {}, '
452 | 'whose dimensions in the model are {} and '
453 | 'whose dimensions in the checkpoint are {}.'
454 | .format(name, own_state[name].size(), param.size()))
455 | elif strict:
456 | if name.find('tail') == -1:
457 | raise KeyError('unexpected key "{}" in state_dict'
458 | .format(name))
459 |
460 | if strict:
461 | missing = set(own_state.keys()) - set(state_dict.keys())
462 | if len(missing) > 0:
463 | raise KeyError('missing keys in state_dict: "{}"'.format(missing))
464 |
--------------------------------------------------------------------------------
/LAM/ModelZoo/NN/san.py:
--------------------------------------------------------------------------------
1 | from ..NN import common
2 | import torch
3 | import torch.nn as nn
4 | import torch.nn.functional as F
5 | from ..NN.MPNCOV.python import MPNCOV
6 |
7 |
8 | #
9 | def make_model(args, parent=False):
10 | return SAN(args)
11 |
12 |
13 | ## non_local module
14 | class _NonLocalBlockND(nn.Module):
15 | def __init__(self, in_channels, inter_channels=None, dimension=3, mode='embedded_gaussian',
16 | sub_sample=True, bn_layer=True):
17 | super(_NonLocalBlockND, self).__init__()
18 | assert dimension in [1, 2, 3]
19 | assert mode in ['embedded_gaussian', 'gaussian', 'dot_product', 'concatenation']
20 |
21 | # print('Dimension: %d, mode: %s' % (dimension, mode))
22 |
23 | self.mode = mode
24 | self.dimension = dimension
25 | self.sub_sample = sub_sample
26 |
27 | self.in_channels = in_channels
28 | self.inter_channels = inter_channels
29 |
30 | if self.inter_channels is None:
31 | self.inter_channels = in_channels // 2
32 | if self.inter_channels == 0:
33 | self.inter_channels = 1
34 |
35 | if dimension == 3:
36 | conv_nd = nn.Conv3d
37 | max_pool = nn.MaxPool3d
38 | bn = nn.BatchNorm3d
39 | elif dimension == 2:
40 | conv_nd = nn.Conv2d
41 | max_pool = nn.MaxPool2d
42 | sub_sample = nn.Upsample
43 | bn = nn.BatchNorm2d
44 | else:
45 | conv_nd = nn.Conv1d
46 | max_pool = nn.MaxPool1d
47 | bn = nn.BatchNorm1d
48 |
49 | self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
50 | kernel_size=1, stride=1, padding=0)
51 |
52 | if bn_layer:
53 | self.W = nn.Sequential(
54 | conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
55 | kernel_size=1, stride=1, padding=0),
56 | bn(self.in_channels)
57 | )
58 | nn.init.constant_(self.W[1].weight, 0)
59 | nn.init.constant_(self.W[1].bias, 0)
60 | else:
61 | self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
62 | kernel_size=1, stride=1, padding=0)
63 | nn.init.constant_(self.W.weight, 0)
64 | nn.init.constant_(self.W.bias, 0)
65 |
66 | self.theta = None
67 | self.phi = None
68 | self.concat_project = None
69 | # self.fc = nn.Linear(64,2304,bias=True)
70 | # self.sub_bilinear = nn.Upsample(size=(48,48),mode='bilinear')
71 | # self.sub_maxpool = nn.AdaptiveMaxPool2d(output_size=(48,48))
72 | if mode in ['embedded_gaussian', 'dot_product', 'concatenation']:
73 | self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
74 | kernel_size=1, stride=1, padding=0)
75 | self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
76 | kernel_size=1, stride=1, padding=0)
77 |
78 | if mode == 'embedded_gaussian':
79 | self.operation_function = self._embedded_gaussian
80 | elif mode == 'dot_product':
81 | self.operation_function = self._dot_product
82 | elif mode == 'concatenation':
83 | self.operation_function = self._concatenation
84 | self.concat_project = nn.Sequential(
85 | nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),
86 | nn.ReLU()
87 | )
88 | elif mode == 'gaussian':
89 | self.operation_function = self._gaussian
90 |
91 | if sub_sample:
92 | self.g = nn.Sequential(self.g, max_pool(kernel_size=2))
93 | if self.phi is None:
94 | self.phi = max_pool(kernel_size=2)
95 | else:
96 | self.phi = nn.Sequential(self.phi, max_pool(kernel_size=2))
97 |
98 | def forward(self, x):
99 | '''
100 | :param x: (b, c, t, h, w)
101 | :return:
102 | '''
103 |
104 | output = self.operation_function(x)
105 | return output
106 |
107 | def _embedded_gaussian(self, x):
108 | batch_size, C, H, W = x.shape
109 |
110 | # x_sub = self.sub_bilinear(x) # bilinear downsample
111 | # x_sub = self.sub_maxpool(x) # maxpool downsample
112 |
113 | ##
114 | # g_x = x.view(batch_size, self.inter_channels, -1)
115 | # g_x = g_x.permute(0, 2, 1)
116 | #
117 | # # theta=>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, thw, 0.5c)
118 | # # phi =>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, 0.5c, thw)
119 | # # f=>(b, thw, 0.5c)dot(b, 0.5c, twh) = (b, thw, thw)
120 | # theta_x = x.view(batch_size, self.inter_channels, -1)
121 | # theta_x = theta_x.permute(0, 2, 1)
122 | # fc = self.fc(theta_x)
123 | # # phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
124 | # # f = torch.matmul(theta_x, phi_x)
125 | # # return f
126 | # # f_div_C = F.softmax(fc, dim=-1)
127 | # return fc
128 |
129 | ##
130 | # g=>(b, c, t, h, w)->(b, 0.5c, t, h, w)->(b, thw, 0.5c)
131 | g_x = self.g(x).view(batch_size, self.inter_channels, -1)
132 | g_x = g_x.permute(0, 2, 1)
133 |
134 | # theta=>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, thw, 0.5c)
135 | # phi =>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, 0.5c, thw)
136 | # f=>(b, thw, 0.5c)dot(b, 0.5c, twh) = (b, thw, thw)
137 | theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
138 | theta_x = theta_x.permute(0, 2, 1)
139 | phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
140 | f = torch.matmul(theta_x, phi_x)
141 | # return f
142 | f_div_C = F.softmax(f, dim=-1)
143 | # return f_div_C
144 | # (b, thw, thw)dot(b, thw, 0.5c) = (b, thw, 0.5c)->(b, 0.5c, t, h, w)->(b, c, t, h, w)
145 | y = torch.matmul(f_div_C, g_x)
146 | y = y.permute(0, 2, 1).contiguous()
147 | y = y.view(batch_size, self.inter_channels, *x.size()[2:])
148 | W_y = self.W(y)
149 | z = W_y + x
150 |
151 | return z
152 |
153 | def _gaussian(self, x):
154 | batch_size = x.size(0)
155 | g_x = self.g(x).view(batch_size, self.inter_channels, -1)
156 | g_x = g_x.permute(0, 2, 1)
157 |
158 | theta_x = x.view(batch_size, self.in_channels, -1)
159 | theta_x = theta_x.permute(0, 2, 1)
160 |
161 | if self.sub_sample:
162 | phi_x = self.phi(x).view(batch_size, self.in_channels, -1)
163 | else:
164 | phi_x = x.view(batch_size, self.in_channels, -1)
165 |
166 | f = torch.matmul(theta_x, phi_x)
167 | f_div_C = F.softmax(f, dim=-1)
168 |
169 | y = torch.matmul(f_div_C, g_x)
170 | y = y.permute(0, 2, 1).contiguous()
171 | y = y.view(batch_size, self.inter_channels, *x.size()[2:])
172 | W_y = self.W(y)
173 | z = W_y + x
174 |
175 | return z
176 |
177 | def _dot_product(self, x):
178 | batch_size = x.size(0)
179 |
180 | g_x = self.g(x).view(batch_size, self.inter_channels, -1)
181 | g_x = g_x.permute(0, 2, 1)
182 |
183 | theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
184 | theta_x = theta_x.permute(0, 2, 1)
185 | phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
186 | f = torch.matmul(theta_x, phi_x)
187 | N = f.size(-1)
188 | f_div_C = f / N
189 |
190 | y = torch.matmul(f_div_C, g_x)
191 | y = y.permute(0, 2, 1).contiguous()
192 | y = y.view(batch_size, self.inter_channels, *x.size()[2:])
193 | W_y = self.W(y)
194 | z = W_y + x
195 |
196 | return z
197 |
198 | def _concatenation(self, x):
199 | batch_size = x.size(0)
200 |
201 | g_x = self.g(x).view(batch_size, self.inter_channels, -1)
202 | g_x = g_x.permute(0, 2, 1)
203 |
204 | # (b, c, N, 1)
205 | theta_x = self.theta(x).view(batch_size, self.inter_channels, -1, 1)
206 | # (b, c, 1, N)
207 | phi_x = self.phi(x).view(batch_size, self.inter_channels, 1, -1)
208 |
209 | h = theta_x.size(2)
210 | w = phi_x.size(3)
211 | theta_x = theta_x.repeat(1, 1, 1, w)
212 | phi_x = phi_x.repeat(1, 1, h, 1)
213 |
214 | concat_feature = torch.cat([theta_x, phi_x], dim=1)
215 | f = self.concat_project(concat_feature)
216 | b, _, h, w = f.size()
217 | f = f.view(b, h, w)
218 |
219 | N = f.size(-1)
220 | f_div_C = f / N
221 |
222 | y = torch.matmul(f_div_C, g_x)
223 | y = y.permute(0, 2, 1).contiguous()
224 | y = y.view(batch_size, self.inter_channels, *x.size()[2:])
225 | W_y = self.W(y)
226 | z = W_y + x
227 |
228 | return z
229 |
230 |
231 | class NONLocalBlock1D(_NonLocalBlockND):
232 | def __init__(self, in_channels, inter_channels=None, mode='embedded_gaussian', sub_sample=True, bn_layer=True):
233 | super(NONLocalBlock1D, self).__init__(in_channels,
234 | inter_channels=inter_channels,
235 | dimension=1, mode=mode,
236 | sub_sample=sub_sample,
237 | bn_layer=bn_layer)
238 |
239 |
240 | class NONLocalBlock2D(_NonLocalBlockND):
241 | def __init__(self, in_channels, inter_channels=None, mode='embedded_gaussian', sub_sample=True, bn_layer=True):
242 | super(NONLocalBlock2D, self).__init__(in_channels,
243 | inter_channels=inter_channels,
244 | dimension=2, mode=mode,
245 | sub_sample=sub_sample,
246 | bn_layer=bn_layer)
247 |
248 |
249 | ## Channel Attention (CA) Layer
250 | class CALayer(nn.Module):
251 | def __init__(self, channel, reduction=8):
252 | super(CALayer, self).__init__()
253 | # global average pooling: feature --> point
254 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
255 | self.max_pool = nn.AdaptiveMaxPool2d(1)
256 | # feature channel downscale and upscale --> channel weight
257 | self.conv_du = nn.Sequential(
258 | nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
259 | nn.ReLU(inplace=True),
260 | nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
261 | # nn.Sigmoid()
262 | # nn.BatchNorm2d(channel)
263 | )
264 |
265 | def forward(self, x):
266 | _, _, h, w = x.shape
267 | y_ave = self.avg_pool(x)
268 | # y_max = self.max_pool(x)
269 | y_ave = self.conv_du(y_ave)
270 | # y_max = self.conv_du(y_max)
271 | # y = y_ave + y_max
272 | # expand y to C*H*W
273 | # expand_y = y.expand(-1,-1,h,w)
274 | return y_ave
275 |
276 |
277 | ## second-order Channel attention (SOCA)
278 | class SOCA(nn.Module):
279 | def __init__(self, channel, reduction=8):
280 | super(SOCA, self).__init__()
281 | # global average pooling: feature --> point
282 | # self.avg_pool = nn.AdaptiveAvgPool2d(1)
283 | # self.max_pool = nn.AdaptiveMaxPool2d(1)
284 | self.max_pool = nn.MaxPool2d(kernel_size=2)
285 |
286 | # feature channel downscale and upscale --> channel weight
287 | self.conv_du = nn.Sequential(
288 | nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
289 | nn.ReLU(inplace=True),
290 | nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
291 | nn.Sigmoid()
292 | # nn.BatchNorm2d(channel)
293 | )
294 |
295 | def forward(self, x):
296 | batch_size, C, h, w = x.shape # x: NxCxHxW
297 | N = int(h * w)
298 | min_h = min(h, w)
299 | h1 = 1000
300 | w1 = 1000
301 | if h < h1 and w < w1:
302 | x_sub = x
303 | elif h < h1 and w > w1:
304 | # H = (h - h1) // 2
305 | W = (w - w1) // 2
306 | x_sub = x[:, :, :, W:(W + w1)]
307 | elif w < w1 and h > h1:
308 | H = (h - h1) // 2
309 | # W = (w - w1) // 2
310 | x_sub = x[:, :, H:H + h1, :]
311 | else:
312 | H = (h - h1) // 2
313 | W = (w - w1) // 2
314 | x_sub = x[:, :, H:(H + h1), W:(W + w1)]
315 | # subsample
316 | # subsample_scale = 2
317 | # subsample = nn.Upsample(size=(h // subsample_scale, w // subsample_scale), mode='nearest')
318 | # x_sub = subsample(x)
319 | # max_pool = nn.MaxPool2d(kernel_size=2)
320 | # max_pool = nn.AvgPool2d(kernel_size=2)
321 | # x_sub = self.max_pool(x)
322 | ##
323 | ## MPN-COV
324 | cov_mat = MPNCOV.CovpoolLayer(x_sub) # Global Covariance pooling layer
325 | cov_mat_sqrt = MPNCOV.SqrtmLayer(cov_mat,
326 | 5) # Matrix square root layer( including pre-norm,Newton-Schulz iter. and post-com. with 5 iteration)
327 | ##
328 | cov_mat_sum = torch.mean(cov_mat_sqrt, 1)
329 | cov_mat_sum = cov_mat_sum.view(batch_size, C, 1, 1)
330 | # y_ave = self.avg_pool(x)
331 | # y_max = self.max_pool(x)
332 | y_cov = self.conv_du(cov_mat_sum)
333 | # y_max = self.conv_du(y_max)
334 | # y = y_ave + y_max
335 | # expand y to C*H*W
336 | # expand_y = y.expand(-1,-1,h,w)
337 | return y_cov * x
338 |
339 |
340 | ## self-attention+ channel attention module
341 | class Nonlocal_CA(nn.Module):
342 | def __init__(self, in_feat=64, inter_feat=32, reduction=8, sub_sample=False, bn_layer=True):
343 | super(Nonlocal_CA, self).__init__()
344 | # second-order channel attention
345 | self.soca = SOCA(in_feat, reduction=reduction)
346 | # nonlocal module
347 | self.non_local = _NonLocalBlockND(in_channels=in_feat, inter_channels=inter_feat, dimension=2, sub_sample=sub_sample, bn_layer=bn_layer)
348 |
349 | self.sigmoid = nn.Sigmoid()
350 |
351 | def forward(self, x):
352 | ## divide feature map into 4 part
353 | batch_size, C, H, W = x.shape
354 | H1 = int(H / 2)
355 | W1 = int(W / 2)
356 | nonlocal_feat = torch.zeros_like(x)
357 |
358 | feat_sub_lu = x[:, :, :H1, :W1]
359 | feat_sub_ld = x[:, :, H1:, :W1]
360 | feat_sub_ru = x[:, :, :H1, W1:]
361 | feat_sub_rd = x[:, :, H1:, W1:]
362 |
363 | nonlocal_lu = self.non_local(feat_sub_lu)
364 | nonlocal_ld = self.non_local(feat_sub_ld)
365 | nonlocal_ru = self.non_local(feat_sub_ru)
366 | nonlocal_rd = self.non_local(feat_sub_rd)
367 | nonlocal_feat[:, :, :H1, :W1] = nonlocal_lu
368 | nonlocal_feat[:, :, H1:, :W1] = nonlocal_ld
369 | nonlocal_feat[:, :, :H1, W1:] = nonlocal_ru
370 | nonlocal_feat[:, :, H1:, W1:] = nonlocal_rd
371 |
372 | return nonlocal_feat
373 |
374 |
375 | ## Residual Block (RB)
376 | class RB(nn.Module):
377 | def __init__(self, conv, n_feat, kernel_size, reduction, bias=True, bn=False, act=nn.ReLU(inplace=True),
378 | res_scale=1, dilation=2):
379 | super(RB, self).__init__()
380 | modules_body = []
381 |
382 | # self.gamma1 = nn.Parameter(torch.ones(1))
383 | self.gamma1 = 1.0
384 | # self.salayer = SALayer(n_feat, reduction=reduction, dilation=dilation)
385 | # self.salayer = SALayer2(n_feat, reduction=reduction, dilation=dilation)
386 |
387 | self.conv_first = nn.Sequential(conv(n_feat, n_feat, kernel_size, bias=bias),
388 | act,
389 | conv(n_feat, n_feat, kernel_size, bias=bias)
390 | )
391 |
392 | self.res_scale = res_scale
393 |
394 | def forward(self, x):
395 | y = self.conv_first(x)
396 | y = y + x
397 |
398 | return y
399 |
400 |
401 | ## Local-source Residual Attention Group (LSRARG)
402 | class LSRAG(nn.Module):
403 | def __init__(self, conv, n_feat, kernel_size, reduction, act, res_scale, n_resblocks):
404 | super(LSRAG, self).__init__()
405 | ##
406 | self.rcab = nn.ModuleList([RB(conv, n_feat, kernel_size, reduction, \
407 | bias=True, bn=False, act=nn.ReLU(inplace=True), res_scale=1) for _ in
408 | range(n_resblocks)])
409 | self.soca = (SOCA(n_feat, reduction=reduction))
410 | self.conv_last = (conv(n_feat, n_feat, kernel_size))
411 | self.n_resblocks = n_resblocks
412 | ##
413 | # modules_body = []
414 | self.gamma = nn.Parameter(torch.zeros(1))
415 | # self.gamma = 0.2
416 | # for i in range(n_resblocks):
417 | # modules_body.append(RCAB(conv, n_feat, kernel_size, reduction, bias=True, bn=False, act=nn.ReLU(inplace=True), res_scale=1))
418 | # modules_body.append(SOCA(n_feat,reduction=reduction))
419 | # # modules_body.append(Nonlocal_CA(in_feat=n_feat, inter_feat=n_feat//8, reduction =reduction, sub_sample=False, bn_layer=False))
420 | # modules_body.append(conv(n_feat, n_feat, kernel_size))
421 | # self.body = nn.Sequential(*modules_body)
422 | ##
423 |
424 | def make_layer(self, block, num_of_layer):
425 | layers = []
426 | for _ in range(num_of_layer):
427 | layers.append(block)
428 | return nn.ModuleList(layers)
429 | # return nn.Sequential(*layers)
430 |
431 | def forward(self, x):
432 | residual = x
433 | # batch_size,C,H,W = x.shape
434 | # y_pre = self.body(x)
435 | # y_pre = y_pre + x
436 | # return y_pre
437 |
438 | ## share-source skip connection
439 |
440 | for i, l in enumerate(self.rcab):
441 | # x = l(x) + self.gamma*residual
442 | x = l(x)
443 | x = self.soca(x)
444 | x = self.conv_last(x)
445 |
446 | x = x + residual
447 |
448 | return x
449 | ##
450 |
451 |
452 | ## Second-order Channel Attention Network (SAN)
453 | class SAN(nn.Module):
454 | def __init__(self, factor=4, num_channels=3, conv=common.default_conv):
455 | super(SAN, self).__init__()
456 | n_resgroups = 20
457 | n_resblocks = 10
458 | n_feats = 64
459 | kernel_size = 3
460 | reduction = 16
461 | scale = factor
462 | act = nn.ReLU(inplace=True)
463 |
464 | # RGB mean for DIV2K
465 | rgb_mean = (0.4488, 0.4371, 0.4040)
466 | rgb_std = (1.0, 1.0, 1.0)
467 | self.sub_mean = common.MeanShift(1.0, rgb_mean, rgb_std)
468 | # self.soca= SOCA(n_feats, reduction=reduction)
469 |
470 | # define head module
471 | modules_head = [conv(num_channels, n_feats, kernel_size)]
472 |
473 | # define body module
474 | ## share-source skip connection
475 |
476 | ##
477 | self.gamma = nn.Parameter(torch.zeros(1))
478 | # self.gamma = 0.2
479 | self.n_resgroups = n_resgroups
480 | self.RG = nn.ModuleList([LSRAG(conv, n_feats, kernel_size, reduction, \
481 | act=act, res_scale=1.0, n_resblocks=n_resblocks) for _ in
482 | range(n_resgroups)])
483 | self.conv_last = conv(n_feats, n_feats, kernel_size)
484 |
485 | # modules_body = [
486 | # ResidualGroup(
487 | # conv, n_feats, kernel_size, reduction, act=act, res_scale=args.res_scale, n_resblocks=n_resblocks) \
488 | # for _ in range(n_resgroups)]
489 | # modules_body.append(conv(n_feats, n_feats, kernel_size))
490 |
491 | # define tail module
492 | modules_tail = [
493 | common.Upsampler(conv, scale, n_feats, act=False),
494 | conv(n_feats, num_channels, kernel_size)]
495 |
496 | self.add_mean = common.MeanShift(1.0, rgb_mean, rgb_std, 1)
497 | self.non_local = Nonlocal_CA(in_feat=n_feats, inter_feat=n_feats // 8, reduction=8, sub_sample=False,
498 | bn_layer=False)
499 |
500 | self.head = nn.Sequential(*modules_head)
501 | # self.body = nn.Sequential(*modules_body)
502 | self.tail = nn.Sequential(*modules_tail)
503 |
504 | def make_layer(self, block, num_of_layer):
505 | layers = []
506 | for _ in range(num_of_layer):
507 | layers.append(block)
508 |
509 | return nn.ModuleList(layers)
510 | # return nn.Sequential(*layers)
511 |
512 | def forward(self, x):
513 | x = self.sub_mean(x * 255.)
514 | x = self.head(x)
515 |
516 | ## add nonlocal
517 | xx = self.non_local(x)
518 |
519 | # share-source skip connection
520 | residual = xx
521 |
522 | # res = self.RG(xx)
523 | # res = res + xx
524 | ## share-source residual gruop
525 | for i, l in enumerate(self.RG):
526 | xx = l(xx) + self.gamma * residual
527 | # xx = self.gamma*xx + residual
528 | # body part
529 | # res = self.body(xx)
530 | ##
531 | ## add nonlocal
532 | res = self.non_local(xx)
533 | ##
534 | # res = self.soca(res)
535 | # res += x
536 | res = res + x
537 |
538 | x = self.tail(res)
539 | x = self.add_mean(x)
540 |
541 | return x / 255.
542 |
543 | def load_state_dict(self, state_dict, strict=False):
544 | own_state = self.state_dict()
545 | for name, param in state_dict.items():
546 | if name in own_state:
547 | if isinstance(param, nn.Parameter):
548 | param = param.data
549 | try:
550 | own_state[name].copy_(param)
551 | except Exception:
552 | if name.find('tail') >= 0:
553 | print('Replace pre-trained upsampler to new one...')
554 | else:
555 | raise RuntimeError('While copying the parameter named {}, '
556 | 'whose dimensions in the model are {} and '
557 | 'whose dimensions in the checkpoint are {}.'
558 | .format(name, own_state[name].size(), param.size()))
559 | elif strict:
560 | if name.find('tail') == -1:
561 | raise KeyError('unexpected key "{}" in state_dict'
562 | .format(name))
563 |
564 | if strict:
565 | missing = set(own_state.keys()) - set(state_dict.keys())
566 | if len(missing) > 0:
567 | raise KeyError('missing keys in state_dict: "{}"'.format(missing))
--------------------------------------------------------------------------------