glenn-jocher commited on
Commit
655895a
1 Parent(s): df224a0

module updates

Browse files
Files changed (2) hide show
  1. models/experimental.py +6 -16
  2. models/yolo.py +1 -1
models/experimental.py CHANGED
@@ -4,12 +4,13 @@ from models.common import *
4
 
5
 
6
  class CrossConv(nn.Module):
7
- # Cross Convolution
8
- def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
 
9
  super(CrossConv, self).__init__()
10
  c_ = int(c2 * e) # hidden channels
11
- self.cv1 = Conv(c1, c_, (1, 3), 1)
12
- self.cv2 = Conv(c_, c2, (3, 1), 1, g=g)
13
  self.add = shortcut and c1 == c2
14
 
15
  def forward(self, x):
@@ -27,7 +28,7 @@ class C3(nn.Module):
27
  self.cv4 = Conv(2 * c_, c2, 1, 1)
28
  self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
29
  self.act = nn.LeakyReLU(0.1, inplace=True)
30
- self.m = nn.Sequential(*[CrossConv(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
31
 
32
  def forward(self, x):
33
  y1 = self.cv3(self.m(self.cv1(x)))
@@ -84,17 +85,6 @@ class GhostBottleneck(nn.Module):
84
  return self.conv(x) + self.shortcut(x)
85
 
86
 
87
- class ConvPlus(nn.Module):
88
- # Plus-shaped convolution
89
- def __init__(self, c1, c2, k=3, s=1, g=1, bias=True): # ch_in, ch_out, kernel, stride, groups
90
- super(ConvPlus, self).__init__()
91
- self.cv1 = nn.Conv2d(c1, c2, (k, 1), s, (k // 2, 0), groups=g, bias=bias)
92
- self.cv2 = nn.Conv2d(c1, c2, (1, k), s, (0, k // 2), groups=g, bias=bias)
93
-
94
- def forward(self, x):
95
- return self.cv1(x) + self.cv2(x)
96
-
97
-
98
  class MixConv2d(nn.Module):
99
  # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
100
  def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
 
4
 
5
 
6
  class CrossConv(nn.Module):
7
+ # Cross Convolution Downsample
8
+ def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
9
+ # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
10
  super(CrossConv, self).__init__()
11
  c_ = int(c2 * e) # hidden channels
12
+ self.cv1 = Conv(c1, c_, (1, k), (1, s))
13
+ self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
14
  self.add = shortcut and c1 == c2
15
 
16
  def forward(self, x):
 
28
  self.cv4 = Conv(2 * c_, c2, 1, 1)
29
  self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
30
  self.act = nn.LeakyReLU(0.1, inplace=True)
31
+ self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
32
 
33
  def forward(self, x):
34
  y1 = self.cv3(self.m(self.cv1(x)))
 
85
  return self.conv(x) + self.shortcut(x)
86
 
87
 
 
 
 
 
 
 
 
 
 
 
 
88
  class MixConv2d(nn.Module):
89
  # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
90
  def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
models/yolo.py CHANGED
@@ -161,7 +161,7 @@ def parse_model(md, ch): # model_dict, input_channels(3)
161
  pass
162
 
163
  n = max(round(n * gd), 1) if n > 1 else n # depth gain
164
- if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, ConvPlus, BottleneckCSP]:
165
  c1, c2 = ch[f], args[0]
166
 
167
  # Normal
 
161
  pass
162
 
163
  n = max(round(n * gd), 1) if n > 1 else n # depth gain
164
+ if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, BottleneckCSP, CrossConv]:
165
  c1, c2 = ch[f], args[0]
166
 
167
  # Normal