爆改YOLOv8 | 利用MB-TaylorFormer提高YOLOv8图像去雾检测

1，本文介绍

MB-TaylorFormer是一种新型多支路线性Transformer网络，用于图像去雾任务。它通过泰勒展开改进了softmax-attention，使用多支路和多尺度结构以获取多条理和多尺度的信息，且比传统方法在性能、盘算量和网络重量上更优。
关于MB-TaylorFormer的详细介绍可以看论文：https://arxiv.org/pdf/2308.14036.pdf
本文将讲解怎样将MB-TaylorFormer融合进yolov8
话不多说，上代码！
２，将MB-TaylorFormer融合进YOLＯv8

2.1 步调一
首先找到如下的目次'ultralytics/nn'，然后在这个目次下创建一个'Addmodules'文件夹，然后在这个目次下创建一个TaylorFormer.py文件，文件名字可以根据你自己的习惯起，然后将MB-TaylorFormer的核心代码复制进去。

1. import torch
2. import torch.nn as nn
3. import torch.nn.functional as F
4. from torchvision.ops.deform_conv import DeformConv2d
5. import numbers
6. import math
7. from einops import rearrange
8. import numpy as np
11. __all__ = ['MB_TaylorFormer']
13. freqs_dict = dict()
16. ##########################################################################
18. def to_3d(x):
19.     return rearrange(x, 'b c h w -> b (h w) c')
22. def to_4d(x, h, w):
23.     return rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
26. class BiasFree_LayerNorm(nn.Module):
27.     def __init__(self, normalized_shape):
28.         super(BiasFree_LayerNorm, self).__init__()
29.         if isinstance(normalized_shape, numbers.Integral):
30.             normalized_shape = (normalized_shape,)
31.         normalized_shape = torch.Size(normalized_shape)
33.         assert len(normalized_shape) == 1
35.         self.weight = nn.Parameter(torch.ones(normalized_shape))
36.         self.normalized_shape = normalized_shape
38.     def forward(self, x):
39.         sigma = x.var(-1, keepdim=True, unbiased=False)
40.         return x / torch.sqrt(sigma + 1e-5) * self.weight
43. class WithBias_LayerNorm(nn.Module):
44.     def __init__(self, normalized_shape):
45.         super(WithBias_LayerNorm, self).__init__()
46.         if isinstance(normalized_shape, numbers.Integral):
47.             normalized_shape = (normalized_shape,)
48.         normalized_shape = torch.Size(normalized_shape)
50.         assert len(normalized_shape) == 1
52.         self.weight = nn.Parameter(torch.ones(normalized_shape))
53.         self.bias = nn.Parameter(torch.zeros(normalized_shape))
54.         self.normalized_shape = normalized_shape
56.     def forward(self, x):
57.         mu = x.mean(-1, keepdim=True)
58.         sigma = x.var(-1, keepdim=True, unbiased=False)
59.         return (x - mu) / torch.sqrt(sigma + 1e-5) * self.weight + self.bias
62. class LayerNorm(nn.Module):
63.     def __init__(self, dim, LayerNorm_type):
64.         super(LayerNorm, self).__init__()
65.         if LayerNorm_type == 'BiasFree':
66.             self.body = BiasFree_LayerNorm(dim)
67.         else:
68.             self.body = WithBias_LayerNorm(dim)
70.     def forward(self, x):
71.         h, w = x.shape[-2:]
72.         return to_4d(self.body(to_3d(x)), h, w)
75. ##########################################################################
76. ## Gated-Dconv Feed-Forward Network (GDFN)
77. class FeedForward(nn.Module):
78.     def __init__(self, dim, ffn_expansion_factor, bias):
79.         super(FeedForward, self).__init__()
81.         hidden_features = int(dim * ffn_expansion_factor)
83.         self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
85.         self.dwconv = nn.Conv2d(hidden_features * 2, hidden_features * 2, kernel_size=3, stride=1, padding=1,
86.                                 groups=hidden_features * 2, bias=bias)
88.         self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
90.     def forward(self, x):
91.         x = self.project_in(x)
92.         x1, x2 = self.dwconv(x).chunk(2, dim=1)
93.         x = F.gelu(x1) * x2
94.         x = self.project_out(x)
95.         return x
98. class refine_att(nn.Module):
99.     """Convolutional relative position encoding."""
101.     def __init__(self, Ch, h, window):
103.         super().__init__()
105.         if isinstance(window, int):
106.             # Set the same window size for all attention heads.
107.             window = {window: h}
108.             self.window = window
109.         elif isinstance(window, dict):
110.             self.window = window
111.         else:
113.             raise ValueError()
115.         self.conv_list = nn.ModuleList()
116.         self.head_splits = []
117.         for cur_window, cur_head_split in window.items():
118.             dilation = 1  # Use dilation=1 at default.
119.             padding_size = (cur_window + (cur_window - 1) *
120.                             (dilation - 1)) // 2
121.             cur_conv = nn.Conv2d(
122.                 cur_head_split * Ch * 2,
123.                 cur_head_split,
124.                 kernel_size=(cur_window, cur_window),
125.                 padding=(padding_size, padding_size),
126.                 dilation=(dilation, dilation),
127.                 groups=cur_head_split,
128.             )
130.             self.conv_list.append(cur_conv)
131.             self.head_splits.append(cur_head_split)
132.         self.channel_splits = [x * Ch * 2 for x in self.head_splits]
134.     def forward(self, q, k, v, size):
135.         """foward function"""
136.         B, h, N, Ch = q.shape
137.         H, W = size
139.         # We don't use CLS_TOKEN
140.         q_img = q
141.         k_img = k
142.         v_img = v
143.         # Shape: [B, h, H*W, Ch] -> [B, h*Ch, H, W].
144.         q_img = rearrange(q_img, "B h (H W) Ch -> B h Ch H W", H=H, W=W)
145.         k_img = rearrange(k_img, "B h Ch (H W) -> B h Ch H W", H=H, W=W)
146.         qk_concat = torch.cat((q_img, k_img), 2)
147.         qk_concat = rearrange(qk_concat, "B h Ch H W -> B (h Ch) H W", H=H, W=W)
148.         # Split according to channels.
149.         qk_concat_list = torch.split(qk_concat, self.channel_splits, dim=1)
150.         qk_att_list = [
151.             conv(x) for conv, x in zip(self.conv_list, qk_concat_list)
152.         ]
153.         qk_att = torch.cat(qk_att_list, dim=1)
154.         # Shape: [B, h*Ch, H, W] -> [B, h, H*W, Ch].
155.         qk_att = rearrange(qk_att, "B (h Ch) H W -> B h (H W) Ch", h=h)
156.         return qk_att
157. ##########################################################################
158. ## Multi-DConv Head Transposed Self-Attention (MDTA)
159. class Attention(nn.Module):
160.     def __init__(self, dim, num_heads, bias, shared_refine_att=None, qk_norm=1):
161.         super(Attention, self).__init__()
162.         self.norm = qk_norm
163.         self.num_heads = num_heads
164.         self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
165.         # self.Leakyrelu=nn.LeakyReLU(negative_slope=0.01,inplace=True)
166.         self.sigmoid = nn.Sigmoid()
167.         self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
168.         self.qkv_dwconv = nn.Conv2d(dim * 3, dim * 3, kernel_size=3, stride=1, padding=1, groups=dim * 3, bias=bias)
169.         self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
170.         if num_heads == 8:
171.             crpe_window = {
172.                 3: 2,
173.                 5: 3,
174.                 7: 3
175.             }
176.         elif num_heads == 1:
177.             crpe_window = {
178.                 3: 1,
179.             }
180.         elif num_heads == 2:
181.             crpe_window = {
182.                 3: 2,
183.             }
184.         elif num_heads == 4:
185.             crpe_window = {
186.                 3: 2,
187.                 5: 2,
188.             }
189.         self.refine_att = refine_att(Ch=dim // num_heads,
190.                                      h=num_heads,
191.                                      window=crpe_window)
192.     def forward(self, x):
193.         b, c, h, w = x.shape
194.         qkv = self.qkv_dwconv(self.qkv(x))
195.         q, k, v = qkv.chunk(3, dim=1)
196.         q = rearrange(q, 'b (head c) h w -> b head (h w) c', head=self.num_heads)
197.         k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
198.         v = rearrange(v, 'b (head c) h w -> b head (h w) c', head=self.num_heads)
199.         # q = torch.nn.functional.normalize(q, dim=-1)
200.         q_norm = torch.norm(q, p=2, dim=-1, keepdim=True) / self.norm + 1e-6
201.         q = torch.div(q, q_norm)
202.         k_norm = torch.norm(k, p=2, dim=-2, keepdim=True) / self.norm + 1e-6
203.         k = torch.div(k, k_norm)
204.         # k = torch.nn.functional.normalize(k, dim=-2)
205.         refine_weight = self.refine_att(q, k, v, size=(h, w))
206.         # refine_weight=self.Leakyrelu(refine_weight)
207.         refine_weight = self.sigmoid(refine_weight)
208.         attn = k @ v
209.         # attn = attn.softmax(dim=-1)
210.         # print(torch.sum(k, dim=-1).unsqueeze(3).shape)
211.         out_numerator = torch.sum(v, dim=-2).unsqueeze(2) + (q @ attn)
212.         out_denominator = torch.full((h * w, c // self.num_heads), h * w).to(q.device) \
213.                           + q @ torch.sum(k, dim=-1).unsqueeze(3).repeat(1, 1, 1, c // self.num_heads) + 1e-6
214.         # out=torch.div(out_numerator,out_denominator)*self.temperature*refine_weight
215.         out = torch.div(out_numerator, out_denominator) * self.temperature
216.         out = out * refine_weight
217.         out = rearrange(out, 'b head (h w) c-> b (head c) h w', head=self.num_heads, h=h, w=w)
218.         out = self.project_out(out)
219.         return out
220. ##########################################################################
221. class TransformerBlock(nn.Module):
222.     def __init__(self, dim, num_heads, ffn_expansion_factor, bias, LayerNorm_type, shared_refine_att=None, qk_norm=1):
223.         super(TransformerBlock, self).__init__()
224.         self.norm1 = LayerNorm(dim, LayerNorm_type)
225.         self.attn = Attention(dim, num_heads, bias, shared_refine_att=shared_refine_att, qk_norm=qk_norm)
226.         self.norm2 = LayerNorm(dim, LayerNorm_type)
227.         self.ffn = FeedForward(dim, ffn_expansion_factor, bias)
228.     def forward(self, x):
229.         x = x + self.attn(self.norm1(x))
230.         x = x + self.ffn(self.norm2(x))
231.         return x
232. class MHCAEncoder(nn.Module):
233.     """Multi-Head Convolutional self-Attention Encoder comprised of `MHCA`
234.     blocks."""
235.     def __init__(
236.             self,
237.             dim,
238.             num_layers=1,
239.             num_heads=8,
240.             ffn_expansion_factor=2.66,
241.             bias=False,
242.             LayerNorm_type='BiasFree',
243.             qk_norm=1
244.     ):
245.         super().__init__()
246.         self.num_layers = num_layers
247.         self.MHCA_layers = nn.ModuleList([
248.             TransformerBlock(
249.                 dim,
250.                 num_heads=num_heads,
251.                 ffn_expansion_factor=ffn_expansion_factor,
252.                 bias=bias,
253.                 LayerNorm_type=LayerNorm_type,
254.                 qk_norm=qk_norm
255.             ) for idx in range(self.num_layers)
256.         ])
257.     def forward(self, x, size):
258.         """foward function"""
259.         H, W = size
260.         B = x.shape[0]
261.         # return x's shape : [B, N, C] -> [B, C, H, W]
262.         x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
264.         for layer in self.MHCA_layers:
265.             x = layer(x)
267.         return x
270. class ResBlock(nn.Module):
271.     """Residual block for convolutional local feature."""
273.     def __init__(
274.             self,
275.             in_features,
276.             hidden_features=None,
277.             out_features=None,
278.             act_layer=nn.Hardswish,
279.             norm_layer=nn.BatchNorm2d,
280.     ):
281.         super().__init__()
283.         out_features = out_features or in_features
284.         hidden_features = hidden_features or in_features
285.         # self.act0 = act_layer()
286.         self.conv1 = Conv2d_BN(in_features,
287.                                hidden_features,
288.                                act_layer=act_layer)
289.         self.dwconv = nn.Conv2d(
290.             hidden_features,
291.             hidden_features,
292.             3,
293.             1,
294.             1,
295.             bias=False,
296.             groups=hidden_features,
297.         )
298.         # self.norm = norm_layer(hidden_features)
299.         self.act = act_layer()
300.         self.conv2 = Conv2d_BN(hidden_features, out_features)
301.         self.apply(self._init_weights)
303.     def _init_weights(self, m):
304.         """
305.         initialization
306.         """
307.         if isinstance(m, nn.Conv2d):
308.             fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
309.             fan_out //= m.groups
310.             m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
311.             if m.bias is not None:
312.                 m.bias.data.zero_()
314.     def forward(self, x):
315.         """foward function"""
316.         identity = x
317.         # x=self.act0(x)
318.         feat = self.conv1(x)
319.         feat = self.dwconv(feat)
320.         # feat = self.norm(feat)
321.         feat = self.act(feat)
322.         feat = self.conv2(feat)
324.         return identity + feat
327. class MHCA_stage(nn.Module):
328.     """Multi-Head Convolutional self-Attention stage comprised of `MHCAEncoder`
329.     layers."""
331.     def __init__(
332.             self,
333.             embed_dim,
334.             out_embed_dim,
335.             num_layers=1,
336.             num_heads=8,
337.             ffn_expansion_factor=2.66,
338.             num_path=4,
339.             bias=False,
340.             LayerNorm_type='BiasFree',
341.             qk_norm=1
343.     ):
344.         super().__init__()
346.         self.mhca_blks = nn.ModuleList([
347.             MHCAEncoder(
348.                 embed_dim,
349.                 num_layers,
350.                 num_heads,
351.                 ffn_expansion_factor=ffn_expansion_factor,
352.                 bias=bias,
353.                 LayerNorm_type=LayerNorm_type,
354.                 qk_norm=qk_norm
356.             ) for _ in range(num_path)
357.         ])
359.         self.aggregate = SKFF(embed_dim, height=num_path)
361.         # self.InvRes = ResBlock(in_features=embed_dim, out_features=embed_dim)
363.     # self.aggregate = Conv2d_aggregate(embed_dim * (num_path + 1),
364.     #                           out_embed_dim,
365.     #                           act_layer=nn.Hardswish)
367.     def forward(self, inputs):
368.         """foward function"""
369.         # att_outputs = [self.InvRes(inputs[0])]
370.         att_outputs = []
372.         for x, encoder in zip(inputs, self.mhca_blks):
373.             # [B, C, H, W] -> [B, N, C]
374.             _, _, H, W = x.shape
375.             x = x.flatten(2).transpose(1, 2).contiguous()
376.             att_outputs.append(encoder(x, size=(H, W)))
378.         # out_concat = torch.cat(att_outputs, dim=1)
379.         out = self.aggregate(att_outputs)
381.         return out
384. ##########################################################################
385. ## Overlapped image patch embedding with 3x3 Conv
386. class Conv2d_BN(nn.Module):
388.     def __init__(
389.             self,
390.             in_ch,
391.             out_ch,
392.             kernel_size=1,
393.             stride=1,
394.             pad=0,
395.             dilation=1,
396.             groups=1,
397.             bn_weight_init=1,
398.             norm_layer=nn.BatchNorm2d,
399.             act_layer=None,
400.     ):
401.         super().__init__()
403.         self.conv = torch.nn.Conv2d(in_ch,
404.                                     out_ch,
405.                                     kernel_size,
406.                                     stride,
407.                                     pad,
408.                                     dilation,
409.                                     groups,
410.                                     bias=False)
411.         # self.bn = norm_layer(out_ch)
413.         # torch.nn.init.constant_(self.bn.weight, bn_weight_init)
414.         # torch.nn.init.constant_(self.bn.bias, 0)
416.         for m in self.modules():
417.             if isinstance(m, nn.Conv2d):
418.                 # Note that there is no bias due to BN
419.                 fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
420.                 m.weight.data.normal_(mean=0.0, std=np.sqrt(2.0 / fan_out))
422.         self.act_layer = act_layer() if act_layer is not None else nn.Identity()
424.     def forward(self, x):
426.         x = self.conv(x)
427.         # x = self.bn(x)
428.         x = self.act_layer(x)
430.         return x
433. class SKFF(nn.Module):
434.     def __init__(self, in_channels, height=2, reduction=8, bias=False):
435.         super(SKFF, self).__init__()
437.         self.height = height
438.         d = max(int(in_channels / reduction), 4)
440.         self.avg_pool = nn.AdaptiveAvgPool2d(1)
441.         self.conv_du = nn.Sequential(nn.Conv2d(in_channels, d, 1, padding=0, bias=bias), nn.PReLU())
443.         self.fcs = nn.ModuleList([])
444.         for i in range(self.height):
445.             self.fcs.append(nn.Conv2d(d, in_channels, kernel_size=1, stride=1, bias=bias))
447.         self.softmax = nn.Softmax(dim=1)
449.     def forward(self, inp_feats):
450.         batch_size = inp_feats[0].shape[0]
451.         n_feats = inp_feats[0].shape[1]
453.         inp_feats = torch.cat(inp_feats, dim=1)
454.         inp_feats = inp_feats.view(batch_size, self.height, n_feats, inp_feats.shape[2], inp_feats.shape[3])
456.         feats_U = torch.sum(inp_feats, dim=1)
457.         feats_S = self.avg_pool(feats_U)
458.         feats_Z = self.conv_du(feats_S)
460.         attention_vectors = [fc(feats_Z) for fc in self.fcs]
461.         attention_vectors = torch.cat(attention_vectors, dim=1)
462.         attention_vectors = attention_vectors.view(batch_size, self.height, n_feats, 1, 1)
463.         # stx()
464.         attention_vectors = self.softmax(attention_vectors)
466.         feats_V = torch.sum(inp_feats * attention_vectors, dim=1)
468.         return feats_V
471. class DWConv2d_BN(nn.Module):
473.     def __init__(
474.             self,
475.             in_ch,
476.             out_ch,
477.             kernel_size=1,
478.             stride=1,
479.             norm_layer=nn.BatchNorm2d,
480.             act_layer=nn.Hardswish,
481.             bn_weight_init=1,
482.             offset_clamp=(-1, 1)
483.     ):
484.         super().__init__()
485.         # dw
486.         # self.conv=torch.nn.Conv2d(in_ch,out_ch,kernel_size,stride,(kernel_size - 1) // 2,bias=False,)
487.         # self.mask_generator = nn.Sequential(nn.Conv2d(in_channels=in_ch, out_channels=in_ch, kernel_size=3,
488.         #                                                 stride=1, padding=1, bias=False, groups=in_ch),
489.         #                                       nn.Conv2d(in_channels=in_ch, out_channels=9,
490.         #                                                 kernel_size=1,
491.         #                                                 stride=1, padding=0, bias=False)
492.         #                                      )
493.         self.offset_clamp = offset_clamp
494.         self.offset_generator = nn.Sequential(nn.Conv2d(in_channels=in_ch, out_channels=in_ch, kernel_size=3,
495.                                                         stride=1, padding=1, bias=False, groups=in_ch),
496.                                               nn.Conv2d(in_channels=in_ch, out_channels=18,
497.                                                         kernel_size=1,
498.                                                         stride=1, padding=0, bias=False)
500.                                               )
501.         self.dcn = DeformConv2d(
502.             in_channels=in_ch,
503.             out_channels=in_ch,
504.             kernel_size=3,
505.             stride=1,
506.             padding=1,
507.             bias=False,
508.             groups=in_ch
509.         )  # .cuda(7)
510.         self.pwconv = nn.Conv2d(in_ch, out_ch, 1, 1, 0, bias=False)
512.         # self.bn = norm_layer(out_ch)
513.         self.act = act_layer() if act_layer is not None else nn.Identity()
514.         for m in self.modules():
515.             if isinstance(m, nn.Conv2d):
516.                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
517.                 m.weight.data.normal_(0, math.sqrt(2.0 / n))
518.                 if m.bias is not None:
519.                     m.bias.data.zero_()
520.                 # print(m)
522.         #   elif isinstance(m, nn.BatchNorm2d):
523.         #     m.weight.data.fill_(bn_weight_init)
524.         #      m.bias.data.zero_()
526.     def forward(self, x):
528.         # x=self.conv(x)
529.         # x = self.bn(x)
530.         # x = self.act(x)
531.         # mask= torch.sigmoid(self.mask_generator(x))
532.         # print('1')
533.         offset = self.offset_generator(x)
534.         # print('2')
535.         if self.offset_clamp:
536.             offset = torch.clamp(offset, min=self.offset_clamp[0], max=self.offset_clamp[1])  # .cuda(7)1
537.         # print(offset)
538.         # print('3')
539.         # x=x.cuda(7)
540.         x = self.dcn(x, offset)
541.         # x=x.cpu()
542.         # print('4')
543.         x = self.pwconv(x)
544.         # print('5')
545.         # x = self.bn(x)
546.         x = self.act(x)
547.         return x
550. class DWCPatchEmbed(nn.Module):
551.     """Depthwise Convolutional Patch Embedding layer Image to Patch
552.     Embedding."""
554.     def __init__(self,
555.                  in_chans=3,
556.                  embed_dim=768,
557.                  patch_size=16,
558.                  stride=1,
559.                  idx=0,
560.                  act_layer=nn.Hardswish,
561.                  offset_clamp=(-1, 1)):
562.         super().__init__()
564.         self.patch_conv = DWConv2d_BN(
565.             in_chans,
566.             embed_dim,
567.             kernel_size=patch_size,
568.             stride=stride,
569.             act_layer=act_layer,
570.             offset_clamp=offset_clamp
571.         )
572.         """
573.         self.patch_conv = DWConv2d_BN(
574.             in_chans,
575.             embed_dim,
576.             kernel_size=patch_size,
577.             stride=stride,
578.             act_layer=act_layer,
579.         )
580.         """
582.     def forward(self, x):
583.         """foward function"""
584.         x = self.patch_conv(x)
586.         return x
589. class Patch_Embed_stage(nn.Module):
590.     """Depthwise Convolutional Patch Embedding stage comprised of
591.     `DWCPatchEmbed` layers."""
593.     def __init__(self, in_chans, embed_dim, num_path=4, isPool=False, offset_clamp=(-1, 1)):
594.         super(Patch_Embed_stage, self).__init__()
596.         self.patch_embeds = nn.ModuleList([
597.             DWCPatchEmbed(
598.                 in_chans=in_chans if idx == 0 else embed_dim,
599.                 embed_dim=embed_dim,
600.                 patch_size=3,
601.                 stride=1,
602.                 idx=idx,
603.                 offset_clamp=offset_clamp
604.             ) for idx in range(num_path)
605.         ])
607.     def forward(self, x):
608.         """foward function"""
609.         att_inputs = []
610.         for pe in self.patch_embeds:
611.             x = pe(x)
612.             att_inputs.append(x)
614.         return att_inputs
617. class OverlapPatchEmbed(nn.Module):
618.     def __init__(self, in_c=3, embed_dim=48, bias=False):
619.         super(OverlapPatchEmbed, self).__init__()
620.         self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=3, stride=1, padding=1, bias=bias)
621.         # self.proj_dw = nn.Conv2d(in_c, in_c, kernel_size=3, stride=1, padding=1,groups=in_c, bias=bias)
622.         # self.proj_pw = nn.Conv2d(in_c, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias)
624.         # self.bn=nn.BatchNorm2d(embed_dim)
625.         # self.act=nn.Hardswish()
627.     def forward(self, x):
628.         x = self.proj(x)
629.         # x = self.proj_dw(x)
630.         # x= self.proj_pw(x)
631.         # x=self.bn(x)
632.         # x=self.act(x)
633.         return x
636. ##########################################################################
637. ## Resizing modules
638. class Downsample(nn.Module):
639.     def __init__(self, input_feat, out_feat):
640.         super(Downsample, self).__init__()
642.         self.body = nn.Sequential(  # nn.Conv2d(n_feat, n_feat // 2, kernel_size=3, stride=1, padding=1, bias=False),
643.             # dw
644.             nn.Conv2d(input_feat, input_feat, kernel_size=3, stride=1, padding=1, groups=input_feat, bias=False, ),
645.             # pw-linear
646.             nn.Conv2d(input_feat, out_feat // 4, 1, 1, 0, bias=False),
647.             # nn.BatchNorm2d(n_feat // 2),
648.             # nn.Hardswish(),
649.             nn.PixelUnshuffle(2))
651.     def forward(self, x):
652.         return self.body(x)
655. class Upsample(nn.Module):
656.     def __init__(self, input_feat, out_feat):
657.         super(Upsample, self).__init__()
659.         self.body = nn.Sequential(  # nn.Conv2d(n_feat, n_feat*2, kernel_size=3, stride=1, padding=1, bias=False),
660.             # dw
661.             nn.Conv2d(input_feat, input_feat, kernel_size=3, stride=1, padding=1, groups=input_feat, bias=False, ),
662.             # pw-linear
663.             nn.Conv2d(input_feat, out_feat * 4, 1, 1, 0, bias=False),
664.             # nn.BatchNorm2d(n_feat*2),
665.             # nn.Hardswish(),
666.             nn.PixelShuffle(2))
668.     def forward(self, x):
669.         return self.body(x)
672. ##########################################################################
673. ##---------- Restormer -----------------------
674. class MB_TaylorFormer(nn.Module):
675.     def __init__(self,
676.                  inp_channels=3,
677.                  dim=[6, 12, 24, 36],
678.                  num_blocks=[1, 1, 1, 1],
679.                  heads=[1, 1, 1, 1],
680.                  bias=False,
681.                  dual_pixel_task=True,
682.                  num_path=[1, 1, 1, 1],  ## True for dual-pixel defocus deblurring only. Also set inp_channels=6
683.                  qk_norm=1,
684.                  offset_clamp=(-1, 1)
685.                  ):
687.         super(MB_TaylorFormer, self).__init__()
689.         self.patch_embed = OverlapPatchEmbed(inp_channels, dim[0])
690.         self.patch_embed_encoder_level1 = Patch_Embed_stage(dim[0], dim[0], num_path=num_path[0], isPool=False,
691.                                                             offset_clamp=offset_clamp)
692.         self.encoder_level1 = MHCA_stage(dim[0], dim[0], num_layers=num_blocks[0], num_heads=heads[0],
693.                                          ffn_expansion_factor=2.66, num_path=num_path[0],
694.                                          bias=False, LayerNorm_type='BiasFree', qk_norm=qk_norm)
696.         self.down1_2 = Downsample(dim[0], dim[1])  ## From Level 1 to Level 2
698.         self.patch_embed_encoder_level2 = Patch_Embed_stage(dim[1], dim[1], num_path=num_path[1], isPool=False,
699.                                                             offset_clamp=offset_clamp)
700.         self.encoder_level2 = MHCA_stage(dim[1], dim[1], num_layers=num_blocks[1], num_heads=heads[1],
701.                                          ffn_expansion_factor=2.66,
702.                                          num_path=num_path[1], bias=False, LayerNorm_type='BiasFree', qk_norm=qk_norm)
704.         self.down2_3 = Downsample(dim[1], dim[2])  ## From Level 2 to Level 3
706.         self.patch_embed_encoder_level3 = Patch_Embed_stage(dim[2], dim[2], num_path=num_path[2],
707.                                                             isPool=False, offset_clamp=offset_clamp)
708.         self.encoder_level3 = MHCA_stage(dim[2], dim[2], num_layers=num_blocks[2], num_heads=heads[2],
709.                                          ffn_expansion_factor=2.66,
710.                                          num_path=num_path[2], bias=False, LayerNorm_type='BiasFree', qk_norm=qk_norm)
712.         self.down3_4 = Downsample(dim[2], dim[3])  ## From Level 3 to Level 4
714.         self.patch_embed_latent = Patch_Embed_stage(dim[3], dim[3], num_path=num_path[3],
715.                                                     isPool=False, offset_clamp=offset_clamp)
716.         self.latent = MHCA_stage(dim[3], dim[3], num_layers=num_blocks[3], num_heads=heads[3],
717.                                  ffn_expansion_factor=2.66, num_path=num_path[3], bias=False,
718.                                  LayerNorm_type='BiasFree', qk_norm=qk_norm)
720.         self.up4_3 = Upsample(int(dim[3]), dim[2])  ## From Level 4 to Level 3
721.         self.reduce_chan_level3 = nn.Sequential(
722.             nn.Conv2d(dim[2] * 2, dim[2], 1, 1, 0, bias=bias),
723.             # nn.BatchNorm2d(dim * 2**2),
724.             # nn.Hardswish(),
725.         )
727.         self.patch_embed_decoder_level3 = Patch_Embed_stage(dim[2], dim[2], num_path=num_path[2],
728.                                                             isPool=False, offset_clamp=offset_clamp)
729.         self.decoder_level3 = MHCA_stage(dim[2], dim[2], num_layers=num_blocks[2], num_heads=heads[2],
730.                                          ffn_expansion_factor=2.66, num_path=num_path[2], bias=False,
731.                                          LayerNorm_type='BiasFree', qk_norm=qk_norm)
733.         self.up3_2 = Upsample(int(dim[2]), dim[1])  ## From Level 3 to Level 2
734.         self.reduce_chan_level2 = nn.Sequential(
735.             nn.Conv2d(dim[1] * 2, dim[1], 1, 1, 0, bias=bias),
736.             # nn.BatchNorm2d( dim * 2),
737.             # nn.Hardswish(),
738.         )
740.         self.patch_embed_decoder_level2 = Patch_Embed_stage(dim[1], dim[1], num_path=num_path[1],
741.                                                             isPool=False, offset_clamp=offset_clamp)
742.         self.decoder_level2 = MHCA_stage(dim[1], dim[1], num_layers=num_blocks[1], num_heads=heads[1],
743.                                          ffn_expansion_factor=2.66, num_path=num_path[1], bias=False,
744.                                          LayerNorm_type='BiasFree', qk_norm=qk_norm)
746.         self.up2_1 = Upsample(int(dim[1]), dim[0])  ## From Level 2 to Level 1  (NO 1x1 conv to reduce channels)
748.         self.patch_embed_decoder_level1 = Patch_Embed_stage(dim[1], dim[1], num_path=num_path[0],
749.                                                             isPool=False, offset_clamp=offset_clamp)
750.         self.decoder_level1 = MHCA_stage(dim[1], dim[1], num_layers=num_blocks[0], num_heads=heads[0],
751.                                          ffn_expansion_factor=2.66, num_path=num_path[0], bias=False,
752.                                          LayerNorm_type='BiasFree', qk_norm=qk_norm)
754.         self.patch_embed_refinement = Patch_Embed_stage(dim[1], dim[1], num_path=num_path[0],
755.                                                         isPool=False, offset_clamp=offset_clamp)
756.         self.refinement = MHCA_stage(dim[1], dim[1], num_layers=num_blocks[0], num_heads=heads[0],
757.                                      ffn_expansion_factor=2.66, num_path=num_path[0], bias=False,
758.                                      LayerNorm_type='BiasFree', qk_norm=qk_norm)
760.         #### For Dual-Pixel Defocus Deblurring Task ####
761.         self.dual_pixel_task = dual_pixel_task
762.         if self.dual_pixel_task:
763.             self.skip_conv = nn.Conv2d(dim[0], dim[1], kernel_size=1, bias=bias)
764.         ###########################
766.         # self.output = nn.Conv2d(dim*2**1, 3, kernel_size=3, stride=1, padding=1, bias=False)
768.         self.output = nn.Sequential(  # nn.Conv2d(n_feat, n_feat*2, kernel_size=3, stride=1, padding=1, bias=False),
769.             # nn.BatchNorm2d(dim*2),
770.             # nn.Hardswish(),
771.             nn.Conv2d(dim[1], 3, kernel_size=3, stride=1, padding=1, bias=False, ),
773.         )
775.     def forward(self, inp_img):
776.         inp_enc_level1 = self.patch_embed(inp_img)
778.         inp_enc_level1_list = self.patch_embed_encoder_level1(inp_enc_level1)
780.         out_enc_level1 = self.encoder_level1(inp_enc_level1_list) + inp_enc_level1
782.         # out_enc_level1 = self.encoder_level1(inp_enc_level1_list)
783.         inp_enc_level2 = self.down1_2(out_enc_level1)
785.         inp_enc_level2_list = self.patch_embed_encoder_level2(inp_enc_level2)
786.         out_enc_level2 = self.encoder_level2(inp_enc_level2_list) + inp_enc_level2
787.         inp_enc_level3 = self.down2_3(out_enc_level2)
789.         inp_enc_level3_list = self.patch_embed_encoder_level3(inp_enc_level3)
790.         out_enc_level3 = self.encoder_level3(inp_enc_level3_list) + inp_enc_level3
791.         inp_enc_level4 = self.down3_4(out_enc_level3)
793.         inp_latent = self.patch_embed_latent(inp_enc_level4)
794.         latent = self.latent(inp_latent) + inp_enc_level4
796.         inp_dec_level3 = self.up4_3(latent)
797.         inp_dec_level3 = torch.cat([inp_dec_level3, out_enc_level3], 1)
798.         inp_dec_level3 = self.reduce_chan_level3(inp_dec_level3)
799.         inp_dec_level3_list = self.patch_embed_decoder_level3(inp_dec_level3)
800.         out_dec_level3 = self.decoder_level3(inp_dec_level3_list) + inp_dec_level3
802.         inp_dec_level2 = self.up3_2(out_dec_level3)
804.         inp_dec_level2 = torch.cat([inp_dec_level2, out_enc_level2], 1)
806.         inp_dec_level2 = self.reduce_chan_level2(inp_dec_level2)
808.         inp_dec_level2_list = self.patch_embed_decoder_level2(inp_dec_level2)
810.         out_dec_level2 = self.decoder_level2(inp_dec_level2_list) + inp_dec_level2
812.         inp_dec_level1 = self.up2_1(out_dec_level2)
814.         inp_dec_level1 = torch.cat([inp_dec_level1, out_enc_level1], 1)
816.         inp_dec_level1_list = self.patch_embed_decoder_level1(inp_dec_level1)
818.         out_dec_level1 = self.decoder_level1(inp_dec_level1_list) + inp_dec_level1
820.         inp_latent_list = self.patch_embed_refinement(out_dec_level1)
822.         out_dec_level1 = self.refinement(inp_latent_list) + out_dec_level1
824.         # nn.Hardswish()
826.         #### For Dual-Pixel Defocus Deblurring Task ####
827.         if self.dual_pixel_task:
828.             out_dec_level1 = out_dec_level1 + self.skip_conv(inp_enc_level1)
829.             out_dec_level1 = self.output(out_dec_level1)
830.         ###########################
831.         else:
832.             out_dec_level1 = self.output(out_dec_level1) + inp_img
834.         return out_dec_level1
837. def count_param(model):
838.     param_count = 0
839.     for param in model.parameters():
840.         param_count += param.view(-1).size()[0]
841.     return param_count
844. if __name__ == "__main__":
845.     from thop import profile
847.     model = MB_TaylorFormer()
848.     model.eval()
849.     print("params", count_param(model))
850.     inputs = torch.randn(1, 3, 640, 640)
851.     output = model(inputs)
852.     print(output.size())

复制代码

2.2 步调二
在Addmodules下创建一个新的py文件名字为'__init__.py',然后在其内部添加如下代码

2.3 步调三
在task.py举行导入

到此注册成功,复制背面的yaml文件直接运行即可
yaml文件
[code]
# Ultralytics YOLO