From 655895a8382fcafd0e346794f401c081f57f2dc4 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 3 Jul 2020 13:46:12 -0700
Subject: [PATCH] module updates

---
 models/experimental.py | 22 ++++++----------------
 models/yolo.py         |  2 +-
 2 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/models/experimental.py b/models/experimental.py
index cff9d14..539e7f9 100644
--- a/models/experimental.py
+++ b/models/experimental.py
@@ -4,12 +4,13 @@ from models.common import *
 
 
 class CrossConv(nn.Module):
-    # Cross Convolution
-    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
+    # Cross Convolution Downsample
+    def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
+        # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
         super(CrossConv, self).__init__()
         c_ = int(c2 * e)  # hidden channels
-        self.cv1 = Conv(c1, c_, (1, 3), 1)
-        self.cv2 = Conv(c_, c2, (3, 1), 1, g=g)
+        self.cv1 = Conv(c1, c_, (1, k), (1, s))
+        self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
         self.add = shortcut and c1 == c2
 
     def forward(self, x):
@@ -27,7 +28,7 @@ class C3(nn.Module):
         self.cv4 = Conv(2 * c_, c2, 1, 1)
         self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
         self.act = nn.LeakyReLU(0.1, inplace=True)
-        self.m = nn.Sequential(*[CrossConv(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
+        self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
 
     def forward(self, x):
         y1 = self.cv3(self.m(self.cv1(x)))
@@ -84,17 +85,6 @@ class GhostBottleneck(nn.Module):
         return self.conv(x) + self.shortcut(x)
 
 
-class ConvPlus(nn.Module):
-    # Plus-shaped convolution
-    def __init__(self, c1, c2, k=3, s=1, g=1, bias=True):  # ch_in, ch_out, kernel, stride, groups
-        super(ConvPlus, self).__init__()
-        self.cv1 = nn.Conv2d(c1, c2, (k, 1), s, (k // 2, 0), groups=g, bias=bias)
-        self.cv2 = nn.Conv2d(c1, c2, (1, k), s, (0, k // 2), groups=g, bias=bias)
-
-    def forward(self, x):
-        return self.cv1(x) + self.cv2(x)
-
-
 class MixConv2d(nn.Module):
     # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
     def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
diff --git a/models/yolo.py b/models/yolo.py
index c9e6c49..7cc86b2 100644
--- a/models/yolo.py
+++ b/models/yolo.py
@@ -161,7 +161,7 @@ def parse_model(md, ch):  # model_dict, input_channels(3)
                 pass
 
         n = max(round(n * gd), 1) if n > 1 else n  # depth gain
-        if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, ConvPlus, BottleneckCSP]:
+        if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, BottleneckCSP, CrossConv]:
             c1, c2 = ch[f], args[0]
 
             # Normal