glenn-jocher
commited on
Commit
•
655895a
1
Parent(s):
df224a0
module updates
Browse files- models/experimental.py +6 -16
- 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,
|
|
|
9 |
super(CrossConv, self).__init__()
|
10 |
c_ = int(c2 * e) # hidden channels
|
11 |
-
self.cv1 = Conv(c1, c_, (1,
|
12 |
-
self.cv2 = Conv(c_, c2, (
|
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_,
|
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,
|
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
|