Spaces:
Runtime error
Runtime error
| # This file contains modules common to various models | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| def autopad(k, p=None): # kernel, padding | |
| # Pad to 'same' | |
| if p is None: | |
| p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad | |
| return p | |
| def DWConv(c1, c2, k=1, s=1, act=True): | |
| # Depthwise convolution | |
| return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) | |
| class Conv(nn.Module): | |
| # Standard convolution | |
| def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups | |
| super(Conv, self).__init__() | |
| self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), 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))) | |
| def fuseforward(self, x): | |
| return self.act(self.conv(x)) | |
| class Bottleneck(nn.Module): | |
| # Standard bottleneck | |
| def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion | |
| super(Bottleneck, self).__init__() | |
| c_ = int(c2 * e) # hidden channels | |
| 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 BottleneckCSP(nn.Module): | |
| # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks | |
| def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion | |
| super(BottleneckCSP, self).__init__() | |
| c_ = int(c2 * e) # hidden channels | |
| 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(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(*[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 SPP(nn.Module): | |
| # Spatial pyramid pooling layer used in YOLOv3-SPP | |
| def __init__(self, c1, c2, k=(5, 9, 13)): | |
| super(SPP, self).__init__() | |
| c_ = c1 // 2 # hidden channels | |
| 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 Focus(nn.Module): | |
| # Focus wh information into c-space | |
| def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups | |
| super(Focus, self).__init__() | |
| self.conv = Conv(c1 * 4, c2, k, s, p, g, act) | |
| def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) | |
| 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): | |
| # Concatenate a list of tensors along dimension | |
| def __init__(self, dimension=1): | |
| super(Concat, self).__init__() | |
| self.d = dimension | |
| def forward(self, x): | |
| return torch.cat(x, self.d) | |
| class Flatten(nn.Module): | |
| # Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions | |
| def forward(x): | |
| return x.view(x.size(0), -1) | |
| class Classify(nn.Module): | |
| # Classification head, i.e. x(b,c1,20,20) to x(b,c2) | |
| def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups | |
| super(Classify, self).__init__() | |
| self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1) | |
| self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) # to x(b,c2,1,1) | |
| self.flat = Flatten() | |
| def forward(self, x): | |
| z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list | |
| return self.flat(self.conv(z)) # flatten to x(b,c2) | |