|
|
|
|
|
|
|
import torch.nn.functional as F |
|
|
|
from utils.utils import * |
|
|
|
|
|
def DWConv(c1, c2, k=1, s=1, act=True): |
|
return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) |
|
|
|
|
|
class Conv(nn.Module): |
|
def __init__(self, c1, c2, k=1, s=1, g=1, act=True): |
|
super(Conv, self).__init__() |
|
self.conv = nn.Conv2d(c1, c2, k, s, k // 2, groups=g, bias=False) |
|
self.bn = nn.BatchNorm2d(c2) |
|
self.act = nn.LeakyReLU(0.1, inplace=True) if act else nn.Identity() |
|
|
|
def forward(self, x): |
|
return self.act(self.bn(self.conv(x))) |
|
|
|
|
|
class Bottleneck(nn.Module): |
|
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): |
|
super(Bottleneck, self).__init__() |
|
c_ = int(c2 * e) |
|
self.cv1 = Conv(c1, c_, 1, 1) |
|
self.cv2 = Conv(c_, c2, 3, 1, g=g) |
|
self.add = shortcut and c1 == c2 |
|
|
|
def forward(self, x): |
|
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) |
|
|
|
|
|
class BottleneckLight(nn.Module): |
|
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): |
|
super(BottleneckLight, self).__init__() |
|
c_ = int(c2 * e) |
|
self.cv1 = Conv(c1, c_, 1, 1) |
|
self.cv2 = nn.Conv2d(c_, c2, 3, 1, 3 // 2, groups=g, bias=False) |
|
self.bn = nn.BatchNorm2d(c2) |
|
self.act = nn.LeakyReLU(0.1, inplace=True) |
|
self.add = shortcut and c1 == c2 |
|
|
|
def forward(self, x): |
|
return self.act(self.bn(x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)))) |
|
|
|
|
|
class BottleneckCSP(nn.Module): |
|
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): |
|
super(BottleneckCSP, self).__init__() |
|
c_ = int(c2 * e) |
|
self.cv1 = Conv(c1, c_, 1, 1) |
|
self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) |
|
self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) |
|
self.cv4 = Conv(c2, c2, 1, 1) |
|
self.bn = nn.BatchNorm2d(2 * c_) |
|
self.act = nn.LeakyReLU(0.1, inplace=True) |
|
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) |
|
|
|
def forward(self, x): |
|
y1 = self.cv3(self.m(self.cv1(x))) |
|
y2 = self.cv2(x) |
|
return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) |
|
|
|
|
|
class Narrow(nn.Module): |
|
def __init__(self, c1, c2, shortcut=True, g=1): |
|
super(Narrow, self).__init__() |
|
c_ = c2 // 2 |
|
self.cv1 = Conv(c1, c_, 1, 1) |
|
self.cv2 = Conv(c_, c2, 3, 1, g=g) |
|
self.add = shortcut and c1 == c2 |
|
|
|
def forward(self, x): |
|
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) |
|
|
|
|
|
class Origami(nn.Module): |
|
def forward(self, x): |
|
y = F.pad(x, [1, 1, 1, 1]) |
|
return torch.cat([x, y[..., :-2, 1:-1], y[..., 1:-1, :-2], y[..., 2:, 1:-1], y[..., 1:-1, 2:]], 1) |
|
|
|
|
|
class ConvPlus(nn.Module): |
|
def __init__(self, c1, c2, k=3, s=1, g=1, bias=True): |
|
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 SPP(nn.Module): |
|
def __init__(self, c1, c2, k=(5, 9, 13)): |
|
super(SPP, self).__init__() |
|
c_ = c1 // 2 |
|
self.cv1 = Conv(c1, c_, 1, 1) |
|
self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) |
|
self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) |
|
|
|
def forward(self, x): |
|
x = self.cv1(x) |
|
return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) |
|
|
|
|
|
class Flatten(nn.Module): |
|
|
|
def forward(self, x): |
|
return x.view(x.size(0), -1) |
|
|
|
|
|
class Focus(nn.Module): |
|
|
|
def __init__(self, c1, c2, k=1): |
|
super(Focus, self).__init__() |
|
self.conv = Conv(c1 * 4, c2, k, 1) |
|
|
|
def forward(self, x): |
|
return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)) |
|
|
|
|
|
class Concat(nn.Module): |
|
|
|
def __init__(self, dimension=1): |
|
super(Concat, self).__init__() |
|
self.d = dimension |
|
|
|
def forward(self, x): |
|
return torch.cat(x, self.d) |
|
|
|
|
|
class MixConv2d(nn.Module): |
|
|
|
def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): |
|
super(MixConv2d, self).__init__() |
|
groups = len(k) |
|
if equal_ch: |
|
i = torch.linspace(0, groups - 1E-6, c2).floor() |
|
c_ = [(i == g).sum() for g in range(groups)] |
|
else: |
|
b = [c2] + [0] * groups |
|
a = np.eye(groups + 1, groups, k=-1) |
|
a -= np.roll(a, 1, axis=1) |
|
a *= np.array(k) ** 2 |
|
a[0] = 1 |
|
c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() |
|
|
|
self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)]) |
|
self.bn = nn.BatchNorm2d(c2) |
|
self.act = nn.LeakyReLU(0.1, inplace=True) |
|
|
|
def forward(self, x): |
|
return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) |
|
|