Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.

134 Zeilen
4.8KB

  1. # This file contains modules common to various models
  2. import torch.nn.functional as F
  3. from utils.utils import *
  4. def DWConv(c1, c2, k=1, s=1, act=True):
  5. # Depthwise convolution
  6. return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
  7. class Conv(nn.Module):
  8. # Standard convolution
  9. def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
  10. super(Conv, self).__init__()
  11. self.conv = nn.Conv2d(c1, c2, k, s, k // 2, groups=g, bias=False)
  12. self.bn = nn.BatchNorm2d(c2)
  13. self.act = nn.LeakyReLU(0.1, inplace=True) if act else nn.Identity()
  14. def forward(self, x):
  15. return self.act(self.bn(self.conv(x)))
  16. def fuseforward(self, x):
  17. return self.act(self.conv(x))
  18. class Bottleneck(nn.Module):
  19. # Standard bottleneck
  20. def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
  21. super(Bottleneck, self).__init__()
  22. c_ = int(c2 * e) # hidden channels
  23. self.cv1 = Conv(c1, c_, 1, 1)
  24. self.cv2 = Conv(c_, c2, 3, 1, g=g)
  25. self.add = shortcut and c1 == c2
  26. def forward(self, x):
  27. return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
  28. class BottleneckCSP(nn.Module):
  29. # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
  30. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
  31. super(BottleneckCSP, self).__init__()
  32. c_ = int(c2 * e) # hidden channels
  33. self.cv1 = Conv(c1, c_, 1, 1)
  34. self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
  35. self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
  36. self.cv4 = Conv(c2, c2, 1, 1)
  37. self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
  38. self.act = nn.LeakyReLU(0.1, inplace=True)
  39. self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
  40. def forward(self, x):
  41. y1 = self.cv3(self.m(self.cv1(x)))
  42. y2 = self.cv2(x)
  43. return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
  44. class ConvPlus(nn.Module):
  45. # Plus-shaped convolution
  46. def __init__(self, c1, c2, k=3, s=1, g=1, bias=True): # ch_in, ch_out, kernel, stride, groups
  47. super(ConvPlus, self).__init__()
  48. self.cv1 = nn.Conv2d(c1, c2, (k, 1), s, (k // 2, 0), groups=g, bias=bias)
  49. self.cv2 = nn.Conv2d(c1, c2, (1, k), s, (0, k // 2), groups=g, bias=bias)
  50. def forward(self, x):
  51. return self.cv1(x) + self.cv2(x)
  52. class SPP(nn.Module):
  53. # Spatial pyramid pooling layer used in YOLOv3-SPP
  54. def __init__(self, c1, c2, k=(5, 9, 13)):
  55. super(SPP, self).__init__()
  56. c_ = c1 // 2 # hidden channels
  57. self.cv1 = Conv(c1, c_, 1, 1)
  58. self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
  59. self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
  60. def forward(self, x):
  61. x = self.cv1(x)
  62. return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
  63. class Flatten(nn.Module):
  64. # Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions
  65. def forward(self, x):
  66. return x.view(x.size(0), -1)
  67. class Focus(nn.Module):
  68. # Focus wh information into c-space
  69. def __init__(self, c1, c2, k=1):
  70. super(Focus, self).__init__()
  71. self.conv = Conv(c1 * 4, c2, k, 1)
  72. def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
  73. return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
  74. class Concat(nn.Module):
  75. # Concatenate a list of tensors along dimension
  76. def __init__(self, dimension=1):
  77. super(Concat, self).__init__()
  78. self.d = dimension
  79. def forward(self, x):
  80. return torch.cat(x, self.d)
  81. class MixConv2d(nn.Module):
  82. # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
  83. def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
  84. super(MixConv2d, self).__init__()
  85. groups = len(k)
  86. if equal_ch: # equal c_ per group
  87. i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices
  88. c_ = [(i == g).sum() for g in range(groups)] # intermediate channels
  89. else: # equal weight.numel() per group
  90. b = [c2] + [0] * groups
  91. a = np.eye(groups + 1, groups, k=-1)
  92. a -= np.roll(a, 1, axis=1)
  93. a *= np.array(k) ** 2
  94. a[0] = 1
  95. c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
  96. self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
  97. self.bn = nn.BatchNorm2d(c2)
  98. self.act = nn.LeakyReLU(0.1, inplace=True)
  99. def forward(self, x):
  100. return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))