Spaces:
Running
Running
import torch | |
import torch.nn as nn | |
import torch.nn.functional as F | |
class L2Norm(nn.Module): | |
def __init__(self, n_channels, scale=1.0): | |
super(L2Norm, self).__init__() | |
self.n_channels = n_channels | |
self.scale = scale | |
self.eps = 1e-10 | |
self.weight = nn.Parameter(torch.Tensor(self.n_channels)) | |
self.weight.data *= 0.0 | |
self.weight.data += self.scale | |
def forward(self, x): | |
norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps | |
x = x / norm * self.weight.view(1, -1, 1, 1) | |
return x | |
class s3fd(nn.Module): | |
def __init__(self): | |
super(s3fd, self).__init__() | |
self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1) | |
self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) | |
self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1) | |
self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1) | |
self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1) | |
self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1) | |
self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1) | |
self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1) | |
self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) | |
self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) | |
self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) | |
self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) | |
self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) | |
self.fc6 = nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=3) | |
self.fc7 = nn.Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0) | |
self.conv6_1 = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0) | |
self.conv6_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1) | |
self.conv7_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0) | |
self.conv7_2 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1) | |
self.conv3_3_norm = L2Norm(256, scale=10) | |
self.conv4_3_norm = L2Norm(512, scale=8) | |
self.conv5_3_norm = L2Norm(512, scale=5) | |
self.conv3_3_norm_mbox_conf = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1) | |
self.conv3_3_norm_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1) | |
self.conv4_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1) | |
self.conv4_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1) | |
self.conv5_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1) | |
self.conv5_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1) | |
self.fc7_mbox_conf = nn.Conv2d(1024, 2, kernel_size=3, stride=1, padding=1) | |
self.fc7_mbox_loc = nn.Conv2d(1024, 4, kernel_size=3, stride=1, padding=1) | |
self.conv6_2_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1) | |
self.conv6_2_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1) | |
self.conv7_2_mbox_conf = nn.Conv2d(256, 2, kernel_size=3, stride=1, padding=1) | |
self.conv7_2_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1) | |
def forward(self, x): | |
h = F.relu(self.conv1_1(x)) | |
h = F.relu(self.conv1_2(h)) | |
h = F.max_pool2d(h, 2, 2) | |
h = F.relu(self.conv2_1(h)) | |
h = F.relu(self.conv2_2(h)) | |
h = F.max_pool2d(h, 2, 2) | |
h = F.relu(self.conv3_1(h)) | |
h = F.relu(self.conv3_2(h)) | |
h = F.relu(self.conv3_3(h)) | |
f3_3 = h | |
h = F.max_pool2d(h, 2, 2) | |
h = F.relu(self.conv4_1(h)) | |
h = F.relu(self.conv4_2(h)) | |
h = F.relu(self.conv4_3(h)) | |
f4_3 = h | |
h = F.max_pool2d(h, 2, 2) | |
h = F.relu(self.conv5_1(h)) | |
h = F.relu(self.conv5_2(h)) | |
h = F.relu(self.conv5_3(h)) | |
f5_3 = h | |
h = F.max_pool2d(h, 2, 2) | |
h = F.relu(self.fc6(h)) | |
h = F.relu(self.fc7(h)) | |
ffc7 = h | |
h = F.relu(self.conv6_1(h)) | |
h = F.relu(self.conv6_2(h)) | |
f6_2 = h | |
h = F.relu(self.conv7_1(h)) | |
h = F.relu(self.conv7_2(h)) | |
f7_2 = h | |
f3_3 = self.conv3_3_norm(f3_3) | |
f4_3 = self.conv4_3_norm(f4_3) | |
f5_3 = self.conv5_3_norm(f5_3) | |
cls1 = self.conv3_3_norm_mbox_conf(f3_3) | |
reg1 = self.conv3_3_norm_mbox_loc(f3_3) | |
cls2 = self.conv4_3_norm_mbox_conf(f4_3) | |
reg2 = self.conv4_3_norm_mbox_loc(f4_3) | |
cls3 = self.conv5_3_norm_mbox_conf(f5_3) | |
reg3 = self.conv5_3_norm_mbox_loc(f5_3) | |
cls4 = self.fc7_mbox_conf(ffc7) | |
reg4 = self.fc7_mbox_loc(ffc7) | |
cls5 = self.conv6_2_mbox_conf(f6_2) | |
reg5 = self.conv6_2_mbox_loc(f6_2) | |
cls6 = self.conv7_2_mbox_conf(f7_2) | |
reg6 = self.conv7_2_mbox_loc(f7_2) | |
# max-out background label | |
chunk = torch.chunk(cls1, 4, 1) | |
bmax = torch.max(torch.max(chunk[0], chunk[1]), chunk[2]) | |
cls1 = torch.cat([bmax, chunk[3]], dim=1) | |
return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6] | |