Você não pode selecionar mais de 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.

experimental.py 5.4KB

4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
4 anos atrás
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. # This file contains experimental modules
  2. from models.common import *
  3. from utils import google_utils
  4. class CrossConv(nn.Module):
  5. # Cross Convolution Downsample
  6. def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
  7. # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
  8. super(CrossConv, self).__init__()
  9. c_ = int(c2 * e) # hidden channels
  10. self.cv1 = Conv(c1, c_, (1, k), (1, s))
  11. self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
  12. self.add = shortcut and c1 == c2
  13. def forward(self, x):
  14. return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
  15. class C3(nn.Module):
  16. # Cross Convolution CSP
  17. def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
  18. super(C3, self).__init__()
  19. c_ = int(c2 * e) # hidden channels
  20. self.cv1 = Conv(c1, c_, 1, 1)
  21. self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
  22. self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
  23. self.cv4 = Conv(2 * c_, c2, 1, 1)
  24. self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
  25. self.act = nn.LeakyReLU(0.1, inplace=True)
  26. self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
  27. def forward(self, x):
  28. y1 = self.cv3(self.m(self.cv1(x)))
  29. y2 = self.cv2(x)
  30. return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
  31. class Sum(nn.Module):
  32. # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
  33. def __init__(self, n, weight=False): # n: number of inputs
  34. super(Sum, self).__init__()
  35. self.weight = weight # apply weights boolean
  36. self.iter = range(n - 1) # iter object
  37. if weight:
  38. self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights
  39. def forward(self, x):
  40. y = x[0] # no weight
  41. if self.weight:
  42. w = torch.sigmoid(self.w) * 2
  43. for i in self.iter:
  44. y = y + x[i + 1] * w[i]
  45. else:
  46. for i in self.iter:
  47. y = y + x[i + 1]
  48. return y
  49. class GhostConv(nn.Module):
  50. # Ghost Convolution https://github.com/huawei-noah/ghostnet
  51. def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
  52. super(GhostConv, self).__init__()
  53. c_ = c2 // 2 # hidden channels
  54. self.cv1 = Conv(c1, c_, k, s, g, act)
  55. self.cv2 = Conv(c_, c_, 5, 1, c_, act)
  56. def forward(self, x):
  57. y = self.cv1(x)
  58. return torch.cat([y, self.cv2(y)], 1)
  59. class GhostBottleneck(nn.Module):
  60. # Ghost Bottleneck https://github.com/huawei-noah/ghostnet
  61. def __init__(self, c1, c2, k, s):
  62. super(GhostBottleneck, self).__init__()
  63. c_ = c2 // 2
  64. self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
  65. DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
  66. GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
  67. self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
  68. Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
  69. def forward(self, x):
  70. return self.conv(x) + self.shortcut(x)
  71. class MixConv2d(nn.Module):
  72. # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
  73. def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
  74. super(MixConv2d, self).__init__()
  75. groups = len(k)
  76. if equal_ch: # equal c_ per group
  77. i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices
  78. c_ = [(i == g).sum() for g in range(groups)] # intermediate channels
  79. else: # equal weight.numel() per group
  80. b = [c2] + [0] * groups
  81. a = np.eye(groups + 1, groups, k=-1)
  82. a -= np.roll(a, 1, axis=1)
  83. a *= np.array(k) ** 2
  84. a[0] = 1
  85. c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
  86. self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
  87. self.bn = nn.BatchNorm2d(c2)
  88. self.act = nn.LeakyReLU(0.1, inplace=True)
  89. def forward(self, x):
  90. return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
  91. class Ensemble(nn.ModuleList):
  92. # Ensemble of models
  93. def __init__(self):
  94. super(Ensemble, self).__init__()
  95. def forward(self, x, augment=False):
  96. y = []
  97. for module in self:
  98. y.append(module(x, augment)[0])
  99. # y = torch.stack(y).max(0)[0] # max ensemble
  100. # y = torch.cat(y, 1) # nms ensemble
  101. y = torch.stack(y).mean(0) # mean ensemble
  102. return y, None # inference, train output
  103. def attempt_load(weights, map_location=None):
  104. # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
  105. model = Ensemble()
  106. for w in weights if isinstance(weights, list) else [weights]:
  107. google_utils.attempt_download(w)
  108. model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model
  109. if len(model) == 1:
  110. return model[-1] # return model
  111. else:
  112. print('Ensemble created with %s\n' % weights)
  113. for k in ['names', 'stride']:
  114. setattr(model, k, getattr(model[-1], k))
  115. return model # return ensemble