devbox
不正经
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

pytorch实现常用的一些即插即用模块(长期更新)_torch_dwconv

作者：不正经 | 2024-03-22 21:20:20

赞

踩

torch_dwconv

1.可分离卷积

  1. #coding:utf-8
  2. import torch.nn as nn
  3.  
  4. class DWConv(nn.Module):
  5.     def __init__(self, in_plane, out_plane):
  6.         super(DWConv, self).__init__()
  7.         self.depth_conv = nn.Conv2d(in_channels=in_plane,
  8.                                     out_channels=in_plane,
  9.                                     kernel_size=3,
  10.                                     stride=1,
  11.                                     padding=1,
  12.                                     groups=in_plane)
  13.         self.point_conv = nn.Conv2d(in_channels=in_plane,
  14.                                     out_channels=out_plane,
  15.                                     kernel_size=1,
  16.                                     stride=1,
  17.                                     padding=0,
  18.                                     groups=1)
  19.  
  20.     def forward(self, x):
  21.         x = self.depth_conv(x)
  22.         x = self.point_conv(x)
  23.         return x
  24.  
  25. def deubg_dw():
  26.     import torch
  27.     DW_model = DWConv(3, 32)
  28.     x = torch.rand((32, 3, 320, 320))
  29.     out = DW_model(x)
  30.     print(out.shape)
  31. if __name__ == '__main__':
  32.     deubg_dw()

2.DBnet论文中的DBhead

  1. #coding:utf-8
  2. import torch
  3. from torch import nn
  4.  
  5. class DBHead(nn.Module):
  6.     def __init__(self, in_channels, out_channels, k=50):
  7.         super().__init__()
  8.         self.k = k
  9.         self.binarize = nn.Sequential(
  10.             nn.Conv2d(in_channels, in_channels // 4, 3, padding=1),
  11.             nn.BatchNorm2d(in_channels // 4),
  12.             nn.ReLU(inplace=True),
  13.             nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 2, 2),
  14.             nn.BatchNorm2d(in_channels // 4),
  15.             nn.ReLU(inplace=True),
  16.             nn.ConvTranspose2d(in_channels // 4, 1, 2, 2),
  17.             nn.Sigmoid())
  18.         self.binarize.apply(self.weights_init)
  19.  
  20.         self.thresh = self._init_thresh(in_channels)
  21.         self.thresh.apply(self.weights_init)
  22.  
  23.     def forward(self, x):
  24.         shrink_maps = self.binarize(x)
  25.         threshold_maps = self.thresh(x)
  26.         if self.training:#从父类继承的变量, train的时候默认是true, eval的时候会变为false
  27.             binary_maps = self.step_function(shrink_maps, threshold_maps)
  28.             y = torch.cat((shrink_maps, threshold_maps, binary_maps), dim=1)
  29.         else:
  30.             y = torch.cat((shrink_maps, threshold_maps), dim=1)
  31.         return y
  32.  
  33.     def weights_init(self, m):
  34.         classname = m.__class__.__name__
  35.         if classname.find('Conv') != -1:
  36.             nn.init.kaiming_normal_(m.weight.data)
  37.         elif classname.find('BatchNorm') != -1:
  38.             m.weight.data.fill_(1.)
  39.             m.bias.data.fill_(1e-4)
  40.  
  41.     def _init_thresh(self, inner_channels, serial=False, smooth=False, bias=False):
  42.         in_channels = inner_channels
  43.         if serial:
  44.             in_channels += 1
  45.         self.thresh = nn.Sequential(
  46.             nn.Conv2d(in_channels, inner_channels // 4, 3, padding=1, bias=bias),
  47.             nn.BatchNorm2d(inner_channels // 4),
  48.             nn.ReLU(inplace=True),
  49.             self._init_upsample(inner_channels // 4, inner_channels // 4, smooth=smooth, bias=bias),
  50.             nn.BatchNorm2d(inner_channels // 4),
  51.             nn.ReLU(inplace=True),
  52.             self._init_upsample(inner_channels // 4, 1, smooth=smooth, bias=bias),
  53.             nn.Sigmoid())
  54.         return self.thresh
  55.  
  56.     def _init_upsample(self, in_channels, out_channels, smooth=False, bias=False):
  57.         if smooth:
  58.             inter_out_channels = out_channels
  59.             if out_channels == 1:
  60.                 inter_out_channels = in_channels
  61.             module_list = [
  62.                 nn.Upsample(scale_factor=2, mode='nearest'),
  63.                 nn.Conv2d(in_channels, inter_out_channels, 3, 1, 1, bias=bias)]
  64.             if out_channels == 1:
  65.                 module_list.append(nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=1, bias=True))
  66.             return nn.Sequential(module_list)
  67.         else:
  68.             return nn.ConvTranspose2d(in_channels, out_channels, 2, 2)
  69.  
  70.     def step_function(self, x, y):
  71.         return torch.reciprocal(1 + torch.exp(-self.k * (x - y)))
  72.  
  73. def debug_main():
  74.     x = torch.rand((8, 256, 160, 160))
  75.  
  76.     head_model = DBHead(in_channels=256, out_channels=2)
  77.     head_model.train()
  78.     y = head_model(x)
  79.     print('==y.shape:', y.shape)
  80.  
  81.     head_model.eval()
  82.     y = head_model(x)
  83.     print('==y.shape:', y.shape)
  84.  
  85. if __name__ == '__main__':
  86.     debug_main()

3.sENet中的attention

目的对于不同通道进行加权,先squeeze将h*w*c　global averge pooling成1*1*c特征，在经过两层线性层，通过sigmoid输出加权在不同通道。

  1.  
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5. class SELayer(nn.Module):
  6.     def __init__(self, channel, reduction=16):
  7.         super(SELayer, self).__init__()
  8.         self.avg_pool = nn.AdaptiveAvgPool2d(1) # 压缩空间
  9.         self.fc = nn.Sequential(
  10.             nn.Linear(channel, channel // reduction, bias=False),
  11.             nn.ReLU(inplace=True),
  12.             nn.Linear(channel // reduction, channel, bias=False),
  13.             nn.Sigmoid()
  14.         )
  15.  
  16.     def forward(self, x):
  17.         b, c, _, _ = x.size()
  18.         y = self.avg_pool(x).view(b, c)
  19.         y = self.fc(y).view(b, c, 1, 1)
  20.         return x * y
  21.  
  22. def debug_attention():
  23.     attention_module = SELayer(channel=128, reduction=16)
  24.     # B,C,H,W
  25.     x = torch.rand((2, 128, 100, 100))
  26.     out = attention_module(x)
  27.     print('==out.shape:', out.shape)
  28.  
  29. if __name__ == '__main__':
  30.     debug_attention()

4.cv中的self-attention

(1).feature map通过1*1卷积获得,q,k,v三个向量,q与v转置相乘得到attention矩阵，进行softmax归一化到0到1,在作用于V,得到每个像素的加权.

(2).softmax

(3).加权求和

  1.  
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5.  
  6. class Self_Attn(nn.Module):
  7.     """ Self attention Layer"""
  8.  
  9.     def __init__(self, in_dim):
  10.         super(Self_Attn, self).__init__()
  11.         self.chanel_in = in_dim
  12.  
  13.         self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
  14.         self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
  15.         self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
  16.         self.gamma = nn.Parameter(torch.zeros(1))
  17.  
  18.         self.softmax = nn.Softmax(dim=-1)
  19.  
  20.     def forward(self, x):
  21.         """
  22.             inputs :
  23.                 x : input feature maps( B * C * W * H)
  24.             returns :
  25.                 out : self attention value + input feature
  26.                 attention: B * N * N (N is Width*Height)
  27.         """
  28.         m_batchsize, C, width, height = x.size()
  29.         proj_query = self.query_conv(x).view(m_batchsize, -1, width * height).permute(0, 2, 1)  # B*N*C
  30.         proj_key = self.key_conv(x).view(m_batchsize, -1, width * height)  # B*C*N
  31.         energy = torch.bmm(proj_query, proj_key)  # batch的matmul B*N*N
  32.         attention = self.softmax(energy)  # B * (N) * (N)
  33.         proj_value = self.value_conv(x).view(m_batchsize, -1, width * height)  # B * C * N
  34.  
  35.         out = torch.bmm(proj_value, attention.permute(0, 2, 1))  # B*C*N
  36.         out = out.view(m_batchsize, C, width, height)  # B*C*H*W
  37.  
  38.         out = self.gamma * out + x
  39.         return out, attention
  40.  
  41.  
  42. def debug_attention():
  43.     attention_module = Self_Attn(in_dim=128)
  44.     #B,C,H,W
  45.     x = torch.rand((2, 128, 100, 100))
  46.     attention_module(x)
  47.  
  48.  
  49. if __name__ == '__main__':
  50.     debug_attention()

5.spp多窗口pooling

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4. class SPP(nn.Module):
  5.     """
  6.         Spatial Pyramid Pooling
  7.     """
  8.  
  9.     def __init__(self):
  10.         super(SPP, self).__init__()
  11.  
  12.     def forward(self, x):
  13.         x_1 = F.max_pool2d(x, kernel_size=5, stride=1, padding=2)
  14.         x_2 = F.max_pool2d(x, kernel_size=9, stride=1, padding=4)
  15.         x_3 = F.max_pool2d(x, kernel_size=13, stride=1, padding=6)
  16.         x = torch.cat([x, x_1, x_2, x_3], dim=1)
  17.  
  18.         return x
  19.  
  20. def debug_spp():
  21.     x = torch.rand((8,3,256,256))
  22.     spp = SPP()
  23.     x = spp(x)
  24.     print('==x.shape:', x.shape)
  25.  
  26. if __name__ == '__main__':
  27.     debug_spp()

6.RetinaFPN

  1. # coding: utf-8
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5.  
  6.  
  7. class RetinaFPN(nn.Module):
  8.     def __init__(self,
  9.                  C3_inplanes,
  10.                  C4_inplanes,
  11.                  C5_inplanes,
  12.                  planes,
  13.                  use_p5=False):
  14.         super(RetinaFPN, self).__init__()
  15.         self.use_p5 = use_p5
  16.         self.P3_1 = nn.Conv2d(C3_inplanes,
  17.                               planes,
  18.                               kernel_size=1,
  19.                               stride=1,
  20.                               padding=0)
  21.         self.P3_2 = nn.Conv2d(planes,
  22.                               planes,
  23.                               kernel_size=3,
  24.                               stride=1,
  25.                               padding=1)
  26.         self.P4_1 = nn.Conv2d(C4_inplanes,
  27.                               planes,
  28.                               kernel_size=1,
  29.                               stride=1,
  30.                               padding=0)
  31.         self.P4_2 = nn.Conv2d(planes,
  32.                               planes,
  33.                               kernel_size=3,
  34.                               stride=1,
  35.                               padding=1)
  36.         self.P5_1 = nn.Conv2d(C5_inplanes,
  37.                               planes,
  38.                               kernel_size=1,
  39.                               stride=1,
  40.                               padding=0)
  41.         self.P5_2 = nn.Conv2d(planes,
  42.                               planes,
  43.                               kernel_size=3,
  44.                               stride=1,
  45.                               padding=1)
  46.         if self.use_p5:
  47.             self.P6 = nn.Conv2d(planes,
  48.                                 planes,
  49.                                 kernel_size=3,
  50.                                 stride=2,
  51.                                 padding=1)
  52.         else:
  53.             self.P6 = nn.Conv2d(C5_inplanes,
  54.                                 planes,
  55.                                 kernel_size=3,
  56.                                 stride=2,
  57.                                 padding=1)
  58.  
  59.         self.P7 = nn.Sequential(
  60.             nn.ReLU(),
  61.             nn.Conv2d(planes, planes, kernel_size=3, stride=2, padding=1))
  62.  
  63.     def forward(self, inputs):
  64.         [C3, C4, C5] = inputs
  65.  
  66.         P5 = self.P5_1(C5)
  67.         P4 = self.P4_1(C4)
  68.         P4 = F.interpolate(P5, size=(P4.shape[2], P4.shape[3]),
  69.                            mode='nearest') + P4
  70.         P3 = self.P3_1(C3)
  71.         P3 = F.interpolate(P4, size=(P3.shape[2], P3.shape[3]),
  72.                            mode='nearest') + P3
  73.  
  74.         P5 = self.P5_2(P5)
  75.         P4 = self.P4_2(P4)
  76.         P3 = self.P3_2(P3)
  77.  
  78.         if self.use_p5:
  79.             P6 = self.P6(P5)
  80.         else:
  81.             P6 = self.P6(C5)
  82.  
  83.         del C3, C4, C5
  84.  
  85.         P7 = self.P7(P6)
  86.  
  87.         return [P3, P4, P5, P6, P7]
  88.  
  89. if __name__ == '__main__':
  90.     image_h, image_w = 640, 640
  91.     fpn = RetinaFPN(512, 1024, 2048, 256)
  92.  
  93.     C3, C4, C5 = torch.randn(3, 512, 80, 80), torch.randn(3, 1024, 40, 40), torch.randn(3, 2048, 20, 20)
  94.     [P3, P4, P5, P6, P7] = fpn([C3, C4, C5])
  95.  
  96.     print("P3", P3.shape)
  97.     print("P4", P4.shape)
  98.     print("P5", P5.shape)
  99.     print("P6", P6.shape)
  100.     print("P7", P7.shape)

7.Focus

  1.  
  2. import torch
  3. import torch.nn as nn
  4.  
  5. def autopad(k, p=None):  # kernel, padding
  6.     # Pad to 'same'
  7.     if p is None:
  8.         p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
  9.     # print('==p:', p)
  10.     return p
  11.  
  12. class Conv(nn.Module):
  13.     # Standard convolution
  14.     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
  15.         super(Conv, self).__init__()
  16.         self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
  17.         self.bn = nn.BatchNorm2d(c2)
  18.         self.act = nn.Hardswish() if act else nn.Identity()
  19.  
  20.     def forward(self, x):
  21.         return self.act(self.bn(self.conv(x)))
  22.  
  23.     def fuseforward(self, x):
  24.         return self.act(self.conv(x))
  25.  
  26. class Focus(nn.Module):
  27.     # Focus wh information into c-space
  28.     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
  29.         super(Focus, self).__init__()
  30.         self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
  31.  
  32.     def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
  33.         return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
  34.  
  35.  
  36. def debug_focus():
  37.     model = Focus(c1=3, c2=24)
  38.     img = torch.rand((8, 3, 124, 124))
  39.     print('==img.shape', img.shape)
  40.     out = model(img)
  41.     print('===out.shape', out.shape)
  42.  
  43. debug_focus()

本文内容由网友自发贡献，转载请注明出处：https://www.wpsshop.cn/w/不正经/article/detail/290957
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号